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

import ArkLib.Data.FieldTheory.BinaryField.Tower.Prelude

import ArkLib.Data.RingTheory.AlgebraTower

import Mathlib.Tactic.DepRewrite

namespace BinaryTower

noncomputable section

open Polynomial AdjoinRoot Module

section BTFieldDefs

structure BinaryTowerResult (F : Type _) (k : ℕ) where
  vec : (List.Vector F (k + 1))
  instField : (Field F)
  instFintype : Fintype F
  specialElement : F
  specialElementNeZero : NeZero specialElement
  firstElementOfVecIsSpecialElement [Inhabited F] : vec.1.headI = specialElement
  instIrreduciblePoly : (Irreducible (p := (definingPoly specialElement)))
  sumZeroIffEq : ∀ (x y : F), x + y = 0 ↔ x = y
  fieldFintypeCard : Fintype.card F = 2^(2^k)
  traceMapEvalAtRootsIs1 : TraceMapProperty F specialElement k

structure BinaryTowerInductiveStepResult (k : ℕ) (prevBTField : Type _)
  (prevBTResult : BinaryTowerResult prevBTField k) [instPrevBTFieldIsField : Field prevBTField]
  (prevPoly : Polynomial prevBTField) (F : Type _) where
  binaryTowerResult : BinaryTowerResult F (k+1)
  eq_adjoin : F = AdjoinRoot prevPoly
  u_is_root : Eq.mp (eq_adjoin) binaryTowerResult.specialElement = AdjoinRoot.root prevPoly
  eval_defining_poly_at_root : Eq.mp (eq_adjoin) binaryTowerResult.specialElement^2 +
    Eq.mp (eq_adjoin) binaryTowerResult.specialElement * (of prevPoly) prevBTResult.specialElement
    + 1 = 0

def binary_tower_inductive_step
  (k : Nat)
  (prevBTField : Type _) [Field prevBTField]
  (prevBTResult : BinaryTowerResult prevBTField k)
: Σ' (F : Type _), BinaryTowerInductiveStepResult (k:=k) (prevBTField:=prevBTField)
  (prevBTResult:=prevBTResult) (prevPoly:=definingPoly (F:=prevBTField)
    (instField:=prevBTResult.instField) (s:=prevBTResult.specialElement)) (F:=F)
  (instPrevBTFieldIsField:=prevBTResult.instField) :=

def BinaryTowerAux (k : ℕ) : (Σ' (F : Type 0), BinaryTowerResult F k) :=
  match k with
  | 0 => 
    let curBTField := GF(2)
    let newList : List.Vector (GF(2)) 1 := List.Vector.cons (1 : GF(2)) List.Vector.nil
    let specialElement : GF(2) := newList.1.headI
    let firstElementOfVecIsSpecialElement : newList.1.headI = specialElement := rfl
    let specialElementIs1 : specialElement = 1 := by admit /- proof elided -/

@[simp]
def BTField (k : ℕ) := (BinaryTowerAux k).1

@[simp]
instance Inhabited (k : ℕ) : Inhabited (BTField k) where
  default := (0 : BTField k)

@[simp]
def sumZeroIffEq (k : ℕ) : ∀ (x y : BTField k),
  x + y = 0 ↔ x = y := (BinaryTowerAux k).2.sumZeroIffEq

@[simp]
def Z (k : ℕ) : BTField k := (BinaryTowerAux k).snd.specialElement 

@[simp]
def poly (k : ℕ) : Polynomial (BTField k) := definingPoly (s:=(Z k))

instance polyMonic (n : ℕ) : Monic (poly n) := definingPoly_is_monic (Z n)

end BTFieldDefs

section BinaryAlgebraTower

def canonicalEmbedding (k : ℕ) : BTField k →+* BTField (k+1) :=
  AdjoinRoot.of (poly k)

def towerAlgebraMap (l r : ℕ) (h_le : l ≤ r) : BTField l →+* BTField r :=

def binaryAlgebraTower {l r : ℕ} (h_le : l ≤ r) : Algebra (BTField l) (BTField r) :=

instance (priority := 1000) algebra_adjacent_tower (l : ℕ) :
  Algebra (BTField l) (BTField (l+1)) :=

end BinaryAlgebraTower

noncomputable section MultilinearBasis

def join_via_add_smul {k : ℕ} (h_pos : k > 0) (hi_btf lo_btf : BTField (k - 1)) :
    BTField k :=

def split (k : ℕ) (h_k : k > 0) (x : BTField k) : BTField (k-1) × BTField (k-1) :=

-- Repository definitions used:
content: notation : 10 "GF(" term : 10 ")" => GaloisField term 1

content: structure TraceMapProperty (F : Type*) [Field F] (u : F) (k : ℕ) : Prop where
  element_trace : ∑ i ∈ Finset.range (2 ^ k), u ^ (2 ^ i) = 1
  inverse_trace : ∑ i ∈ Finset.range (2 ^ k), (u⁻¹) ^ (2 ^ i) = 1

content: noncomputable def definingPoly {F : Type*} [instField : Field F] (s : F)
  := X ^ 2 + C s * X + 1

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)

