-- 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 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 -/

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

-- 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_symm (ha : a = a') (hb : dcast ha b = b') : dcast (ha.symm) b' = b

-- 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 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 -/

-- Local lemmas used:
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: 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: 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))

-- Target theorem and proof:
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)))
:= by
  have lhs_lo_case := BitVec.extractLsb_eq_and_pow_2_minus_1_ofNat (num_bits:=2 ^ (k - 1))
    (n:=2 ^ k) (Nat.two_pow_pos (k - 1)) (x:=x)
  have rhs_hi_case_bitvec_eq := BitVec.extractLsb_eq_shift_ofNat (n:=2 ^ k) (r:=2 ^ k - 1)
    (l:=2 ^ (k - 1)) (x:=x)
  constructor
  · -- Forward direction : split x = (hi_btf, lo_btf) → bitwise operations
    intro h_split
    unfold split at h_split
    have ⟨h_hi, h_lo⟩ := Prod.ext_iff.mp h_split
    simp only at h_hi h_lo
    have hi_eq : hi_btf = fromNat (k:=k - 1) (x.toNat >>> 2 ^ (k - 1)) := by
      unfold fromNat
      rw [←h_hi]
      rw [dcast_symm (h_sub_middle h_pos).symm]
      rw [rhs_hi_case_bitvec_eq]
      rw [BitVec.dcast_bitvec_eq]
    have lo_eq : lo_btf = fromNat (k:=k - 1) (x.toNat &&& ((2 ^ (2 ^ (k - 1)) - 1))) := by
      unfold fromNat
      rw [←h_lo]
      have rhs_lo_case_bitvec_eq :=
        BitVec.extractLsb_eq_shift_ofNat (n:=2 ^ k) (r:=2 ^ (k - 1) - 1) (l:=0) (x:=x)
      rw [dcast_symm (h_middle_sub).symm]
      rw [rhs_lo_case_bitvec_eq]
      rw [BitVec.dcast_bitvec_eq] -- remove dcast
      rw [←lhs_lo_case]
      exact rhs_lo_case_bitvec_eq
    exact ⟨hi_eq, lo_eq⟩
  · -- Backward direction : bitwise operations → split x = (hi_btf, lo_btf)
    intro h_bits
    unfold split
    have ⟨h_hi, h_lo⟩ := h_bits
    have hi_extract_eq : dcast (h_sub_middle h_pos) (BitVec.extractLsb (hi := 2 ^ k - 1)
      (lo := 2 ^ (k - 1)) x) = fromNat (k:=k - 1) (x.toNat >>> 2 ^ (k - 1)) := by
      unfold fromNat
      rw [dcast_symm (h_sub_middle h_pos).symm]
      rw [rhs_hi_case_bitvec_eq]
      rw [BitVec.dcast_bitvec_eq]

    have lo_extract_eq : dcast (h_middle_sub) (BitVec.extractLsb (hi := 2 ^ (k - 1) - 1)
      (lo := 0) x) = fromNat (k:=k - 1) (x.toNat &&& ((2 ^ (2 ^ (k - 1)) - 1))) := by
      unfold fromNat
      rw [lhs_lo_case]
      rw [BitVec.dcast_bitvec_eq]

    simp only [hi_extract_eq, Nat.sub_zero, lo_extract_eq, Nat.and_two_pow_sub_one_eq_mod, h_hi,
      h_lo]