-- ArkLib/Data/FieldTheory/BinaryField/Tower/Impl.lean

import ArkLib.Data.Classes.DCast

import ArkLib.Data.FieldTheory.BinaryField.Tower.Basic

namespace ConcreteBinaryTower

open Polynomial

def ConcreteBTField : ℕ → Type := fun k => BitVec (2 ^ k)

section BitVecDCast

instance BitVec.instDCast : DCast Nat BitVec where
  dcast h := BitVec.cast h
  dcast_id := by admit /- proof elided -/

end BitVecDCast

section ConversionUtils

def fromNat {k : ℕ} (n : Nat) : ConcreteBTField k :=
  BitVec.ofNat (2 ^ k) n

instance ConcreteBTField.instDCast_local : DCast ℕ ConcreteBTField where
  dcast h_k_eq term_k1 := BitVec.cast (congrArg (fun n => 2 ^ n) h_k_eq) term_k1
  dcast_id := by admit /- proof elided -/

end ConversionUtils

section NumericLemmas

end NumericLemmas

section FieldOperationsAndInstances

def add {k : ℕ} (x y : ConcreteBTField k) : ConcreteBTField k := BitVec.xor x y

instance (k : ℕ) : HAdd (ConcreteBTField k) (ConcreteBTField k) (ConcreteBTField k) where
  hAdd := add

instance (k : ℕ) : Add (ConcreteBTField k) where
  add := add

-- split extracts the high and low halves of a bitvector using BitVec.extractLsb,
-- then casts them to the correct width using dcast. It returns (hi, lo).
def split {k : ℕ} (h : k > 0) (x : ConcreteBTField k) :
    ConcreteBTField (k - 1) × ConcreteBTField (k - 1) :=
  let lo_bits : BitVec (2 ^ (k - 1) - 1 - 0 + 1) :=
    BitVec.extractLsb (hi := 2 ^ (k - 1) - 1) (lo := 0) x
  let hi_bits : BitVec (2 ^ k - 1 - 2 ^ (k - 1) + 1) :=
    BitVec.extractLsb (hi := 2 ^ k - 1) (lo := 2 ^ (k - 1)) x
  have h_lo : 2 ^ (k - 1) - 1 - 0 + 1 = 2 ^ (k - 1) := by admit /- proof elided -/
  have h_hi : 2 ^ k - 1 - 2 ^ (k - 1) + 1 = 2 ^ (k - 1) := by admit /- proof elided -/
  let lo : ConcreteBTField (k - 1) := dcast h_lo lo_bits
  let hi : ConcreteBTField (k - 1) := dcast h_hi hi_bits
  (hi, lo)

def join {k : ℕ} (h_pos : k > 0) (hi lo : ConcreteBTField (k - 1)) : ConcreteBTField k :=