content: structure SpecialElementRelation {F_prev : Type*} [Field F_prev] (t1 : F_prev)
  {F_cur : Type*} [Field F_cur] (u : F_cur) [Algebra F_prev F_cur] : Prop where
    sum_inv_eq : u + u⁻¹ = algebraMap F_prev F_cur t1
    h_u_square : u^2 = u * (algebraMap F_prev F_cur t1) + 1

content: instance irreducible_quadratic_defining_poly_of_traceMap_eq_1
  {F : Type*} [Field F] [Fintype F] [CharP F 2] (s : F) [NeZero s] (k : ℕ)
  (trace_map_prop : TraceMapProperty F s k)
  (fintypeCard : Fintype.card F = 2 ^ (2 ^ k))
  : Irreducible (definingPoly s) :=

content: instance charP_eq_2_of_add_self_eq_zero {F : Type*} [Field F]
    (sumZeroIffEq : ∀ (x y : F), x + y = 0 ↔ x = y) : CharP F 2 :=
  have h_two : (2 : (F)) = 0 := by admit /- proof elided -/

content: def coeff.{u} {F : Type u} [Semiring F] (f : F[X][Y]) (i j : ℕ) : F := (f.coeff j).coeff i

content: instance GF_2_fintype : Fintype (GF(2)) := Fintype.ofFinite (GF(2))

content: @[simp]
def AlgebraTower.toAlgebra {i j : ι} (h : i ≤ j) : Algebra (A i) (A j) :=
  (AlgebraTower.algebraMap (i:=i) (j:=j) (h:=h)).toAlgebra

-- Repository lemmas used:
content: lemma degree_definingPoly {F : Type*} [Field F] [Fintype F] (s : F) [NeZero s] :
  (definingPoly s).degree = 2

content: lemma degree_s_smul_X_add_1 {F : Type*} [Field F] [Fintype F] (s : F) [NeZero s] :
  (C s * (X : (F)[X]) + 1).degree = 1

content: lemma degree_s_smul_X {F : Type*} [Field F] [Fintype F] (s : F) [NeZero s] :
  (C s * (X : (F)[X])).degree = 1

content: theorem unique_linear_form_of_elements_in_adjoined_commring
    {R : Type*} [CommRing R] (f : R[X]) (hf_deg : f.natDegree = 2)
    (hf_monic : Monic f) (c1 : AdjoinRoot f) :
    ∃! p : R × R, c1 = (AdjoinRoot.of f) p.1 * root f + (AdjoinRoot.of f) p.2

content: theorem unique_linear_sum_repr (R : Type*) [CommRing R] (S : Type*) [CommRing S] [Algebra R S]
    (pb : PowerBasis R S) (h_dim : pb.dim = 2) (s : S) :
    ∃! p : R × R, s = p.fst • pb.gen + algebraMap R S p.snd

content: theorem unique_repr {R : Type*} [CommRing R] {S : Type*} [CommRing S] [Algebra R S]
    (pb : PowerBasis R S) (repr1 repr2 : Fin pb.dim →₀ R)
    (h : ∑ i : Fin pb.dim, repr1 i • pb.basis i = ∑ i : Fin pb.dim, repr2 i • pb.basis i) :
    repr1 = repr2

