κ-improvement development, part 4: array model and networks

-- OpenAI-preprint / κ-improvement development (integer-multiplication-kappa), generated by skeleton subtraction.

import Definitions.Def_IntMul_Kappa_K_Defs3

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.NetComplexPhase
/-!
# Proposition 8 (binary phase interface): the algebra over `ℂ`

1. The `2 × 2` matrices `H = (1/√2)[[1,1],[1,-1]]`, `C = H diag(1,i) H = aI + bX`
   (`a = (1+i)/2`, `b = (1-i)/2`), `Z = diag(1,-1)`; `C⁻¹ = āI + b̄X`; `C = i Z C⁻¹ Z`,
   `C⁻¹ = -i Z C Z`; `H Z H = X`.
2. Binary weights: `wt(z + w) = wt z + wt w - 2 |supp z ∩ supp w|`; if `z · w = 0` over `𝔽₂`,
   weight is additive mod 4; for a family orthogonal to `y` and pairwise orthogonal,
   `wt(y + ∑ cⱼ vⱼ) ≡ wt y + ∑ [cⱼ] wt vⱼ (mod 4)`; `2[u·x] ≡ 2 ∑_{j ∈ supp u} xⱼ (mod 4)`.
3. The phase identity `i^{ε[b]} = (1+iε)/2 + (1-iε)/2 · (-1)^b` for `ε = ±1`.
4. The normalized Walsh–Hadamard transform `WH` on functions `(Fin m → ZMod 2) → ℂ` (the
   matrix `H^{⊗m}`, `WH_eq_tensor`), `WH ∘ WH = id`, the general tensor-conjugation formula
   `WH (∏ᵢ dᵢ(xᵢ) · WH g) = (⊗ᵢ Kᵢ) g`, and its instances:
   * `WH_char`: `W diag((-1)^{v·x}) W⁻¹ = X_v` (translation by `v`);
   * `WH_phase`: `W diag(i^{ε[v·x]}) W⁻¹ = (1+iε)/2 · I + (1-iε)/2 · X_v`, i.e. `aI + bX_v` for
     `ε = 1` and its inverse `āI + b̄X_v` for `ε = -1`;
   * `WH_full`: `W diag(i^{wt x}) W⁻¹ = C^{⊗m}` (the frame of the full label);
   * `WH_mixed`: `W diag(i^{wt x + 2[u·x]}) W⁻¹ = (C⁻¹ on supp u) ⊗ (C elsewhere)`, and with
     `Z` conjugation on `supp u` and the scalar `i^{|supp u|}` this becomes `C^{⊗m}` again.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace NetComplexPhase

open Matrix Finset Complex

/-! ## 1. The `2 × 2` matrices -/

lemma sqrt2_sq : ((Real.sqrt 2 : ℝ) : ℂ) * ((Real.sqrt 2 : ℝ) : ℂ) = 2 := by
  rw [← Complex.ofReal_mul, Real.mul_self_sqrt (by norm_num : (0 : ℝ) ≤ 2)]
  norm_num

noncomputable def ca : ℂ := (1 + I) / 2
noncomputable def cb : ℂ := (1 - I) / 2

/-! ## 2. Binary weights -/

lemma zmod2_cases (a : ZMod 2) : a = 0 ∨ a = 1 := by revert a; decide

@[simp] lemma val0 : (0 : ZMod 2).val = 0 := rfl
@[simp] lemma val1 : (1 : ZMod 2).val = 1 := rfl

section Weight

variable {m : ℕ}

/-- Support of a binary vector. -/
def supp (z : Fin m → ZMod 2) : Finset (Fin m) := univ.filter fun i => z i ≠ 0

/-- Integer Hamming weight. -/
def wt (z : Fin m → ZMod 2) : ℕ := #(supp z)

lemma ite_ne_zero_val (a : ZMod 2) : (if a ≠ 0 then 1 else 0 : ℕ) = a.val := by
  fin_cases a <;> rfl

lemma wt_eq_sum (z : Fin m → ZMod 2) : NetComplexPhase.wt z = ∑ i, (z i).val := by
  unfold NetComplexPhase.wt supp
  rw [Finset.card_filter]
  exact Finset.sum_congr rfl fun i _ => ite_ne_zero_val _

lemma card_supp_inter (z w : Fin m → ZMod 2) :
    #(supp z ∩ supp w) = ∑ i, (z i).val * (w i).val := by
  unfold supp
  rw [← Finset.filter_and, Finset.card_filter]
  refine Finset.sum_congr rfl fun i _ => ?_
  generalize z i = a; generalize w i = b
  fin_cases a <;> fin_cases b <;> rfl

lemma val_add_zmod2 (a b : ZMod 2) :
    ((a + b).val : ℤ) = a.val + b.val - 2 * (a.val * b.val) := by
  fin_cases a <;> fin_cases b <;> rfl

/-- `wt(z + w) = wt z + wt w - 2 |supp z ∩ supp w|`. -/
theorem wt_add (z w : Fin m → ZMod 2) :
    (NetComplexPhase.wt (z + w) : ℤ) = NetComplexPhase.wt z + NetComplexPhase.wt w - 2 * #(supp z ∩ supp w) := by
  simp only [wt_eq_sum, card_supp_inter, Pi.add_apply]
  push_cast
  simp only [val_add_zmod2, Finset.sum_sub_distrib, Finset.sum_add_distrib, Finset.mul_sum]