-- Repository definitions used:
content: class DCast (α : Sort*) (β : α → Sort*) where
  dcast : ∀ {a a' : α}, a = a' → β a → β a'
  dcast_id : ∀ {a : α}, dcast (Eq.refl a) = id

content: def AlgebraTowerEquiv.symm (e : AlgebraTowerEquiv A B) : AlgebraTowerEquiv B A where
  toRingEquiv := fun i => (e.toRingEquiv i).symm
  commutesLeft' := fun i j h r => by admit /- proof elided -/

content: structure AlgebraTowerEquiv (A : ι → Type*) [∀ i, CommSemiring (A i)] [a : AlgebraTower A]
  (B : ι → Type*) [∀ i, CommSemiring (B i)] [b : AlgebraTower B]
  where
    toRingEquiv : ∀ i, (A i ≃+* B i)
    commutesLeft' : ∀ (i j : ι) (h : i ≤ j) (r : A i),
      (b.algebraMap (i:=i) (j:=j) (h:=h)) ((toRingEquiv i) r) =
      (toRingEquiv j) (a.algebraMap (i:=i) (j:=j) (h:=h) r)

-- Repository lemmas used:
content: theorem one_le_two_pow_n (n : ℕ) : 1 ≤ 2 ^ n

content: @[simp]
theorem dcast_eq : dcast (Eq.refl a) b = b

content: theorem one_le_sub_consecutive_two_pow (n : ℕ): 1 ≤ 2^(n+1) - 2^n

content: theorem dcast_eq_root_cast (h : a = a') : dcast h b = _root_.cast (congrArg β h) b

content: theorem dcast_symm (ha : a = a') (hb : dcast ha b = b') : dcast (ha.symm) b' = b

content: theorem dcast_eq_dcast_iff (h : a = a'') (h' : a' = a'') :
    dcast h b = dcast h' b' ↔ b = dcast (h'.trans h.symm) b'

content: @[simp]
theorem dcast_trans (h : a = a') (h' : a' = a'') :
    dcast h' (dcast h b) = dcast (h.trans h') b

content: theorem heq_of_dcast (ha : a = a') (hb : dcast ha b = b') : HEq b b'

content: -- Note: For BitVec, dcast h is just BitVec.cast h (see BitVec.instDCast).
-- This means for any h : n = m, we have: dcast h (x ^^^ y) = dcast h x ^^^ dcast h y
-- This can be proved by: subst h; rfl
-- Similarly, BitVec.extractLsb distributes over xor because shift and mask distribute over xor.

-- Local definitions used:
content: def ConcreteBTField : ℕ → Type := fun k => BitVec (2 ^ k)

content: instance BitVec.instDCast : DCast Nat BitVec where
  dcast h := BitVec.cast h
  dcast_id := by admit /- proof elided -/

content: def fromNat {k : ℕ} (n : Nat) : ConcreteBTField k :=
  BitVec.ofNat (2 ^ k) n

content: instance ConcreteBTField.instDCast_local : DCast ℕ ConcreteBTField where
  dcast h_k_eq term_k1 := BitVec.cast (congrArg (fun n => 2 ^ n) h_k_eq) term_k1
  dcast_id := by admit /- proof elided -/

content: def add {k : ℕ} (x y : ConcreteBTField k) : ConcreteBTField k := BitVec.xor x y

content: def split {k : ℕ} (h : k > 0) (x : ConcreteBTField k) :
    ConcreteBTField (k - 1) × ConcreteBTField (k - 1) :=
  let lo_bits : BitVec (2 ^ (k - 1) - 1 - 0 + 1) :=
    BitVec.extractLsb (hi := 2 ^ (k - 1) - 1) (lo := 0) x
  let hi_bits : BitVec (2 ^ k - 1 - 2 ^ (k - 1) + 1) :=
    BitVec.extractLsb (hi := 2 ^ k - 1) (lo := 2 ^ (k - 1)) x
  have h_lo : 2 ^ (k - 1) - 1 - 0 + 1 = 2 ^ (k - 1) := by admit /- proof elided -/

content: def join {k : ℕ} (h_pos : k > 0) (hi lo : ConcreteBTField (k - 1)) : ConcreteBTField k :=

-- Local lemmas used:
content: theorem BitVec.dcast_id {n : Nat} (bv : BitVec n) :
  DCast.dcast (Eq.refl n) bv = bv

content: theorem BitVec.dcast_bitvec_eq {l r val : ℕ} (h_width_eq : l = r) :
    dcast h_width_eq (BitVec.ofNat l val) = BitVec.ofNat r val

content: theorem BitVec.dcast_bitvec_toNat_eq {w w2 : ℕ} (x : BitVec w) (h_width_eq : w = w2) :
    BitVec.toNat x = BitVec.toNat (dcast (h_width_eq) x)

content: theorem BitVec.dcast_bitvec_extractLsb_eq {w hi1 lo1 hi2 lo2 : ℕ}
    (x : BitVec w) (h_lo_eq : lo1 = lo2)
    (h_width_eq : hi1 - lo1 + 1 = hi2 - lo2 + 1) :
    dcast h_width_eq (BitVec.extractLsb (hi:=hi1) (lo:=lo1) x)
      = BitVec.extractLsb (hi:=hi2) (lo:=lo2) (x)

content: theorem BitVec.eq_mp_eq_dcast {w w2 : ℕ} (x : BitVec w) (h_width_eq : w = w2)
  (h_bitvec_eq : BitVec w = BitVec w2 := by rw [h_width_eq]) :
  Eq.mp (h:=h_bitvec_eq) (a:=x) = dcast (h_width_eq) (x)

content: theorem BitVec.extractLsb_concat_hi {hi_size lo_size : ℕ} (hi : BitVec hi_size)
  (lo : BitVec lo_size) (h_hi : hi_size > 0) :
  BitVec.extractLsb (hi:=hi_size + lo_size - 1) (lo:=lo_size)
  (BitVec.append (msbs:=hi) (lsbs:=lo)) = dcast (by
    rw [←Nat.sub_add_comm (by omega), Nat.sub_add_cancel (by omega), Nat.add_sub_cancel]
  ) hi

content: theorem BitVec.extractLsb_concat_lo {hi_size lo_size : ℕ} (hi : BitVec hi_size)
  (lo : BitVec lo_size) (h_lo : lo_size > 0) : BitVec.extractLsb (hi:=lo_size - 1) (lo:=0)
  (BitVec.append (msbs:=hi) (lsbs:=lo)) = dcast (by
    rw [←Nat.sub_add_comm (h:=by omega), Nat.sub_add_cancel (h:=by omega), Nat.sub_zero]
  ) lo

content: theorem Nat.shiftRight_lo_mod_2_pow_hi_shiftLeft_lo (n hi_len lo_len : ℕ)
  (h_n : n < 2 ^ (hi_len + lo_len)) :
  (((n >>> lo_len) % (2 ^ hi_len)) <<< lo_len) = (n - n % 2 ^ lo_len)

content: theorem Nat.reconstruct_from_hi_and_lo_parts (n hi_len lo_len : ℕ)
    (h_n : n < 2 ^ (hi_len + lo_len)) :
    n = (((n >>> lo_len) % (2 ^ hi_len)) <<< lo_len) + (n % (2 ^ lo_len))

content: theorem Nat.reconstruct_from_hi_and_lo_parts_or_ver (n hi_len lo_len : ℕ)
  (h_n : n < 2 ^ (hi_len + lo_len)) :
    n = (((n >>> lo_len) % (2 ^ hi_len)) <<< lo_len) ||| (n % (2 ^ lo_len))

content: theorem BitVec.eq_append_iff_extract {lo_size hi_size : ℕ} (lo : BitVec lo_size)
  (hi : BitVec hi_size) (h_hi_gt_0 : hi_size > 0) (h_lo_gt_0 : lo_size > 0)
  (x : BitVec (hi_size + lo_size)) : x = dcast (by rfl) (BitVec.append (msbs:=hi) (lsbs:=lo)) ↔
  hi = dcast (by omega) (BitVec.extractLsb (hi:=hi_size + lo_size - 1) (lo:=lo_size) x) ∧
  lo = dcast (by omega) (BitVec.extractLsb (hi:=lo_size - 1) (lo:=0) x)

content: lemma one_le_sub_middle_of_pow2 {k : ℕ} (h_k : 1 ≤ k) : 1 ≤ 2 ^ k - 2 ^ (k - 1)

content: lemma sub_middle_of_pow2_with_one_canceled {k : ℕ} (h_k : 1 ≤ k) : 2 ^ k - 1 - 2 ^ (k - 1) + 1
  = 2 ^ (k - 1)

content: lemma h_sub_middle {k : ℕ} (h_pos : k > 0) : 2 ^ k - 1 - 2 ^ (k - 1) + 1 = 2 ^ (k - 1)

content: lemma h_middle_sub {k : ℕ} : 2 ^ (k - 1) - 1 - 0 + 1 = 2 ^ (k - 1)

content: lemma h_sum_two_same_pow2 {k : ℕ} (h_pos : k > 0) : 2 ^ (k - 1) + 2 ^ (k - 1) = 2 ^ k

content: theorem sum_fromNat_eq_from_xor_Nat {k : ℕ} (x y : Nat) :
  fromNat (k:=k) (x ^^^ y) = fromNat (k:=k) x + fromNat (k:=k) y

content: theorem BitVec.extractLsb_eq_shift_ofNat {n : Nat} (x : BitVec n) (l r : Nat) :
    BitVec.extractLsb r l x = BitVec.ofNat (r - l + 1) (x.toNat >>> l)

content: theorem setWidth_eq_ofNat_mod {n num_bits : Nat} (x : BitVec n) :
  BitVec.setWidth num_bits x = BitVec.ofNat num_bits (x.toNat % 2 ^ num_bits)

content: theorem BitVec.extractLsb_eq_and_pow_2_minus_1_ofNat {n num_bits : Nat}
  (h_num_bits : num_bits > 0) (x : BitVec n) :
  BitVec.extractLsb (hi:= num_bits - 1) (lo := 0) x =
    BitVec.ofNat (num_bits - 1 - 0 + 1) (x.toNat &&& (2 ^ num_bits - 1))

content: theorem split_bitvec_eq_iff_fromNat {k : ℕ} (h_pos : k > 0) (x : ConcreteBTField k)
  (hi_btf lo_btf : ConcreteBTField (k - 1)) :
  split h_pos x = (hi_btf, lo_btf) ↔
  (hi_btf = fromNat (k:=k - 1) (x.toNat >>> 2 ^ (k - 1)) ∧
  lo_btf = fromNat (k:=k - 1) (x.toNat &&& (2 ^ (2 ^ (k - 1)) - 1)))

content: theorem join_eq_iff_dcast_extractLsb {k : ℕ} (h_pos : k > 0) (x : ConcreteBTField k)
  (hi_btf lo_btf : ConcreteBTField (k - 1)) :
  x = 《 hi_btf, lo_btf 》 ↔
  (hi_btf = dcast (h_sub_middle h_pos) (BitVec.extractLsb (hi := 2 ^ k - 1) (lo := 2 ^ (k - 1)) x) ∧
  lo_btf = dcast (h_middle_sub) (BitVec.extractLsb (hi := 2 ^ (k - 1) - 1) (lo := 0) x))

content: theorem join_eq_bitvec_iff_fromNat {k : ℕ} (h_pos : k > 0) (x : ConcreteBTField k)
  (hi_btf lo_btf : ConcreteBTField (k - 1)) :
  x = 《 hi_btf, lo_btf 》 ↔
  (hi_btf = fromNat (k:=k - 1) (x.toNat >>> 2 ^ (k - 1)) ∧
  lo_btf = fromNat (k:=k - 1) (x.toNat &&& (2 ^ (2 ^ (k - 1)) - 1)))

content: theorem join_of_split {k : ℕ} (h_pos : k > 0) (x : ConcreteBTField k)
    (hi_btf lo_btf : ConcreteBTField (k - 1))
    (h_split_eq : split h_pos x = (hi_btf, lo_btf)) :
    x = 《 hi_btf, lo_btf 》

-- Target theorem and proof:
theorem split_sum_eq_sum_split {k : ℕ} (h_pos : k > 0) (x₀ x₁ : ConcreteBTField k)
  (hi₀ lo₀ hi₁ lo₁ : ConcreteBTField (k - 1))
  (h_split_x₀ : split h_pos x₀ = (hi₀, lo₀))
  (h_split_x₁ : split h_pos x₁ = (hi₁, lo₁)) :
  split h_pos (x₀ + x₁) = (hi₀ + hi₁, lo₀ + lo₁)
:= by
  have h_x₀ := join_of_split h_pos x₀ hi₀ lo₀ h_split_x₀
  have h_x₁ := join_of_split h_pos x₁ hi₁ lo₁ h_split_x₁
  -- Approach : convert equation to Nat realm for simple proof
  have h₀ := (split_bitvec_eq_iff_fromNat (k:=k) (h_pos:=h_pos) x₀ hi₀ lo₀).mp h_split_x₀
  have h₁ := (split_bitvec_eq_iff_fromNat (k:=k) (h_pos:=h_pos) x₁ hi₁ lo₁).mp h_split_x₁
  have h_sum_hi : (hi₀ + hi₁) = fromNat (BitVec.toNat (x₀ + x₁) >>> 2 ^ (k - 1)) := by
    rw [h₀.1, h₁.1]
    rw [←sum_fromNat_eq_from_xor_Nat]
    have h_nat_eq : BitVec.toNat x₀ >>> 2 ^ (k - 1) ^^^ BitVec.toNat x₁ >>> 2 ^ (k - 1)
      = BitVec.toNat (x₀ + x₁) >>> 2 ^ (k - 1) := by
      -- unfold Concrete BTF addition into BitVec.xor
      simp only [instHAddConcreteBTField, add, BitVec.xor_eq]
      rw [Nat.shiftRight_xor_distrib.symm]
      rw [BitVec.toNat_xor] -- distribution of BitVec.xor over BitVec.toNat
    rw [h_nat_eq]
  have h_sum_lo : (lo₀ + lo₁) = fromNat (BitVec.toNat (x₀ + x₁) &&& 2 ^ 2 ^ (k - 1) - 1) := by
    rw [h₀.2, h₁.2]
    rw [←sum_fromNat_eq_from_xor_Nat]
    have h_nat_eq : BitVec.toNat x₀ &&& 2 ^ 2 ^ (k - 1) - 1 ^^^ BitVec.toNat x₁
      &&& 2 ^ 2 ^ (k - 1) - 1 = BitVec.toNat (x₀ + x₁) &&& 2 ^ 2 ^ (k - 1) - 1 := by
      simp only [instHAddConcreteBTField, add, BitVec.xor_eq]
      rw [BitVec.toNat_xor]
      rw [Nat.and_xor_distrib_right.symm]
    rw [h_nat_eq]
  have h_sum_hi_lo : (hi₀ + hi₁, lo₀ + lo₁) = split h_pos (x₀ + x₁) := by
    rw [(split_bitvec_eq_iff_fromNat (k:=k) (h_pos:=h_pos) (x₀ + x₁)
      (hi₀ + hi₁) (lo₀ + lo₁)).mpr ⟨h_sum_hi, h_sum_lo⟩]
  exact h_sum_hi_lo.symm