-- Local definitions used:
content: structure BinaryTowerResult (F : Type _) (k : ℕ) where
  vec : (List.Vector F (k + 1))
  instField : (Field F)
  instFintype : Fintype F
  specialElement : F
  specialElementNeZero : NeZero specialElement
  firstElementOfVecIsSpecialElement [Inhabited F] : vec.1.headI = specialElement
  instIrreduciblePoly : (Irreducible (p := (definingPoly specialElement)))
  sumZeroIffEq : ∀ (x y : F), x + y = 0 ↔ x = y
  fieldFintypeCard : Fintype.card F = 2^(2^k)
  traceMapEvalAtRootsIs1 : TraceMapProperty F specialElement k

content: structure BinaryTowerInductiveStepResult (k : ℕ) (prevBTField : Type _)
  (prevBTResult : BinaryTowerResult prevBTField k) [instPrevBTFieldIsField : Field prevBTField]
  (prevPoly : Polynomial prevBTField) (F : Type _) where
  binaryTowerResult : BinaryTowerResult F (k+1)
  eq_adjoin : F = AdjoinRoot prevPoly
  u_is_root : Eq.mp (eq_adjoin) binaryTowerResult.specialElement = AdjoinRoot.root prevPoly
  eval_defining_poly_at_root : Eq.mp (eq_adjoin) binaryTowerResult.specialElement^2 +
    Eq.mp (eq_adjoin) binaryTowerResult.specialElement * (of prevPoly) prevBTResult.specialElement
    + 1 = 0

content: def binary_tower_inductive_step
  (k : Nat)
  (prevBTField : Type _) [Field prevBTField]
  (prevBTResult : BinaryTowerResult prevBTField k)
: Σ' (F : Type _), BinaryTowerInductiveStepResult (k:=k) (prevBTField:=prevBTField)
  (prevBTResult:=prevBTResult) (prevPoly:=definingPoly (F:=prevBTField)
    (instField:=prevBTResult.instField) (s:=prevBTResult.specialElement)) (F:=F)
  (instPrevBTFieldIsField:=prevBTResult.instField) :=