lemma card_supp_inter_cast (z w : Fin m → ZMod 2) :
    ((#(supp z ∩ supp w) : ℕ) : ZMod 2) = z ⬝ᵥ w := by
  rw [card_supp_inter, dotProduct]
  push_cast
  simp only [ZMod.natCast_val, ZMod.cast_id', id]

/-- If `z · w = 0` over `𝔽₂`, `|supp z ∩ supp w|` is even. -/
theorem even_inter_of_dot (z w : Fin m → ZMod 2) (h : z ⬝ᵥ w = 0) :
    Even #(supp z ∩ supp w) := by
  rw [← ZMod.natCast_eq_zero_iff_even, card_supp_inter_cast, h]

/-- If `z · w = 0` over `𝔽₂`, weight is additive modulo four. -/
theorem wt_add_mod4 (z w : Fin m → ZMod 2) (h : z ⬝ᵥ w = 0) :
    ((NetComplexPhase.wt (z + w) : ℕ) : ZMod 4) = NetComplexPhase.wt z + NetComplexPhase.wt w := by
  obtain ⟨k, hk⟩ := even_inter_of_dot z w h
  have e := wt_add z w
  rw [hk] at e
  have e' := congrArg (Int.cast : ℤ → ZMod 4) e
  push_cast at e'
  have h4 : (2 * ((k : ZMod 4) + k)) = 0 := by
    rw [← two_mul, ← mul_assoc]
    rw [show (2 * 2 : ZMod 4) = 0 from rfl, zero_mul]
  rw [h4, sub_zero] at e'
  exact_mod_cast e'

lemma wt_smul (c : ZMod 2) (v : Fin m → ZMod 2) : NetComplexPhase.wt (c • v) = c.val * NetComplexPhase.wt v := by
  rcases NetComplexPhase.zmod2_cases c with rfl | rfl
  · simp [NetComplexPhase.wt, supp]
  · simp

/-- Orthogonal expansion: if the `vⱼ` are pairwise orthogonal and orthogonal to `y`, then
`wt(y + ∑ cⱼ vⱼ) ≡ wt y + ∑ [cⱼ] wt vⱼ (mod 4)`. With `y = P_U x`, `cⱼ = vⱼ · x` for an
orthonormal basis of the residual `E` (`V = U ⊥ E`), this is
`q_V(x) - q_U(x) ≡ ∑_{v ∈ B} wt(v) [v · x] (mod 4)`. -/
theorem wt_orth_sum {ι : Type*} [DecidableEq ι] (y : Fin m → ZMod 2) (s : Finset ι)
    (v : ι → Fin m → ZMod 2) (c : ι → ZMod 2) (hy : ∀ j ∈ s, y ⬝ᵥ v j = 0)
    (hv : ∀ j ∈ s, ∀ k ∈ s, j ≠ k → v j ⬝ᵥ v k = 0) :
    ((NetComplexPhase.wt (y + ∑ j ∈ s, c j • v j) : ℕ) : ZMod 4) =
      NetComplexPhase.wt y + ((∑ j ∈ s, (c j).val * NetComplexPhase.wt (v j) : ℕ) : ZMod 4) := by
  induction s using Finset.induction_on with
  | empty => simp
  | insert j s hj ih =>
    rw [Finset.sum_insert hj, Finset.sum_insert hj,
      show y + (c j • v j + ∑ k ∈ s, c k • v k) = (y + ∑ k ∈ s, c k • v k) + c j • v j by abel,
      wt_add_mod4, ih (fun k hk => hy k (Finset.mem_insert_of_mem hk))
        (fun a ha b hb => hv a (Finset.mem_insert_of_mem ha) b (Finset.mem_insert_of_mem hb)),
      wt_smul]
    · push_cast; ring
    · rw [add_dotProduct, dotProduct_smul, hy j (Finset.mem_insert_self j s), smul_zero,
        zero_add, sum_dotProduct]
      refine Finset.sum_eq_zero fun k hk => ?_
      rw [smul_dotProduct, dotProduct_smul, hv k (Finset.mem_insert_of_mem hk) j
        (Finset.mem_insert_self j s) (fun e => hj (e ▸ hk)), smul_zero, smul_zero]

lemma two_mul_eq_of_mod (a b : ℕ) (h : a % 2 = b % 2) : (2 * a : ZMod 4) = 2 * b := by
  have ha := Nat.div_add_mod a 2
  have hb := Nat.div_add_mod b 2
  rw [← ha, ← hb, h]
  push_cast
  ring_nf
  rw [show (4 : ZMod 4) = 0 from rfl]
  ring

/-- `2[u · x] ≡ 2 ∑_{j ∈ supp u} xⱼ (mod 4)`. -/
theorem two_dot_mod4 (u x : Fin m → ZMod 2) :
    (2 * (u ⬝ᵥ x).val : ZMod 4) = 2 * ((∑ j ∈ supp u, (x j).val : ℕ) : ZMod 4) := by
  have := two_mul_eq_of_mod (u ⬝ᵥ x).val (∑ j ∈ supp u, (x j).val) ?_
  · exact_mod_cast this
  rw [← ZMod.natCast_eq_natCast_iff']
  push_cast
  simp only [ZMod.natCast_val, ZMod.cast_id', id, dotProduct, supp]
  rw [Finset.sum_filter]
  refine Finset.sum_congr rfl fun j _ => ?_
  rcases NetComplexPhase.zmod2_cases (u j) with h | h <;> simp [h]

end Weight

/-! ## 3. The phase identity -/

/-- `i^{ε[b]} = (1+iε)/2 + (1-iε)/2 · (-1)^b` for `ε = ±1`. -/
theorem phase_identity (ε : ℤ) (hε : ε = 1 ∨ ε = -1) (b : ZMod 2) :
    I ^ (ε * (b.val : ℤ)) = (1 + I * ε) / 2 + (1 - I * ε) / 2 * (-1) ^ b.val := by
  rcases hε with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases b with rfl | rfl <;> simp <;> ring_nf

/-! ## 4. The Walsh–Hadamard transform `H^{⊗m}` and its conjugations -/

section WH

variable {m : ℕ}

/-- The character `(-1)^b`. -/
def chi (b : ZMod 2) : ℂ := (-1) ^ b.val

lemma chi_add (a b : ZMod 2) : chi (a + b) = chi a * chi b := by
  rcases NetComplexPhase.zmod2_cases a with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases b with rfl | rfl <;>
    simp [chi] <;> rfl

lemma chi_sum {ι : Type*} (s : Finset ι) (f : ι → ZMod 2) :
    chi (∑ i ∈ s, f i) = ∏ i ∈ s, chi (f i) := by
  classical
  induction s using Finset.induction_on with
  | empty => simp [chi]
  | insert j s hj ih => rw [Finset.sum_insert hj, Finset.prod_insert hj, chi_add, ih]

lemma chi_dot (x y : Fin m → ZMod 2) : chi (x ⬝ᵥ y) = ∏ i, chi (x i * y i) :=
  chi_sum _ _

lemma sum_zmod2 (f : ZMod 2 → ℂ) : ∑ a, f a = f 0 + f 1 := Fin.sum_univ_two f

/-- `√2` as a complex number. -/
noncomputable def s2 : ℂ := ((Real.sqrt 2 : ℝ) : ℂ)

/-- The normalized Walsh–Hadamard transform: `(W f)(y) = 2^{-m/2} ∑ₓ (-1)^{x·y} f(x)`. -/
noncomputable def WH (f : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) : ℂ :=
  (s2 ^ m)⁻¹ * ∑ x, chi (x ⬝ᵥ y) * f x

noncomputable def CK (a b : ZMod 2) : ℂ := if a = b then NetComplexPhase.ca else NetComplexPhase.cb
noncomputable def CKinv (a b : ZMod 2) : ℂ := if a = b then (1 - I) / 2 else (1 + I) / 2

lemma norm_sq : (s2 ^ m)⁻¹ * (s2 ^ m)⁻¹ = ∏ _i : Fin m, (1 / 2 : ℂ) := by
  rw [Finset.prod_const, Finset.card_univ, Fintype.card_fin, ← mul_inv, ← mul_pow, s2,
    sqrt2_sq, one_div, inv_pow]

/-- One-coordinate kernel of `H diag(δ) H`: `(1/2) ∑ₐ (-1)^{ya} δ(a) (-1)^{az}`. -/
noncomputable def K1 (δ : ZMod 2 → ℂ) (y z : ZMod 2) : ℂ :=
  (1 / 2) * ∑ a, chi (y * a) * δ a * chi (a * z)

/-- **Tensor conjugation.** Conjugating a product diagonal `∏ᵢ dᵢ(xᵢ)` by `W = H^{⊗m}` gives
the tensor product of the one-coordinate kernels `H diag(dᵢ) H`. -/
theorem WH_conj (d : Fin m → ZMod 2 → ℂ) (g : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => (∏ i, d i (x i)) * WH g x) y = ∑ z, (∏ i, K1 (d i) (y i) (z i)) * g z := by
  have hK : ∀ z : Fin m → ZMod 2, ∏ i, K1 (d i) (y i) (z i) =
      (s2 ^ m)⁻¹ * (s2 ^ m)⁻¹ * ∑ x, chi (x ⬝ᵥ y) * (∏ i, d i (x i)) * chi (z ⬝ᵥ x) := by
    intro z
    rw [norm_sq]
    simp only [K1, Finset.prod_mul_distrib]
    congr 1
    rw [Finset.prod_univ_sum, Fintype.piFinset_univ]
    refine Finset.sum_congr rfl fun x _ => ?_
    rw [chi_dot, chi_dot, Finset.prod_mul_distrib, Finset.prod_mul_distrib]
    congr 1
    · congr 1
      exact Finset.prod_congr rfl fun i _ => by rw [mul_comm]
    · exact Finset.prod_congr rfl fun i _ => by rw [mul_comm]
  simp only [hK, WH, Finset.mul_sum, Finset.sum_mul]
  rw [Finset.sum_comm]
  refine Finset.sum_congr rfl fun z _ => Finset.sum_congr rfl fun x _ => ?_
  ring

lemma K1_char (c y z : ZMod 2) :
    K1 (fun a => chi (c * a)) y z = if z = y + c then 1 else 0 := by
  rcases NetComplexPhase.zmod2_cases c with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases y with rfl | rfl <;>
    rcases NetComplexPhase.zmod2_cases z with rfl | rfl <;> simp [K1, sum_zmod2, chi] <;> norm_num <;> rfl

/-- **Character to translation.** `W diag((-1)^{v·x}) W⁻¹ = X_v` (`W⁻¹ = W`, see `WH_WH`). -/
theorem WH_char (v : Fin m → ZMod 2) (g : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => chi (v ⬝ᵥ x) * WH g x) y = g (y + v) := by
  have := WH_conj (fun i a => chi (v i * a)) g y
  simp only [← chi_dot, K1_char, Fintype.prod_boole] at this
  rw [this]
  have e : ∀ z : Fin m → ZMod 2, (∀ i, z i = y i + v i) ↔ z = y + v := fun z =>
    ⟨fun h => funext h, fun h i => by rw [h]; rfl⟩
  simp only [e, ite_mul, one_mul, zero_mul, Finset.sum_ite_eq', Finset.mem_univ, if_true]

/-- `W ∘ W = id` (so `W⁻¹ = W`). -/
theorem WH_WH (g : (Fin m → ZMod 2) → ℂ) : WH (WH g) = g := by
  funext y
  have := WH_char 0 g y
  simp only [zero_dotProduct, chi, val0, pow_zero, one_mul, add_zero] at this
  exact this

lemma WH_lin (α β : ℂ) (f₁ f₂ : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => α * f₁ x + β * f₂ x) y = α * WH f₁ y + β * WH f₂ y := by
  simp only [WH, mul_add, Finset.sum_add_distrib, Finset.mul_sum]
  congr 1 <;> refine Finset.sum_congr rfl fun x _ => ?_ <;> ring

/-- **Translation kernel (Proposition 8).** For `ε = ±1`,
`W diag(i^{ε[v·x]}) W⁻¹ = (1+iε)/2 · I + (1-iε)/2 · X_v`: this is `aI + bX_v` for `ε = 1` and
its inverse `āI + b̄X_v` for `ε = -1`. -/
theorem WH_phase (ε : ℤ) (hε : ε = 1 ∨ ε = -1) (v : Fin m → ZMod 2)
    (g : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => I ^ (ε * ((v ⬝ᵥ x).val : ℤ)) * WH g x) y =
      (1 + I * ε) / 2 * g y + (1 - I * ε) / 2 * g (y + v) := by
  simp only [phase_identity ε hε]
  have : (fun x => ((1 + I * ε) / 2 + (1 - I * ε) / 2 * (-1) ^ (v ⬝ᵥ x).val) * WH g x) =
      fun x => (1 + I * ε) / 2 * WH g x + (1 - I * ε) / 2 * (chi (v ⬝ᵥ x) * WH g x) := by
    funext x; simp only [chi]; ring
  rw [this, WH_lin, WH_WH, WH_char]

/-- Entries of `C` and `C⁻¹` indexed by `ZMod 2 = Fin 2`. -/
lemma K1_I (y z : ZMod 2) : K1 (fun a => I ^ a.val) y z = CK y z := by
  rcases NetComplexPhase.zmod2_cases y with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases z with rfl | rfl <;>
    simp [K1, CK, NetComplexPhase.ca, NetComplexPhase.cb, sum_zmod2, chi] <;> ring_nf

lemma K1_I3 (y z : ZMod 2) : K1 (fun a => I ^ (a.val * 3)) y z = CKinv y z := by
  have h3 : I ^ 3 = -I := by rw [pow_succ, I_sq]; ring
  rcases NetComplexPhase.zmod2_cases y with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases z with rfl | rfl <;>
    simp [K1, CKinv, sum_zmod2, chi, h3] <;> ring_nf

lemma I_pow_wt (x : Fin m → ZMod 2) : I ^ NetComplexPhase.wt x = ∏ i, I ^ (x i).val := by
  rw [wt_eq_sum, Finset.prod_pow_eq_pow_sum]

/-- **Full-label frame.** `W diag(i^{wt x}) W⁻¹ = C^{⊗m}` (the full label has `q_F = wt`). -/
theorem WH_full (g : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => I ^ NetComplexPhase.wt x * WH g x) y = ∑ z, (∏ i, CK (y i) (z i)) * g z := by
  simp only [I_pow_wt]
  rw [WH_conj (fun _ a => I ^ a.val)]
  simp only [K1_I]

lemma I_pow_mod (a b : ℕ) (h : (a : ZMod 4) = b) : I ^ a = I ^ b := by
  rw [ZMod.natCast_eq_natCast_iff'] at h
  rw [← Nat.div_add_mod a 4, ← Nat.div_add_mod b 4, h, pow_add, pow_add, pow_mul, pow_mul,
    show I ^ 4 = 1 by rw [show (4 : ℕ) = 2 * 2 from rfl, pow_mul, I_sq]; ring]
  simp

/-- **Mixed frame (the pair `X_a → Y_a`).** With `q_{U⊥} - q_U ≡ wt(x) + 2[u·x] (mod 4)`, the
conjugated phase is `C⁻¹` on the coordinates of `supp u` and `C` elsewhere. -/
theorem WH_mixed (u : Fin m → ZMod 2) (g : (Fin m → ZMod 2) → ℂ) (y : Fin m → ZMod 2) :
    WH (fun x => I ^ (NetComplexPhase.wt x + 2 * (u ⬝ᵥ x).val) * WH g x) y =
      ∑ z, (∏ i, (if i ∈ supp u then CKinv else CK) (y i) (z i)) * g z := by
  have hph : ∀ x : Fin m → ZMod 2, I ^ (NetComplexPhase.wt x + 2 * (u ⬝ᵥ x).val) =
      ∏ i, I ^ ((x i).val * (if i ∈ supp u then 3 else 1)) := by
    intro x
    rw [Finset.prod_pow_eq_pow_sum]
    apply I_pow_mod
    push_cast
    rw [two_dot_mod4, wt_eq_sum]
    push_cast
    rw [← Finset.univ_inter (supp u), ← Finset.sum_ite_mem, Finset.mul_sum,
      ← Finset.sum_add_distrib]
    refine Finset.sum_congr rfl fun i _ => ?_
    simp only [Finset.mem_inter, Finset.mem_univ, true_and]
    split_ifs <;> ring
  simp only [hph]
  rw [WH_conj (fun i a => I ^ (a.val * (if i ∈ supp u then 3 else 1)))]
  refine Finset.sum_congr rfl fun z _ => ?_
  congr 1
  refine Finset.prod_congr rfl fun i _ => ?_
  split_ifs
  · exact K1_I3 _ _
  · simpa using K1_I _ _

/-- Scalar content of `C = i Z C⁻¹ Z` on `ZMod 2` indices (`Z = diag(chi)`). -/
lemma CK_eq_IZCKinvZ (a b : ZMod 2) : I * (chi a * CKinv a b * chi b) = CK a b := by
  rcases NetComplexPhase.zmod2_cases a with rfl | rfl <;> rcases NetComplexPhase.zmod2_cases b with rfl | rfl <;>
    simp [CK, CKinv, NetComplexPhase.ca, NetComplexPhase.cb, chi] <;> ring_nf <;> simp [I_sq] <;> ring

/-- **Z-correction.** Applying `Z` on the coordinates of `supp u` before and after, and the
scalar `i^{|supp u|}` once, turns the mixed kernel back into `C^{⊗m}`:
`i^{wt u} · ∏ᵢ zᵢ(yᵢ) Kᵢ(yᵢ,zᵢ) zᵢ(zᵢ) = ∏ᵢ C(yᵢ,zᵢ)`. -/
theorem Z_correction (u y z : Fin m → ZMod 2) :
    I ^ NetComplexPhase.wt u * ∏ i, ((if i ∈ supp u then chi (y i) else 1) *
        (if i ∈ supp u then CKinv else CK) (y i) (z i) *
        (if i ∈ supp u then chi (z i) else 1)) = ∏ i, CK (y i) (z i) := by
  have hI : I ^ NetComplexPhase.wt u = ∏ i, (if i ∈ supp u then I else 1) := by
    rw [Finset.prod_ite, Finset.prod_const_one, mul_one, Finset.prod_const, NetComplexPhase.wt]
    congr 2
    ext; simp
  rw [hI, ← Finset.prod_mul_distrib]
  refine Finset.prod_congr rfl fun i _ => ?_
  split_ifs
  · exact CK_eq_IZCKinvZ _ _
  · ring

/-- The translation kernel `K_{ε,v} = (1+iε)/2 · I + (1-iε)/2 · X_v` (`aI + bX_v` for `ε = 1`,
its inverse for `ε = -1`). -/
noncomputable def transK (ε : ℤ) (v : Fin m → ZMod 2) (g : (Fin m → ZMod 2) → ℂ)
    (y : Fin m → ZMod 2) : ℂ :=
  (1 + I * ε) / 2 * g y + (1 - I * ε) / 2 * g (y + v)

/-- **Proposition 8, per edge.** If the phase-frame difference of an edge is
`i^{∑_{(ε,v) ∈ L} ε [v·x]}` (`ε = ±1`, one entry per residual basis vector), then its conjugation
by `W = H^{⊗m}` is the product of the translation kernels `K_{ε,v}`, one per entry. -/
theorem WH_phase_list (L : List (ℤ × (Fin m → ZMod 2))) (hL : ∀ p ∈ L, p.1 = 1 ∨ p.1 = -1)
    (g : (Fin m → ZMod 2) → ℂ) :
    WH (fun x => (L.map fun p => I ^ (p.1 * ((p.2 ⬝ᵥ x).val : ℤ))).prod * WH g x) =
      L.foldr (fun p g' => transK p.1 p.2 g') g := by
  induction L with
  | nil => simp [WH_WH]
  | cons p L ih =>
    have ih' := ih (fun q hq => hL q (List.mem_cons_of_mem _ hq))
    funext y
    simp only [List.map_cons, List.prod_cons, List.foldr_cons]
    have e : (fun x => I ^ (p.1 * ((p.2 ⬝ᵥ x).val : ℤ)) *
        (L.map fun p => I ^ (p.1 * ((p.2 ⬝ᵥ x).val : ℤ))).prod * WH g x) =
        fun x => I ^ (p.1 * ((p.2 ⬝ᵥ x).val : ℤ)) * WH (WH (fun x =>
          (L.map fun p => I ^ (p.1 * ((p.2 ⬝ᵥ x).val : ℤ))).prod * WH g x)) x := by
      funext x; rw [WH_WH, mul_assoc]
    rw [e, WH_phase p.1 (hL p List.mem_cons_self), ih']
    rfl

end WH

end NetComplexPhase

end

section
-- Source module: Solutions.P8LiftCol
/-!
# Column tensor powers, wire operators and the translation-kernel chunks

* `colT A`: the operator `A^{⊗κ}` on `κ` columns of `m` address bits (§3.5, last paragraph:
  for `k` columns the frame is the `k`-th tensor power). `colT` is multiplicative.
* `blockOp`, `wireOp`: role-block-diagonal operators; restricted steps as wire operators.
* `Tk ε v`: the matrix of the translation kernel `transK ε v = aI + bX_v` (or its inverse).
* `colT_Tk_rowAdd`: conjugating by a binary row addition `M` turns `X_v` into `X_{M v}`
  (columnwise, (15)).
* `colT_Tk_single`: `(aI + bX_{e_h})^{⊗κ}` is `C^{⊗κ}` on slot `h` (`ε = 1`) or `(C⁻¹)^{⊗κ}`
  (`ε = -1`), identity on the other slots.
* `chunk_den`: the restricted step chunk "row additions sending `v` to `e_h`; one child call on
  slot `h`; the row additions undone" computes `(transK ε v)^{⊗κ}` on its wire.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace P8Lift

open Matrix Finset Complex Butterfly

variable {m : ℕ} {κ : Type} [Fintype κ] [DecidableEq κ]

/-- Single-column operators. -/
abbrev M1 (m : ℕ) := Matrix (Fin m → ZMod 2) (Fin m → ZMod 2) ℂ

/-- Column `c` of an address. -/
def col (x : Addr m κ) (c : κ) : Fin m → ZMod 2 := fun r => x (r, c)

/-- Addresses as families of columns. -/
def colEquiv : Addr m κ ≃ (κ → Fin m → ZMod 2) where
  toFun x c := col x c
  invFun g p := g p.2 p.1
  left_inv x := rfl
  right_inv g := rfl

lemma col_ext {x y : Addr m κ} : (∀ c, col x c = col y c) ↔ x = y := by
  constructor
  · intro h; funext p; exact congrFun (h p.2) p.1
  · rintro rfl c; rfl

/-- The column tensor power `A^{⊗κ}`. -/
def colT (A : M1 m) : Op (Fin m × κ) := fun x y => ∏ c, A (col x c) (col y c)

theorem colT_mul (A B : M1 m) : (colT (A * B) : Op (Fin m × κ)) = colT A * colT B := by
  ext x z
  simp only [colT, Matrix.mul_apply]
  rw [Fintype.prod_sum]
  rw [← (colEquiv (m := m) (κ := κ)).sum_comp]
  refine Finset.sum_congr rfl fun y _ => ?_
  rw [← Finset.prod_mul_distrib]
  rfl

theorem colT_one : (colT (1 : M1 m) : Op (Fin m × κ)) = 1 := by
  ext x y
  simp only [colT, one_apply]
  rw [Finset.prod_boole]
  simp only [mem_univ, true_implies]
  by_cases h : x = y
  · simp [h]
  · rw [if_neg (fun h' => h (col_ext.1 h')), if_neg h]

theorem colT_smul (a : ℂ) (A : M1 m) :
    (colT (a • A) : Op (Fin m × κ)) = a ^ Fintype.card κ • colT A := by
  ext x y
  simp only [colT, Matrix.smul_apply, smul_eq_mul, Finset.prod_mul_distrib, Finset.prod_const,
    card_univ]

theorem colT_diagonal (d : (Fin m → ZMod 2) → ℂ) :
    (colT (diagonal d) : Op (Fin m × κ)) = diagonal fun x => ∏ c, d (col x c) := by
  ext x y
  simp only [colT, diagonal_apply]
  by_cases h : x = y
  · subst h; simp
  · rw [if_neg h]
    obtain ⟨c, hc⟩ : ∃ c, col x c ≠ col y c := by
      by_contra hc; push Not at hc; exact h (col_ext.1 hc)
    exact Finset.prod_eq_zero (mem_univ c) (if_neg hc)

/-! ## Block-diagonal operators on roles -/

section Block

variable {R : Type} [Fintype R] [DecidableEq R]

/-- Role-block-diagonal operator. -/
def blockOp (A : R → Op (Fin m × κ)) : ROp R m κ := fun p q => if p.1 = q.1 then A p.1 p.2 q.2 else 0

theorem blockOp_mul (A B : R → Op (Fin m × κ)) :
    blockOp A * blockOp B = blockOp fun r => A r * B r := by
  ext p q
  simp only [blockOp, Matrix.mul_apply]
  rw [Fintype.sum_prod_type, Finset.sum_eq_single p.1]
  · by_cases h : p.1 = q.1
    · simp [h]
    · simp [h]
  · intro b _ hb; simp [Ne.symm hb]
  · simp

theorem blockOp_one : (blockOp fun _ : R => (1 : Op (Fin m × κ))) = 1 := by
  ext p q
  simp only [blockOp, one_apply]
  by_cases h : p.1 = q.1
  · by_cases h' : p.2 = q.2
    · simp [h, h', Prod.ext h h']
    · simp [h, h', Prod.ext_iff]
  · simp [h, Prod.ext_iff]

/-- An operator on one wire. -/
def wireOp (w : R) (A : Op (Fin m × κ)) : ROp R m κ := blockOp fun r => if r = w then A else 1

theorem wireOp_mul (w : R) (A B : Op (Fin m × κ)) :
    wireOp w A * wireOp w B = wireOp w (A * B) := by
  rw [wireOp, wireOp, blockOp_mul, wireOp]
  congr 1; funext r; split_ifs <;> simp

theorem wireOp_one (w : R) : wireOp w (1 : Op (Fin m × κ)) = 1 := by
  rw [wireOp, ← blockOp_one]; congr 1; funext r; split_ifs <;> rfl

/-- Address permutation by a row addition. -/
def Pm (i j : Fin m) : Op (Fin m × κ) := fun x y => if y = rowAddr i j x then 1 else 0

/-- The slot operator: `T` on the bits of slot `h`, identity on the other slots. -/
def slotOp (h : Fin m) (T : Op κ) : Op (Fin m × κ) := fun x y =>
  if ∀ h', h' ≠ h → ∀ j, x (h', j) = y (h', j) then T (slotBits x h) (slotBits y h) else 0

theorem den_rowAdd (T : Op κ) (w : R) (i j : Fin m) (hij : i ≠ j) :
    (RStep.rowAdd w i j hij).den T = wireOp w (Pm i j) := by
  ext ⟨p1, p2⟩ ⟨q1, q2⟩
  simp only [RStep.den, wireOp, blockOp, Pm, one_apply]
  by_cases h1 : p1 = q1
  · subst h1; by_cases h2 : p1 = w <;> simp [h2, Pm, slotOp, one_apply]
  · simp [h1]

theorem childR_eq (w : R) (h : Fin m) (T : Op κ) :
    RStep.childR w h T = wireOp w (slotOp h T) := by
  ext ⟨p1, p2⟩ ⟨q1, q2⟩
  simp only [RStep.childR, wireOp, blockOp, slotOp, one_apply]
  by_cases h1 : p1 = q1
  · subst h1; by_cases h2 : p1 = w <;> simp [h2, Pm, slotOp, one_apply]
  · simp [h1]

theorem den_child_true (w : R) (h : Fin m) :
    (RStep.child w h true : RStep R m).den (Butterfly.tens Butterfly.Cm : Op κ) = wireOp w (slotOp h (Butterfly.tens Butterfly.Cm)) :=
  childR_eq w h _

theorem den_child_false (w : R) (h : Fin m) :
    (RStep.child w h false : RStep R m).den (Butterfly.tens Butterfly.Cm : Op κ) = wireOp w (slotOp h (Butterfly.tens Cinv)) := by
  rw [RStep.den, ← tens_Cinv', childR_eq]

end Block

/-! ## Translation kernels and row additions -/

/-- The matrix of `transK ε v`. -/
noncomputable def Tk (ε : ℤ) (v : Fin m → ZMod 2) : M1 m := fun x y =>
  (1 + I * ε) / 2 * (if y = x then 1 else 0) + (1 - I * ε) / 2 * (if y = x + v then 1 else 0)

theorem Tk_mulVec (ε : ℤ) (v : Fin m → ZMod 2) (g : (Fin m → ZMod 2) → ℂ) :
    Tk ε v *ᵥ g = NetComplexPhase.transK ε v g := by
  funext x
  simp only [mulVec, dotProduct, Tk, add_mul, Finset.sum_add_distrib, mul_assoc,
    ← Finset.mul_sum, ite_mul, one_mul, zero_mul, Finset.sum_ite_eq', mem_univ, if_true,
    NetComplexPhase.transK]

/-- The single-column row addition `z ↦ z + z_i e_j`. -/
def ra (i j : Fin m) (z : Fin m → ZMod 2) : Fin m → ZMod 2 := fun r => if r = j then z r + z i else z r

lemma col_rowAddr (i j : Fin m) (x : Addr m κ) (c : κ) :
    col (rowAddr i j x) c = ra i j (col x c) := rfl

lemma zmod2_add_self (a : ZMod 2) : a + a = 0 := by revert a; decide

lemma ra_invol {i j : Fin m} (hij : i ≠ j) (z : Fin m → ZMod 2) : ra i j (ra i j z) = z := by
  funext r; simp only [ra]
  by_cases hr : r = j
  · simp only [hr, if_true, if_neg hij, add_assoc, P8Lift.zmod2_add_self, add_zero]
  · simp [hr]

lemma ra_add (i j : Fin m) (a b : Fin m → ZMod 2) : ra i j (a + b) = ra i j a + ra i j b := by
  funext r; simp only [ra, Pi.add_apply]; split_ifs <;> ring

lemma ra_inj {i j : Fin m} (hij : i ≠ j) {a b : Fin m → ZMod 2} : ra i j a = ra i j b ↔ a = b :=
  ⟨fun h => by rw [← ra_invol hij a, h, ra_invol hij], fun h => h ▸ rfl⟩

lemma Tk_ra {i j : Fin m} (hij : i ≠ j) (ε : ℤ) (v a b : Fin m → ZMod 2) :
    Tk ε v (ra i j a) (ra i j b) = Tk ε (ra i j v) a b := by
  simp only [Tk]
  have e1 : (ra i j b = ra i j a) ↔ b = a := ra_inj hij
  have e2 : (ra i j b = ra i j a + v) ↔ b = a + ra i j v := by
    constructor
    · intro h
      have := congrArg (ra i j) h
      rwa [ra_invol hij, ra_add, ra_invol hij] at this
    · intro h; rw [h, ra_add, ra_invol hij]
  simp only [e1, e2]

theorem Pm_conj {i j : Fin m} (hij : i ≠ j) (B : Op (Fin m × κ)) :
    (Pm i j * B * Pm i j : Op (Fin m × κ)) = fun x z => B (rowAddr i j x) (rowAddr i j z) := by
  ext x z
  simp only [Matrix.mul_apply, Pm, ite_mul, one_mul, zero_mul, Finset.sum_ite_eq', mem_univ,
    if_true, mul_ite, mul_one, mul_zero]
  have : ∀ y : Addr m κ, (z = rowAddr i j y) ↔ (y = rowAddr i j z) := fun y =>
    ⟨fun h => by rw [h, rowAddr_invol hij], fun h => by rw [h, rowAddr_invol hij]⟩
  simp only [this, Finset.sum_ite_eq', mem_univ, if_true]

theorem colT_Tk_rowAdd {i j : Fin m} (hij : i ≠ j) (ε : ℤ) (v : Fin m → ZMod 2) :
    Pm i j * (colT (Tk ε v) : Op (Fin m × κ)) * Pm i j = colT (Tk ε (ra i j v)) := by
  rw [Pm_conj hij]
  ext x z
  simp only [colT, col_rowAddr, Tk_ra hij]

lemma Cm_apply' (s t : Fin 2) : Butterfly.Cm s t = if s = t then ca else cb := by
  fin_cases s <;> fin_cases t <;> simp [Butterfly.Cm, Xm]

lemma Cinv_apply' (s t : Fin 2) : Cinv s t = if s = t then cb else ca := by
  fin_cases s <;> fin_cases t <;> simp [Cinv, Xm]

lemma zmod2_eq_or (a b : ZMod 2) : a = b ∨ b = a + 1 := by revert a b; decide

lemma zmod2_ne (a : ZMod 2) : a ≠ a + 1 := by revert a; decide

/-- One column of `aI + bX_{e_h}`: identity off `h`, the `2 × 2` kernel on bit `h`. -/
lemma Tk_single_col (ε : ℤ) (h : Fin m) (a b : Fin m → ZMod 2) :
    Tk ε (Pi.single h 1) a b = (if ∀ r, r ≠ h → a r = b r then 1 else 0) *
      (if a h = b h then (1 + I * ε) / 2 else (1 - I * ε) / 2) := by
  simp only [Tk]
  by_cases hoff : ∀ r, r ≠ h → a r = b r
  · rw [if_pos hoff, one_mul]
    rcases zmod2_eq_or (a h) (b h) with e | e
    · have hba : b = a := funext fun r => by
        by_cases hr : r = h
        · rw [hr, e]
        · exact (hoff r hr).symm
      have hne : b ≠ a + Pi.single h 1 := fun h' => by
        have := congrFun h' h; simp [hba] at this
      rw [if_pos hba, if_neg hne, if_pos e]; ring
    · have hba : b = a + Pi.single h 1 := funext fun r => by
        by_cases hr : r = h
        · subst hr; simp [e]
        · simp [Pi.single_apply, hr, hoff r hr]
      have hne : b ≠ a := fun h' => zmod2_ne (a h) (by rw [← e, h'])
      have hne' : a h ≠ b h := fun h' => zmod2_ne (a h) (by rw [← e, h'])
      rw [if_neg hne, if_pos hba, if_neg hne']; ring
  · rw [if_neg hoff, zero_mul]
    push Not at hoff
    obtain ⟨r, hr, hab⟩ := hoff
    have h1 : b ≠ a := fun h' => hab (by rw [h'])
    have h2 : b ≠ a + Pi.single h 1 := fun h' => hab (by rw [h']; simp [Pi.single_apply, hr])
    rw [if_neg h1, if_neg h2]; ring

lemma colT_Tk_single_gen (ε : ℤ) (h : Fin m) (M : M2)
    (hM : ∀ s t : Fin 2, M s t = if s = t then (1 + I * ε) / 2 else (1 - I * ε) / 2) :
    (colT (Tk ε (Pi.single h 1)) : Op (Fin m × κ)) = slotOp h (Butterfly.tens M) := by
  ext x y
  simp only [colT, Tk_single_col, Finset.prod_mul_distrib, slotOp, Butterfly.tens, tensF, slotBits]
  rw [Finset.prod_boole]
  by_cases H : ∀ h', h' ≠ h → ∀ j, x (h', j) = y (h', j)
  · have H' : ∀ i ∈ (univ : Finset κ), ∀ r : Fin m, r ≠ h → col x i r = col y i r :=
      fun i _ r hr => H r hr i
    rw [if_pos H', if_pos H, one_mul]
    refine Finset.prod_congr rfl fun c _ => ?_
    rw [hM]; rfl
  · have H' : ¬ ∀ i ∈ (univ : Finset κ), ∀ r : Fin m, r ≠ h → col x i r = col y i r :=
      fun H' => H fun r hr i => H' i (mem_univ _) r hr
    rw [if_neg H', if_neg H, zero_mul]

theorem colT_Tk_single_pos (h : Fin m) :
    (colT (Tk 1 (Pi.single h 1)) : Op (Fin m × κ)) = slotOp h (Butterfly.tens Butterfly.Cm) :=
  colT_Tk_single_gen 1 h Butterfly.Cm fun s t => by rw [Cm_apply']; simp [ca, cb]

theorem colT_Tk_single_neg (h : Fin m) :
    (colT (Tk (-1) (Pi.single h 1)) : Op (Fin m × κ)) = slotOp h (Butterfly.tens Cinv) :=
  colT_Tk_single_gen (-1) h Cinv fun s t => by
    rw [Cinv_apply']; simp only [ca, cb]; push_cast; split_ifs <;> ring

/-! ## The kernel chunk -/

section Chunk

variable {R : Type} [Fintype R] [DecidableEq R]

/-- Row additions `x_j += x_h` for `j ∈ js` (`j ≠ h`) on wire `w`. -/
def rowAddsFor (w : R) (h : Fin m) (js : List (Fin m)) : List (RStep R m) :=
  js.flatMap fun j => if hj : h = j then [] else [RStep.rowAdd w h j hj]

/-- The effect of one of these row additions on a single column. -/
def raStep (h j : Fin m) (z : Fin m → ZMod 2) : Fin m → ZMod 2 := if h = j then z else ra h j z

/-- Their composite effect, in list order. -/
def raList (h : Fin m) (js : List (Fin m)) (z : Fin m → ZMod 2) : Fin m → ZMod 2 :=
  js.foldl (fun z j => raStep h j z) z

lemma den_rowAddsFor_single (w : R) (h j : Fin m) :
    denListR (Butterfly.tens Butterfly.Cm : Op κ) (rowAddsFor w h [j]) =
      if hj : h = j then 1 else wireOp w (Pm h j) := by
  simp only [rowAddsFor, List.flatMap_cons, List.flatMap_nil, List.append_nil]
  split_ifs with hj
  · rfl
  · rw [show [RStep.rowAdd w h j hj] = [] ++ [RStep.rowAdd w h j hj] from rfl,
      denListR_append, denListR_nil, mul_one]
    simp only [denListR, List.foldl_cons, List.foldl_nil, mul_one]
    exact den_rowAdd _ w h j hj

lemma rowAddsFor_cons (w : R) (h j : Fin m) (js : List (Fin m)) :
    rowAddsFor w h (j :: js) = rowAddsFor w h [j] ++ rowAddsFor w h js := by
  simp [rowAddsFor]

lemma rowAddsFor_append (w : R) (h : Fin m) (js js' : List (Fin m)) :
    rowAddsFor w h (js ++ js') = rowAddsFor w h js ++ rowAddsFor w h js' := by
  simp [rowAddsFor]

/-- **The kernel chunk.** If the row additions send `v` to `e_h`, then
"row additions; child call on slot `h` (forward for `ε = 1`, inverse for `ε = -1`); row
additions in reverse" computes `(transK ε v)^{⊗κ}` on wire `w`. -/
theorem chunk_den (w : R) (h : Fin m) (ε : ℤ) (hε : ε = 1 ∨ ε = -1) :
    ∀ (js : List (Fin m)) (v : Fin m → ZMod 2), raList h js v = Pi.single h 1 →
      denListR (Butterfly.tens Butterfly.Cm : Op κ) (rowAddsFor w h js ++ [RStep.child w h (decide (ε = 1))] ++
        rowAddsFor w h js.reverse) = wireOp w (colT (Tk ε v)) := by
  intro js
  induction js with
  | nil =>
    intro v hv
    simp only [raList, List.foldl_nil] at hv
    subst hv
    simp only [rowAddsFor, List.flatMap_nil, List.nil_append, List.reverse_nil, List.append_nil]
    simp only [denListR, List.foldl_cons, List.foldl_nil, mul_one]
    rcases hε with rfl | rfl
    · rw [colT_Tk_single_pos]; exact den_child_true w h
    · rw [colT_Tk_single_neg]; exact den_child_false w h
  | cons j js ih =>
    intro v hv
    have hv' : raList h js (raStep h j v) = Pi.single h 1 := hv
    have IH := ih _ hv'
    have e : rowAddsFor w h (j :: js) ++ [RStep.child w h (decide (ε = 1))] ++
        rowAddsFor w h (j :: js).reverse = rowAddsFor w h [j] ++ (rowAddsFor w h js ++
          [RStep.child w h (decide (ε = 1))] ++ rowAddsFor w h js.reverse) ++
          rowAddsFor w h [j] := by
      rw [rowAddsFor_cons, List.reverse_cons, rowAddsFor_append]; simp
    rw [e, denListR_append, denListR_append, IH, den_rowAddsFor_single]
    unfold raStep
    split_ifs with hj
    · rw [one_mul, mul_one]
    · rw [wireOp_mul, wireOp_mul, ← mul_assoc, colT_Tk_rowAdd hj, ra_invol hj]

/-- Distinct row additions from a pivot with `z_h = 1`. -/
lemma raList_eq (h : Fin m) :
    ∀ (js : List (Fin m)), js.Nodup → ∀ z : Fin m → ZMod 2, z h = 1 →
      raList h js z = fun r => if r ∈ js ∧ r ≠ h then z r + 1 else z r := by
  intro js
  induction js with
  | nil => intro _ z _; funext r; simp [raList]
  | cons j js ih =>
    intro hnd z hz
    rw [List.nodup_cons] at hnd
    have hz' : raStep h j z h = 1 := by
      unfold raStep ra; split_ifs <;> simp_all
    show raList h js (raStep h j z) = _
    rw [ih hnd.2 _ hz']
    funext r
    by_cases hrj : r = j
    · subst hrj
      simp only [hnd.1, false_and, if_false, List.mem_cons, true_or, true_and]
      unfold raStep ra
      by_cases hh : h = r
      · subst hh; simp
      · simp [hh, Ne.symm hh, hz]
    · have e : raStep h j z r = z r := by
        unfold raStep ra; split_ifs <;> simp [hrj]
      simp only [e, List.mem_cons, hrj, false_or]

lemma zmod2_ne_zero (a : ZMod 2) (h : a ≠ 0) : a = 1 := by revert a; decide

/-- The pivot row additions send a nonzero `v` to `e_h`. -/
theorem raList_pivot (v : Fin m → ZMod 2) (h : Fin m) (hh : v h ≠ 0) :
    raList h (univ.filter fun j => v j ≠ 0).toList v = Pi.single h 1 := by
  have hh1 := zmod2_ne_zero _ hh
  rw [raList_eq h _ (Finset.nodup_toList _) v hh1]
  funext r
  simp only [Finset.mem_toList, mem_filter, mem_univ, true_and, Pi.single_apply]
  by_cases hr : r = h
  · subst hr; simp [hh1]
  · simp only [hr, if_false]
    by_cases hv : v r = 0
    · simp [hv]
    · simp [hv, zmod2_ne_zero _ hv, P8Lift.zmod2_add_self, hr]

/-- The chunk for one kernel `transK ε v`. -/
noncomputable def chunkFor (w : R) (ε : ℤ) (v : Fin m → ZMod 2) : List (RStep R m) := by
  classical
  exact if hv : ∃ h, v h ≠ 0 then
    rowAddsFor w (Classical.choose hv) (univ.filter fun j => v j ≠ 0).toList ++
      [RStep.child w (Classical.choose hv) (decide (ε = 1))] ++
      rowAddsFor w (Classical.choose hv) (univ.filter fun j => v j ≠ 0).toList.reverse
  else []

theorem chunkFor_den (w : R) (ε : ℤ) (hε : ε = 1 ∨ ε = -1) (v : Fin m → ZMod 2) (hv : v ≠ 0) :
    denListR (Butterfly.tens Butterfly.Cm : Op κ) (chunkFor w ε v) = wireOp w (colT (Tk ε v)) := by
  have hv' : ∃ h, v h ≠ 0 := by
    by_contra H; push Not at H; exact hv (funext H)
  unfold chunkFor
  rw [dif_pos hv']
  convert chunk_den w (Classical.choose hv') ε hε _ v
    (raList_pivot v _ (Classical.choose_spec hv'))

lemma countP_rowAddsFor (w : R) (h : Fin m) (js : List (Fin m)) :
    (rowAddsFor w h js).countP RStep.isChild = 0 := by
  rw [List.countP_eq_zero]
  intro st hst
  simp only [rowAddsFor, List.mem_flatMap] at hst
  obtain ⟨j, _, hj⟩ := hst
  split_ifs at hj
  · simp at hj
  · simp only [List.mem_singleton] at hj; subst hj; simp [RStep.isChild]

theorem chunkFor_count (w : R) (ε : ℤ) (v : Fin m → ZMod 2) (hv : v ≠ 0) :
    (chunkFor w ε v).countP RStep.isChild = 1 := by
  have hv' : ∃ h, v h ≠ 0 := by
    by_contra H; push Not at H; exact hv (funext H)
  unfold chunkFor
  rw [dif_pos hv']
  simp only [List.countP_append, countP_rowAddsFor]
  simp [RStep.isChild]

end Chunk

end P8Lift

end

section
-- Source module: Solutions.NetScalar
/-!
# Lemma 5 (bit construction) as a scalar gate network over `𝔽₂`

Parameters: a ground set `Fin h`, a finite type `T` of "triples" with `ind : T → Finset (Fin h)`,
`#(ind t) = 3`, `ind` injective.

1. Generic scalar gates `SGate W = Matrix W W (ZMod 2) × Finset W`, their run `sRun`, the
   matrix `runMat`, elementary gates `1 + egM A r c` (add control `c x` into target `r x`), and the
   bridge `labelled` / `scalarRun_labelled` / `ok_labelled` to `AModel.NGate`.
2. One invocation `Inv` (sources, targets, side wires `A_{ST}` for neighbor pairs, central
   wires `C_i`) with the eight-row forward schedule; `Inv.runMat_gates` / `forward_invocation`
   (Lemma 5): net change `y_S += x_S`, everything else unchanged, for arbitrary initial values.
3. The wires `Wr ind`, the gate list `bitGates` / named `bitNamed` / labelled `bitGatesL`
   (stage 1 `X → Y`, stage 2 `Y → X`, stage 3 `X → Y`, each over all `|𝒯|²` banks), and
   `bit_exchange` / `bit_exchange_net`: the circuit routes `w` to `ρ w` (`X a ↔ Y a`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false
set_option synthInstance.maxSize 512

namespace NetScalar

open Matrix Finset AModel

/-! ## 1. Generic scalar gates over `ZMod 2` -/

section Generic

variable {W : Type*} [Fintype W] [DecidableEq W]

/-- A scalar gate: matrix and touched set. -/
abbrev SGate (W : Type*) := Matrix W W (ZMod 2) × Finset W

/-- `NGate.OK` without the label. -/
def SOK (g : SGate W) : Prop :=
  ∀ w w', (w ∉ g.2 → g.1 w w' = if w = w' then 1 else 0) ∧
    (w ∈ g.2 → w' ∉ g.2 → g.1 w w' = 0)

/-- Scalar run of unlabelled gates. -/
def sRun (gs : List (SGate W)) (v : W → ZMod 2) : W → ZMod 2 :=
  gs.foldl (fun v g => g.1 *ᵥ v) v

/-- Matrix of a gate list (later gates on the left). -/
def runMat (gs : List (SGate W)) : Matrix W W (ZMod 2) :=
  gs.foldl (fun M g => g.1 * M) 1

/-- Attach labels to unlabelled gates. -/
def labelled {N M : Type*} (lab : N → M) (gs : List (N × SGate W)) : List (NGate W (ZMod 2) M) :=
  gs.map fun g => (g.2.1, g.2.2, lab g.1)

/-- Elementary gate matrix: add control `c x` into target `r x`, for `x ∈ A`. -/
def egM {ι : Type*} (A : Finset ι) (r c : ι → W) : Matrix W W (ZMod 2) :=
  ∑ x ∈ A, single (r x) (c x) 1

end Generic

/-- The forward-invocation identity in any ring: with the paper's signs, the eight rows
`−D, −C, +B, +G, +C, +D, −G, −B` multiply to `1 + C G + D B`, given that every other
product of two row matrices vanishes. -/
lemma eight_rows {R : Type*} [Ring R] (B G C D : R)
    (hBB : B * B = 0) (hBG : B * G = 0) (hBC : B * C = 0) (hBD : B * D = 0)
    (hGB : G * B = 0) (hGG : G * G = 0) (hGC : G * C = 0) (hGD : G * D = 0)
    (hCB : C * B = 0) (hCC : C * C = 0) (hCD : C * D = 0)
    (hDG : D * G = 0) (hDC : D * C = 0) (hDD : D * D = 0) :
    (1 - B) * (1 - G) * (1 + D) * (1 + C) * (1 + G) * (1 + B) * (1 - C) * (1 - D) =
      1 + (C * G + D * B) := by
  have t : ∀ X Y Z : R, X * Y = 0 → X * (Y * Z) = 0 := fun X Y Z h => by
    rw [← mul_assoc, h, zero_mul]
  simp only [mul_add, add_mul, mul_sub, sub_mul, one_mul, mul_one, mul_assoc, hBB, hBG, hBC, hBD,
    hGB, hGG, hGC, hGD, hCB, hCC, hCD, hDG, hDC, hDD, mul_zero, zero_mul, t _ _ _ hBC,
    t _ _ _ hGC, t _ _ _ hCC, t _ _ _ hDC, t _ _ _ hBD, t _ _ _ hGD, t _ _ _ hCD, t _ _ _ hDD,
    add_zero, zero_add, sub_zero, zero_sub]
  abel

/-! ## 2. One invocation (Lemma 5) -/

section Invocation

variable {h : ℕ} {T : Type*} [Fintype T] [DecidableEq T] (ind : T → Finset (Fin h))
variable {W : Type*} [Fintype W] [DecidableEq W]

/-- Ordered neighbor pairs `(S, T)`, `|S ∩ T| = 1`. -/
abbrev Nb := {st : T × T // #(ind st.1 ∩ ind st.2) = 1}

/-- Incidences `(i, T)` with `i ∈ T`. -/
abbrev Pinc : Finset (Fin h × T) := univ.filter fun q => q.1 ∈ ind q.2

/-- The wires of one invocation: sources `x_T`, targets `y_S`, side wires `A_{ST}`,
central wires `C_i`. -/
structure Inv (W : Type*) where
  src : T → W
  tgt : T → W
  side : NetScalar.Nb ind → W
  cen : Fin h → W

variable {ind}

/-- Wires of an invocation are distinct. -/
structure Inv.Good (I : NetScalar.Inv ind W) : Prop where
  src_inj : Function.Injective I.src
  tgt_inj : Function.Injective I.tgt
  side_inj : Function.Injective I.side
  cen_inj : Function.Injective I.cen
  st : ∀ a b, I.src a ≠ I.tgt b
  ss : ∀ a n, I.src a ≠ I.side n
  sc : ∀ a i, I.src a ≠ I.cen i
  ts : ∀ a n, I.tgt a ≠ I.side n
  tc : ∀ a i, I.tgt a ≠ I.cen i
  sic : ∀ n i, I.side n ≠ I.cen i

/-- Side injection into target `y_S` (one gate per target). -/
def Inv.sideG (I : NetScalar.Inv ind W) (S : T) : SGate W :=
  (1 + NetScalar.egM (univ.filter fun n : NetScalar.Nb ind => n.1.1 = S) (fun n => I.tgt n.1.1) I.side,
   insert (I.tgt S) ((univ.filter fun n : NetScalar.Nb ind => n.1.1 = S).image I.side))

/-- Copy of source `x_T` into the side wires `A_{ST}` (one gate per source). -/
def Inv.copyG (I : NetScalar.Inv ind W) (t : T) : SGate W :=
  (1 + NetScalar.egM (univ.filter fun n : NetScalar.Nb ind => n.1.2 = t) I.side (fun n => I.src n.1.2),
   insert (I.src t) ((univ.filter fun n : NetScalar.Nb ind => n.1.2 = t).image I.side))

/-- Scatter `y_S += ∑_{i ∈ S} C_i` (one gate on targets and center). -/
def Inv.scatG (I : NetScalar.Inv ind W) : SGate W :=
  (1 + NetScalar.egM (NetScalar.Pinc ind) (fun q => I.tgt q.2) (fun q => I.cen q.1),
   univ.image I.tgt ∪ univ.image I.cen)

/-- Gather `C_i += ∑_{T ∋ i} x_T` (one gate on sources and center). -/
def Inv.gathG (I : NetScalar.Inv ind W) : SGate W :=
  (1 + NetScalar.egM (NetScalar.Pinc ind) (fun q => I.cen q.1) (fun q => I.src q.2),
   univ.image I.src ∪ univ.image I.cen)

/-- Gate names inside an invocation: schedule row and, for per-wire gates, the wire's triple. -/
abbrev IName (T : Type*) := Fin 8 × Option T

/-- The forward schedule of one invocation, rows 0–7, with names. Over `𝔽₂` subtracting equals
adding, so rows 0/5, 1/4, 3/6 and 2/7 use the same gates. -/
noncomputable def Inv.named (I : NetScalar.Inv ind W) : List (IName T × SGate W) :=
  ((univ : Finset T).toList.map fun S => ((0, some S), I.sideG S)) ++ [((1, none), I.scatG)] ++
  ((univ : Finset T).toList.map fun t => ((2, some t), I.copyG t)) ++ [((3, none), I.gathG)] ++
  [((4, none), I.scatG)] ++
  ((univ : Finset T).toList.map fun S => ((5, some S), I.sideG S)) ++ [((6, none), I.gathG)] ++
  ((univ : Finset T).toList.map fun t => ((7, some t), I.copyG t))

noncomputable def Inv.gates (I : NetScalar.Inv ind W) : List (SGate W) := I.named.map Prod.snd

/-- Row matrices. -/
def Inv.Dm (I : NetScalar.Inv ind W) : Matrix W W (ZMod 2) := NetScalar.egM univ (fun n : NetScalar.Nb ind => I.tgt n.1.1) I.side
def Inv.Bm (I : NetScalar.Inv ind W) : Matrix W W (ZMod 2) := NetScalar.egM univ I.side (fun n : NetScalar.Nb ind => I.src n.1.2)
def Inv.Cm (I : NetScalar.Inv ind W) : Matrix W W (ZMod 2) :=
  NetScalar.egM (NetScalar.Pinc ind) (fun q => I.tgt q.2) (fun q => I.cen q.1)
def Inv.Gm (I : NetScalar.Inv ind W) : Matrix W W (ZMod 2) :=
  NetScalar.egM (NetScalar.Pinc ind) (fun q => I.cen q.1) (fun q => I.src q.2)
/-- Net change `J`: `y_S += x_S`. -/
def Inv.Jm (I : NetScalar.Inv ind W) : Matrix W W (ZMod 2) := NetScalar.egM univ I.tgt I.src

variable {I : Inv ind W}

end Invocation

/-! ## 3. The bit construction: three stages of invocations on the data banks -/

section Bit

variable {h : ℕ} {T : Type*} [Fintype T] [DecidableEq T] (ind : T → Finset (Fin h))

/-- Wires: `X a`, `Y a` (`a ∈ 𝒯³`), side wires `A_{ST}` and central wires `C_i` of each
invocation `(stage j, bank p)` (the bank `p` lists the two fixed coordinates). -/
abbrev Wr := (T × T × T) ⊕ (T × T × T) ⊕ (Fin 3 × (T × T) × NetScalar.Nb ind) ⊕ (Fin 3 × (T × T) × Fin h)

variable {ind}

def wX (a : T × T × T) : NetScalar.Wr ind := Sum.inl a
def wY (a : T × T × T) : NetScalar.Wr ind := Sum.inr (Sum.inl a)
def wS (j : Fin 3) (p : T × T) (n : NetScalar.Nb ind) : NetScalar.Wr ind := Sum.inr (Sum.inr (Sum.inl (j, p, n)))
def wC (j : Fin 3) (p : T × T) (i : Fin h) : NetScalar.Wr ind := Sum.inr (Sum.inr (Sum.inr (j, p, i)))

/-- The bank element with coordinate `j` equal to `t` and the other two equal to `p`. -/
def emb (j : Fin 3) (p : T × T) (t : T) : T × T × T :=
  if j = 0 then (t, p.1, p.2) else if j = 1 then (p.1, t, p.2) else (p.1, p.2, t)

lemma emb_inj (j : Fin 3) (p : T × T) : Function.Injective (emb j p) := by
  intro t t' e
  unfold emb at e
  split_ifs at e <;> simp_all [Prod.ext_iff]

/-- `(p, t) ↦ emb j p t` is a bijection `(𝒯 × 𝒯) × 𝒯 ≃ 𝒯³`. -/
def embEquiv (j : Fin 3) : (T × T) × T ≃ T × T × T where
  toFun x := emb j x.1 x.2
  invFun a := if j = 0 then ((a.2.1, a.2.2), a.1) else if j = 1 then ((a.1, a.2.2), a.2.1)
    else ((a.1, a.2.1), a.2.2)
  left_inv x := by unfold emb; split_ifs <;> rfl
  right_inv a := by unfold emb; split_ifs <;> rfl

lemma sum_emb {M : Type*} [AddCommMonoid M] (j : Fin 3) (F : T × T × T → M) :
    ∑ p, ∑ t, F (emb j p t) = ∑ a, F a := by
  rw [← Fintype.sum_prod_type']
  exact Fintype.sum_equiv (embEquiv j) _ _ fun _ => rfl

/-- Sources and targets of stage `j`: `X → Y` at stages 1 and 3, `Y → X` at stage 2. -/
def srcW (j : Fin 3) (a : T × T × T) : NetScalar.Wr ind := if j = 1 then NetScalar.wY a else NetScalar.wX a
def tgtW (j : Fin 3) (a : T × T × T) : NetScalar.Wr ind := if j = 1 then NetScalar.wX a else NetScalar.wY a

/-- The invocation of stage `j` on bank `p`. -/
def inv (j : Fin 3) (p : T × T) : NetScalar.Inv ind (NetScalar.Wr ind) :=
  ⟨fun t => NetScalar.srcW j (emb j p t), fun t => NetScalar.tgtW j (emb j p t), NetScalar.wS j p, NetScalar.wC j p⟩

/-- Gate names: stage, invocation (bank), and the name inside the invocation. -/
abbrev GName (T : Type*) := Fin 3 × (T × T) × IName T

variable (ind)

/-- The named gate list: stage 1 (`X → Y`), stage 2 (`Y → X`), stage 3 (`X → Y`); in each stage
the invocations (banks) one after another, each running the forward schedule (rows 0–7). -/
noncomputable def bitNamed : List (NetScalar.GName T × SGate (NetScalar.Wr ind)) :=
  ([0, 1, 2] : List (Fin 3)).flatMap fun j => (univ : Finset (T × T)).toList.flatMap fun p =>
    (NetScalar.inv j p : NetScalar.Inv ind (NetScalar.Wr ind)).named.map fun g => ((j, p, g.1), g.2)

/-- The scalar gates of the bit construction. -/
noncomputable def bitGates : List (SGate (NetScalar.Wr ind)) := (bitNamed ind).map Prod.snd

/-- The gates with labels attached by a naming function. -/
noncomputable def bitGatesL {M : Type*} (lab : NetScalar.GName T → M) : List (NGate (NetScalar.Wr ind) (ZMod 2) M) :=
  NetScalar.labelled lab (bitNamed ind)

/-- The routing: `X a ↔ Y a`, auxiliaries fixed. -/
def ρ : NetScalar.Wr ind → NetScalar.Wr ind
  | Sum.inl a => Sum.inr (Sum.inl a)
  | Sum.inr (Sum.inl a) => Sum.inl a
  | Sum.inr (Sum.inr x) => Sum.inr (Sum.inr x)

variable {ind}

/-- Stage matrix `K_j = ∑_a e_{tgt a, src a}`. -/
def Kst (j : Fin 3) : Matrix (NetScalar.Wr ind) (NetScalar.Wr ind) (ZMod 2) := NetScalar.egM univ (NetScalar.tgtW j) (NetScalar.srcW j)

end Bit

end NetScalar

end