content: def BinaryTowerAux (k : ℕ) : (Σ' (F : Type 0), BinaryTowerResult F k) :=
  match k with
  | 0 => 
    let curBTField := GF(2)
    let newList : List.Vector (GF(2)) 1 := List.Vector.cons (1 : GF(2)) List.Vector.nil
    let specialElement : GF(2) := newList.1.headI
    let firstElementOfVecIsSpecialElement : newList.1.headI = specialElement := rfl
    let specialElementIs1 : specialElement = 1 := by admit /- proof elided -/

content: @[simp]
def BTField (k : ℕ) := (BinaryTowerAux k).1

content: @[simp]
instance Inhabited (k : ℕ) : Inhabited (BTField k) where
  default := (0 : BTField k)

content: @[simp]
def sumZeroIffEq (k : ℕ) : ∀ (x y : BTField k),
  x + y = 0 ↔ x = y := (BinaryTowerAux k).2.sumZeroIffEq

content: @[simp]
def Z (k : ℕ) : BTField k := (BinaryTowerAux k).snd.specialElement

content: @[simp]
def poly (k : ℕ) : Polynomial (BTField k) := definingPoly (s:=(Z k))

content: instance polyMonic (n : ℕ) : Monic (poly n) := definingPoly_is_monic (Z n)

content: def canonicalEmbedding (k : ℕ) : BTField k →+* BTField (k+1) :=
  AdjoinRoot.of (poly k)

content: def towerAlgebraMap (l r : ℕ) (h_le : l ≤ r) : BTField l →+* BTField r :=

content: def binaryAlgebraTower {l r : ℕ} (h_le : l ≤ r) : Algebra (BTField l) (BTField r) :=

content: instance (priority := 1000) algebra_adjacent_tower (l : ℕ) :
  Algebra (BTField l) (BTField (l+1)) :=

content: def join_via_add_smul {k : ℕ} (h_pos : k > 0) (hi_btf lo_btf : BTField (k - 1)) :
    BTField k :=

content: def split (k : ℕ) (h_k : k > 0) (x : BTField k) : BTField (k-1) × BTField (k-1) :=

-- Local lemmas used:
content: lemma poly_natDegree_eq_2 (k : ℕ) : (poly (k:=k)).natDegree = 2

content: lemma BTField.cast_BTField_eq (k m : ℕ) (h_eq : k = m) :
  BTField k = BTField m

content: lemma towerAlgebraMap_id (k : ℕ) : towerAlgebraMap (h_le:=by omega) = RingHom.id (BTField k)

content: lemma towerAlgebraMap_succ_1 (k : ℕ) :
  towerAlgebraMap (l:=k) (r:=k+1) (h_le:=by omega) = canonicalEmbedding k

content: lemma binaryTowerAlgebra_def (l r : ℕ) (h_le : l ≤ r) :
    @binaryAlgebraTower (l:=l) (r:=r) (h_le:=h_le)
    = (towerAlgebraMap l r h_le).toAlgebra

content: lemma algebraMap_adjacent_tower_def (l : ℕ) :
  (algebraMap (BTField l) (BTField (l + 1))) = canonicalEmbedding l

content: lemma algebraMap_adjacent_tower_succ_eq_Adjoin_of (k : ℕ) :
  (algebraMap (BTField k) (BTField (k + 1))) = of (poly k)

content: theorem unique_linear_decomposition_succ (k : ℕ) :
  ∀ (x : BTField (k+1)), ∃! (p : BTField k × BTField k),
    x = ⋘ p.1, p.2 ⋙

-- Target theorem and proof:
theorem eq_join_via_add_smul_eq_iff_split (k : ℕ) (h_pos : k > 0)
    (x : BTField k) (hi_btf lo_btf : BTField (k - 1)) :
    x = ⋘ hi_btf, lo_btf ⋙ ↔
  split (k:=k) (h_k:=h_pos) x = (hi_btf, lo_btf)
:= by
  have h_k_sub_1_add_1_eq_k : k - 1 + 1 = k := by omega
  have h_BTField_eq := BTField.cast_BTField_eq (k:=k) (m:=k-1+1) (h_eq:=by omega)
  set p := unique_linear_decomposition_succ (k:=(k-1)) (x:=(Eq.mp (h:=h_BTField_eq) x)) with hp
  -- -- ⊢ x = join_via_add_smul k h_pos hi lo
  have h_p_satisfy := p.choose_spec
  set xPair := p.choose
  constructor
  · intro h_x_eq_join
    -- Due to `unique_linear_decomposition_succ`, there must be exactly one pair
    -- `(hi, lo)` that satisfies the equation : `x = join_via_add_smul k h_pos hi lo`
    -- Now we prove `⟨hi_btf, lo_btf⟩` is `Exists.choose` of `unique_linear_decomposition_succ`
    -- which is actually same as the definition of the `split` function
    have h_must_eq := h_p_satisfy.2 (⟨hi_btf, lo_btf⟩)
    simp only [eq_mp_eq_cast] at h_must_eq
    have h_hibtf_lobtf_eq_xPair := h_must_eq (by
      rw! (castMode := .all) [h_k_sub_1_add_1_eq_k]
      simp only [cast_eq]
      convert h_x_eq_join
      · rw [eqRec_eq_cast]; rfl
      · rw [eqRec_eq_cast]; rfl
    )
    have h_split_eq_xPair : split k h_pos x = xPair := by rfl
    rw [h_split_eq_xPair, h_hibtf_lobtf_eq_xPair.symm]
  · intro h_split_eq
    unfold split at h_split_eq
    have h_hibtf_lobtf_eq_xPair : ⟨hi_btf, lo_btf⟩ = xPair := by rw [←h_split_eq]
    have h_xPair_satisfy_join_via_add_smul := h_p_satisfy.1
    rw [←h_hibtf_lobtf_eq_xPair] at h_xPair_satisfy_join_via_add_smul
    rw [eq_mp_eq_cast] at h_xPair_satisfy_join_via_add_smul
    rw! (castMode := .all) [h_k_sub_1_add_1_eq_k] at h_xPair_satisfy_join_via_add_smul
    simp only [cast_eq] at h_xPair_satisfy_join_via_add_smul
    convert h_xPair_satisfy_join_via_add_smul
    · rw [eqRec_eq_cast]; rfl
    · rw [eqRec_eq_cast]; rfl