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

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

import Definitions.Def_IntMul_Kappa_K_Defs1

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.AModelFac
/-!
# Discharging the factorization hypothesis of Proposition 11 (Lemma 10 data)

`AModelProp11.prop11` assumes, as data, integer factorizations `fac A = (E₁, Π, E₂)` with
`A = E₁ Π E₂` **over `ℤ`** (`FacAt`). This cannot be produced in general from Lemma 9 (the
factors are rational, and their reductions mod `q ^ b` are `q`-adic, not integral). This file:

0. `bruhat_lower`: a copy of `Solutions/Bruhat.lean` (Lemma 9 over a field), renamed.
1. `ratFac`, `RatFacSpec`, `GoodAt`, `exists_good_primes`, `exists_good_prime`: for any finite set
   of integer matrices, every prime `q > q₀` avoids all denominators of the `E`'s and all
   numerators of their diagonals; `#ones Π = rank_ℚ A`.
2. `goodSub N` (rationals with denominator coprime to `N`), the ring hom `redHom N` to `ZMod N`,
   the total map `redQ N`, `redQ_matrix_identity`, `redQ_lower`, `isUnit_redQ`.
3. `FacAtZ` (the identity only mod `N`), `facAtZ_of_facAt`, the integer lifts `facZ Q`
   (entries `(redQ Q r).val`), the bridge `facAtZ_facZ`, and the re-proved chain
   `implΦ_eqZ`, `levelR_correct_amapZ`, `levelR_correctZ`, `prog_correctZ` (bounded level),
   `levelCost_eqZ`, `prog_costZ`, `prop11Z`, and the assembled `prop11_rat`.
-/

section BruhatCopy

open Matrix Finset

namespace FacBruhat

variable {K : Type*} [Field K] {m : ℕ}

/-- Lower triangular with nonzero diagonal. -/
def LowerND (M : Matrix (Fin m) (Fin m) K) : Prop :=
  (∀ i j, i < j → M i j = 0) ∧ ∀ i, M i i ≠ 0

lemma lowerND_one : LowerND (1 : Matrix (Fin m) (Fin m) K) :=
  ⟨fun i j h => by simp [one_apply, h.ne], fun i => by simp⟩

lemma lowerND_mul {M N : Matrix (Fin m) (Fin m) K} (hM : LowerND M) (hN : LowerND N) :
    LowerND (M * N) := by
  have diag : ∀ i, (M * N) i i = M i i * N i i := by
    intro i
    rw [mul_apply, Finset.sum_eq_single i]
    · intro k _ hk
      rcases lt_or_gt_of_ne hk with h | h
      · rw [hN.1 k i h, mul_zero]
      · rw [hM.1 i k h, zero_mul]
    · simp
  refine ⟨fun i j hij => ?_, fun i => by rw [diag]; exact mul_ne_zero (hM.2 i) (hN.2 i)⟩
  rw [mul_apply]
  refine Finset.sum_eq_zero fun k _ => ?_
  rcases lt_or_ge i k with h | h
  · rw [hM.1 i k h, zero_mul]
  · rw [hN.1 k j (lt_of_le_of_lt h hij), mul_zero]

lemma isLower_of (M : Matrix (Fin m) (Fin m) K) (h : LowerND M) : M.IsLowerTriangular :=
  fun i j hij => h.1 i j hij

lemma det_ne_zero_of (M : Matrix (Fin m) (Fin m) K) (h : LowerND M) : M.det ≠ 0 := by
  rw [det_of_isLowerTriangular M (isLower_of M h)]
  exact Finset.prod_ne_zero_iff.2 fun i _ => h.2 i

lemma lowerND_inv (M : Matrix (Fin m) (Fin m) K) (h : LowerND M) : LowerND M⁻¹ := by
  have hd : IsUnit M.det := isUnit_iff_ne_zero.2 (det_ne_zero_of M h)
  have : Invertible M := invertibleOfIsUnitDet M hd
  have hl : M⁻¹.IsLowerTriangular := blockTriangular_inv_of_blockTriangular (isLower_of M h)
  refine ⟨fun i j hij => hl hij, fun i => ?_⟩
  -- diagonal of M * M⁻¹ = 1
  have h1 : (M * M⁻¹) i i = 1 := by rw [mul_inv_of_invertible]; simp
  have h2 : (M * M⁻¹) i i = M i i * M⁻¹ i i := by
    rw [mul_apply, Finset.sum_eq_single i]
    · intro k _ hk
      rcases lt_or_gt_of_ne hk with hk | hk
      · rw [hl hk, mul_zero]
      · rw [h.1 i k hk, zero_mul]
    · simp
  intro h0
  rw [h2, h0, mul_zero] at h1
  exact zero_ne_one h1

/-- Rows with index `< k` are clean: zero, or a single `1` whose column is otherwise zero. -/
def Clean (B : Matrix (Fin m) (Fin m) K) (k : ℕ) : Prop :=
  ∀ i : Fin m, i.val < k → (∀ j, B i j = 0) ∨
    ∃ j, B i j = 1 ∧ (∀ j', j' ≠ j → B i j' = 0) ∧ ∀ i', i' ≠ i → B i' j = 0

lemma clean_zero (B : Matrix (Fin m) (Fin m) K) : Clean B 0 := fun _ h => absurd h (Nat.not_lt_zero _)

/-- In a clean row `r < k`, the entry in any column `j` where row `k'` (with `k' ≥ k`) is nonzero vanishes. -/
lemma clean_col_zero {B : Matrix (Fin m) (Fin m) K} {k : ℕ} (hB : Clean B k) (r k' : Fin m)
    (hr : r.val < k) (hk' : k ≤ k'.val) (j : Fin m) (hj : B k' j ≠ 0) : B r j = 0 := by
  rcases hB r hr with h | ⟨σ, -, hrow, hcol⟩
  · exact h j
  · by_cases hσ : j = σ
    · subst hσ
      exact absurd (hcol k' (fun h => by subst h; omega)) hj
    · exact hrow j hσ

lemma step (B : Matrix (Fin m) (Fin m) K) (k : Fin m) (hB : Clean B k.val) :
    ∃ L R : Matrix (Fin m) (Fin m) K, LowerND L ∧ LowerND R ∧ Clean (L * B * R) (k.val + 1) := by
  classical
  by_cases hz : ∀ j, B k j = 0
  · refine ⟨1, 1, lowerND_one, lowerND_one, ?_⟩
    simp only [Matrix.one_mul, Matrix.mul_one]
    intro i hi
    rcases Nat.lt_succ_iff_lt_or_eq.1 hi with h | h
    · exact hB i h
    · left; have : i = k := Fin.ext h; subst this; exact hz
  push Not at hz
  have hne : (univ.filter fun j => B k j ≠ 0).Nonempty := by
    obtain ⟨j, hj⟩ := hz; exact ⟨j, by simp [hj]⟩
  set j0 := (univ.filter fun j => B k j ≠ 0).max' hne with hj0def
  have hj0 : B k j0 ≠ 0 := by
    have := Finset.max'_mem _ hne
    simpa using this
  have hright : ∀ j, j0 < j → B k j = 0 := by
    intro j hj
    by_contra h
    have : j ≤ j0 := Finset.le_max' _ j (by simp [h])
    exact absurd hj (not_lt.2 this)
  set b := B k j0 with hb
  -- column operations
  set N : Matrix (Fin m) (Fin m) K := Matrix.of fun a c => if a = j0 ∧ c < j0 then -(B k c / b) else 0
  set R1 : Matrix (Fin m) (Fin m) K := 1 + N
  have hR1 : LowerND R1 := by
    refine ⟨fun i j hij => ?_, fun i => ?_⟩
    · simp only [R1, N, Matrix.add_apply, one_apply, of_apply, hij.ne, if_false, zero_add]
      rw [if_neg]; rintro ⟨rfl, h⟩; exact absurd (hij.trans h) (lt_irrefl _)
    · simp only [R1, N, Matrix.add_apply, of_apply, one_apply_eq]
      rw [if_neg (by rintro ⟨rfl, h⟩; exact lt_irrefl _ h)]; simp
  set B1 := B * R1 with hB1def
  have hB1 : ∀ r c, B1 r c = B r c + (if c < j0 then -(B k c / b) * B r j0 else 0) := by
    intro r c
    simp only [B1, R1, Matrix.mul_add, Matrix.mul_one, Matrix.add_apply]
    congr 1
    rw [mul_apply, Finset.sum_eq_single j0]
    · simp only [N, of_apply, true_and]
      split_ifs <;> ring
    · intro a _ ha; simp [N, ha]
    · simp
  have hB1k : ∀ c, B1 k c = if c = j0 then b else 0 := by
    intro c
    rw [hB1]
    rcases lt_trichotomy c j0 with h | h | h
    · rw [if_pos h, if_neg h.ne]; field_simp; ring
    · subst h; simp [hb]
    · rw [if_neg (not_lt.2 h.le), if_neg h.ne', hright c h, add_zero]
  -- row operations
  set M : Matrix (Fin m) (Fin m) K := Matrix.of fun r a => if a = k ∧ k < r then -(B1 r j0 / b) else 0
  set L1 : Matrix (Fin m) (Fin m) K := 1 + M
  have hL1 : LowerND L1 := by
    refine ⟨fun i j hij => ?_, fun i => ?_⟩
    · simp only [L1, M, Matrix.add_apply, one_apply, of_apply, hij.ne, if_false, zero_add]
      rw [if_neg]; rintro ⟨rfl, h⟩; exact absurd (hij.trans h) (lt_irrefl _)
    · simp only [L1, M, Matrix.add_apply, of_apply, one_apply_eq]
      rw [if_neg (by rintro ⟨rfl, h⟩; exact lt_irrefl _ h)]; simp
  set B2 := L1 * B1 with hB2def
  have hB2 : ∀ r c, B2 r c = B1 r c + (if k < r then -(B1 r j0 / b) * B1 k c else 0) := by
    intro r c
    simp only [B2, L1, Matrix.add_mul, Matrix.one_mul, Matrix.add_apply]
    congr 1
    rw [mul_apply, Finset.sum_eq_single k]
    · simp only [M, of_apply, true_and]
      split_ifs <;> ring
    · intro a _ ha; simp [M, ha]
    · simp
  -- scaling
  set D : Matrix (Fin m) (Fin m) K := diagonal fun i => if i = k then b⁻¹ else 1
  have hD : LowerND D := by
    refine ⟨fun i j hij => by simp [D, hij.ne], fun i => ?_⟩
    simp only [D, diagonal_apply_eq]
    split_ifs
    · exact inv_ne_zero hj0
    · exact one_ne_zero
  refine ⟨D * L1, R1, lowerND_mul hD hL1, hR1, ?_⟩
  have hfin : D * L1 * B * R1 = D * B2 := by
    simp only [B2, B1, Matrix.mul_assoc]
  rw [hfin]
  have hDB : ∀ r c, (D * B2) r c = (if r = k then b⁻¹ else 1) * B2 r c := by
    intro r c; simp [D, diagonal_mul]
  -- clean rows below k are untouched
  have hold : ∀ r : Fin m, r.val < k.val → ∀ c, (D * B2) r c = B r c := by
    intro r hr c
    have hrk : r ≠ k := fun h => by subst h; omega
    have hkr : ¬ k < r := fun h => by rw [Fin.lt_def] at h; omega
    have hrj0 : B r j0 = 0 := clean_col_zero hB r k hr le_rfl j0 hj0
    rw [hDB, if_neg hrk, one_mul, hB2, if_neg hkr, add_zero, hB1, hrj0]
    simp
  intro i hi
  rcases Nat.lt_succ_iff_lt_or_eq.1 hi with hi | hi
  · rcases hB i hi with h | ⟨σ, h1, hrow, hcol⟩
    · left; intro j; rw [hold i hi, h]
    · right
      refine ⟨σ, by rw [hold i hi, h1], fun j hj => by rw [hold i hi, hrow j hj], fun i' hi' => ?_⟩
      have hσj0 : σ ≠ j0 := by
        rintro rfl; exact hj0 (hcol k (fun h => by subst h; omega))
      have hkσ : B k σ = 0 := hcol k (fun h => by subst h; omega)
      rw [hDB, hB2, hB1k, if_neg hσj0, mul_zero]
      simp only [ite_self, add_zero]
      rw [hB1, hcol i' hi', hkσ]
      simp
  · have hik : i = k := Fin.ext hi
    subst hik
    right
    refine ⟨j0, ?_, fun j hj => ?_, fun i' hi' => ?_⟩
    · rw [hDB, if_pos rfl, hB2, if_neg (lt_irrefl _), add_zero, hB1k, if_pos rfl]
      exact inv_mul_cancel₀ hj0
    · rw [hDB, hB2, if_neg (lt_irrefl _), add_zero, hB1k, if_neg hj, mul_zero]
    · rw [hDB, if_neg hi', hB2, hB1k, if_pos rfl, one_mul]
      rcases lt_trichotomy i' i with h | h | h
      · rw [if_neg (not_lt.2 h.le), add_zero]
        have := hold i' h j0
        rw [hDB, if_neg hi', one_mul, hB2, if_neg (not_lt.2 h.le), add_zero] at this
        rw [this]
        exact clean_col_zero hB i' i h le_rfl j0 hj0
      · exact absurd h hi'
      · rw [if_pos h]; field_simp; ring

lemma reduce (A : Matrix (Fin m) (Fin m) K) :
    ∀ k ≤ m, ∃ L R : Matrix (Fin m) (Fin m) K, LowerND L ∧ LowerND R ∧ Clean (L * A * R) k := by
  intro k
  induction k with
  | zero => intro _; exact ⟨1, 1, lowerND_one, lowerND_one, clean_zero _⟩
  | succ k ih =>
    intro hk
    obtain ⟨L, R, hL, hR, hc⟩ := ih (by omega)
    obtain ⟨L', R', hL', hR', hc'⟩ := step (L * A * R) ⟨k, by omega⟩ hc
    refine ⟨L' * L, R * R', lowerND_mul hL' hL, lowerND_mul hR hR', ?_⟩
    simpa [Matrix.mul_assoc] using hc'

end FacBruhat

open FacBruhat in
theorem bruhat_lower {m : ℕ} {K : Type*} [Field K] [DecidableEq K] (A : Matrix (Fin m) (Fin m) K) :
    ∃ E₁ P E₂ : Matrix (Fin m) (Fin m) K,
      (∀ i j, i < j → E₁ i j = 0) ∧ (∀ i, E₁ i i ≠ 0) ∧
      (∀ i j, i < j → E₂ i j = 0) ∧ (∀ i, E₂ i i ≠ 0) ∧
      (∀ i j, P i j = 0 ∨ P i j = 1) ∧
      (∀ i j j', P i j = 1 → P i j' = 1 → j = j') ∧
      (∀ i i' j, P i j = 1 → P i' j = 1 → i = i') ∧
      (Finset.univ.filter fun x : Fin m × Fin m => P x.1 x.2 = 1).card = A.rank ∧
      A = E₁ * P * E₂ := by
  classical
  obtain ⟨L, R, hL, hR, hc⟩ := reduce A m le_rfl
  set P := L * A * R with hP
  have hLd := det_ne_zero_of L hL
  have hRd := det_ne_zero_of R hR
  have : Invertible L := invertibleOfIsUnitDet L (isUnit_iff_ne_zero.2 hLd)
  have : Invertible R := invertibleOfIsUnitDet R (isUnit_iff_ne_zero.2 hRd)
  have hA : A = L⁻¹ * P * R⁻¹ := by
    simp only [hP, Matrix.mul_assoc, inv_mul_cancel_left_of_invertible]
    rw [Matrix.mul_inv_of_invertible, Matrix.mul_one]
  have hrow : ∀ i, (∀ j, P i j = 0) ∨
      ∃ j, P i j = 1 ∧ (∀ j', j' ≠ j → P i j' = 0) ∧ ∀ i', i' ≠ i → P i' j = 0 :=
    fun i => hc i i.isLt
  have h01 : ∀ i j, P i j = 0 ∨ P i j = 1 := by
    intro i j
    rcases hrow i with h | ⟨σ, h1, hr, -⟩
    · exact Or.inl (h j)
    · by_cases hj : j = σ
      · subst hj; exact Or.inr h1
      · exact Or.inl (hr j hj)
  have hrowu : ∀ i j j', P i j = 1 → P i j' = 1 → j = j' := by
    intro i j j' h h'
    rcases hrow i with hz | ⟨σ, -, hr, -⟩
    · rw [hz j] at h; exact absurd h zero_ne_one
    · by_contra hne
      by_cases hj : j = σ
      · subst hj; rw [hr j' (Ne.symm hne)] at h'; exact zero_ne_one h'
      · rw [hr j hj] at h; exact zero_ne_one h
  have hcolu : ∀ i i' j, P i j = 1 → P i' j = 1 → i = i' := by
    intro i i' j h h'
    rcases hrow i with hz | ⟨σ, -, hr, hcl⟩
    · rw [hz j] at h; exact absurd h zero_ne_one
    · have hjσ : j = σ := by
        by_contra hj; rw [hr j hj] at h; exact zero_ne_one h
      subst hjσ
      by_contra hne
      rw [hcl i' (Ne.symm hne)] at h'; exact zero_ne_one h'
  -- rank count
  have hrankP : P.rank = A.rank := by
    rw [hP, rank_mul_eq_left_of_isUnit_det R _ (isUnit_iff_ne_zero.2 hRd),
      rank_mul_eq_right_of_isUnit_det L A (isUnit_iff_ne_zero.2 hLd)]
  have hPPt : P * Pᵀ = diagonal fun i => if ∃ j, P i j = 1 then (1 : K) else 0 := by
    ext i i'
    rw [mul_apply, diagonal_apply]
    simp only [transpose_apply]
    by_cases hii : i = i'
    · subst hii
      rw [if_pos rfl]
      split_ifs with hex
      · obtain ⟨j, hj⟩ := hex
        rw [Finset.sum_eq_single j]
        · rw [hj, one_mul]
        · intro j' _ hj'
          rcases h01 i j' with h | h
          · rw [h, zero_mul]
          · exact absurd (hrowu i j j' hj h).symm hj'
        · simp
      · push Not at hex
        refine Finset.sum_eq_zero fun j _ => ?_
        rcases h01 i j with h | h
        · rw [h, zero_mul]
        · exact absurd h (hex j)
    · rw [if_neg hii]
      refine Finset.sum_eq_zero fun j _ => ?_
      rcases h01 i j with h | h
      · rw [h, zero_mul]
      rcases h01 i' j with h' | h'
      · rw [h', mul_zero]
      · exact absurd (hcolu i i' j h h') hii
  have hPPtP : P * Pᵀ * P = P := by
    rw [hPPt]
    ext i j
    rw [diagonal_mul]
    split_ifs with hex
    · rw [one_mul]
    · push Not at hex
      rw [zero_mul]
      rcases h01 i j with h | h
      · exact h.symm
      · exact absurd h (hex j)
  have hrank1 : P.rank = (P * Pᵀ).rank := by
    apply le_antisymm
    · calc P.rank = (P * Pᵀ * P).rank := by rw [hPPtP]
        _ ≤ (P * Pᵀ).rank := rank_mul_le_left _ _
    · exact rank_mul_le_left _ _
  have hrank2 : (P * Pᵀ).rank = (univ.filter fun i => ∃ j, P i j = 1).card := by
    rw [hPPt, rank_diagonal, Fintype.card_subtype]
    congr 1
    ext i; simp
  have hcount : (univ.filter fun x : Fin m × Fin m => P x.1 x.2 = 1).card =
      (univ.filter fun i => ∃ j, P i j = 1).card := by
    refine Finset.card_bij (fun x _ => x.1) ?_ ?_ ?_
    · intro x hx; simp only [mem_filter, mem_univ, true_and] at hx ⊢; exact ⟨x.2, hx⟩
    · intro x hx y hy hxy
      simp only [mem_filter, mem_univ, true_and] at hx hy
      have h2 : x.2 = y.2 := hrowu x.1 x.2 y.2 hx (hxy ▸ hy)
      exact Prod.ext hxy h2
    · intro i hi
      simp only [mem_filter, mem_univ, true_and] at hi
      obtain ⟨j, hj⟩ := hi
      exact ⟨(i, j), by simp [hj], rfl⟩
  have hLi := lowerND_inv L hL
  have hRi := lowerND_inv R hR
  exact ⟨L⁻¹, P, R⁻¹, hLi.1, hLi.2, hRi.1, hRi.2, h01, hrowu, hcolu,
    by rw [hcount, ← hrank2, ← hrank1, hrankP], hA⟩
end BruhatCopy

namespace AModelFac

open Matrix AModel

/-! ## 1. Rational Lemma 9 factorizations and good primes -/

section Rational

variable {m : ℕ}

/-- Rational `m × m` matrices. -/
abbrev QMat (m : ℕ) := Matrix (Fin m) (Fin m) ℚ

/-- An integer matrix viewed over `ℚ`. -/
def ratM (A : IMat m) : QMat m := A.map (Int.cast)

/-- The Lemma 9 properties of a triple `(E₁, Π, E₂)` for `A` over `ℚ`. -/
def RatFacSpec (A : QMat m) (F : QMat m × QMat m × QMat m) : Prop :=
  (∀ i j, i < j → F.1 i j = 0) ∧ (∀ i, F.1 i i ≠ 0) ∧
  (∀ i j, i < j → F.2.2 i j = 0) ∧ (∀ i, F.2.2 i i ≠ 0) ∧
  (∀ i j, F.2.1 i j = 0 ∨ F.2.1 i j = 1) ∧
  (∀ i j j', F.2.1 i j = 1 → F.2.1 i j' = 1 → j = j') ∧
  (∀ i i' j, F.2.1 i j = 1 → F.2.1 i' j = 1 → i = i') ∧
  (Finset.univ.filter fun x : Fin m × Fin m => F.2.1 x.1 x.2 = 1).card = A.rank ∧
  A = F.1 * F.2.1 * F.2.2

lemma exists_ratFac (A : IMat m) : ∃ F, RatFacSpec (ratM A) F := by
  obtain ⟨E₁, P, E₂, h⟩ := bruhat_lower (ratM A)
  exact ⟨(E₁, P, E₂), h⟩

/-- The rational factorization of `A`, chosen once (Lemma 9 over `ℚ`). -/
noncomputable def ratFac (A : IMat m) : QMat m × QMat m × QMat m := (exists_ratFac A).choose

/-- `q` avoids every denominator of `E₁, E₂` and every numerator of their diagonals. -/
def GoodAt (q : ℕ) (F : QMat m × QMat m × QMat m) : Prop :=
  (∀ i j, Nat.Coprime (F.1 i j).den q) ∧ (∀ i j, Nat.Coprime (F.2.2 i j).den q) ∧
  (∀ i, Nat.Coprime (F.1 i i).num.natAbs q) ∧ (∀ i, Nat.Coprime (F.2.2 i i).num.natAbs q)

/-- A bound beyond which every prime is good for `ratFac A`. -/
noncomputable def facBound (A : IMat m) : ℕ :=
  (∑ i, ∑ j, (((ratFac A).1 i j).den + ((ratFac A).2.2 i j).den)) +
  ∑ i, (((ratFac A).1 i i).num.natAbs + ((ratFac A).2.2 i i).num.natAbs)

end Rational

/-! ## 2. Reduction `ℚ_(N) → ZMod N` -/

section Reduction

/-- Rationals whose denominator is coprime to `N` (a subring). -/
def goodSub (N : ℕ) : Subring ℚ where
  carrier := {r | Nat.Coprime r.den N}
  zero_mem' := by simp
  one_mem' := by simp
  add_mem' := fun {a b} ha hb =>
    Nat.Coprime.coprime_dvd_left (Rat.add_den_dvd a b) (Nat.Coprime.mul_left ha hb)
  mul_mem' := fun {a b} ha hb =>
    Nat.Coprime.coprime_dvd_left (Rat.mul_den_dvd a b) (Nat.Coprime.mul_left ha hb)
  neg_mem' := fun {a} ha => by simpa [Rat.neg_den] using ha

/-- `r ↦ num · den⁻¹` in `ZMod N`. -/
def redFun (N : ℕ) (r : ℚ) : ZMod N := (r.num : ZMod N) * ((r.den : ℕ) : ZMod N)⁻¹

lemma cancel_coprime {N d : ℕ} (h : Nat.Coprime d N) {x y : ZMod N}
    (hxy : x * d = y * d) : x = y := by
  have := congrArg (· * (d : ZMod N)⁻¹) hxy
  simpa only [mul_assoc, ZMod.coe_mul_inv_eq_one d h, mul_one] using this

lemma redFun_mul_den {N : ℕ} {r : ℚ} (h : Nat.Coprime r.den N) :
    redFun N r * r.den = r.num := by
  rw [redFun, mul_assoc, mul_comm _ (r.den : ZMod N), ZMod.coe_mul_inv_eq_one _ h, mul_one]

/-- Characterization: if `r * d = x` with `d` coprime to `N`, then `red r * d = x`. -/
lemma redFun_spec {N : ℕ} {r : ℚ} (h : Nat.Coprime r.den N) (x : ℤ) (d : ℕ)
    (hr : r * d = x) : redFun N r * d = x := by
  have h1 : (r.num : ℚ) * d = x * r.den := by
    rw [← Rat.mul_den_eq_num r, ← hr]; ring
  have h2 : r.num * (d : ℤ) = x * (r.den : ℤ) := by exact_mod_cast h1
  apply cancel_coprime h
  rw [mul_right_comm, redFun_mul_den h]
  have := congrArg (Int.cast : ℤ → ZMod N) h2
  push_cast at this
  exact this

/-- **Task 2.** The reduction `ℚ_(N) → ZMod N` is a ring hom. -/
def redHom (N : ℕ) : goodSub N →+* ZMod N where
  toFun r := redFun N r
  map_one' := by
    have := redFun_spec (N := N) (r := 1) (by simp) 1 1 (by simp)
    simpa using this
  map_zero' := by
    have := redFun_spec (N := N) (r := 0) (by simp) 0 1 (by simp)
    simpa using this
  map_mul' a b := by
    have ha : Nat.Coprime a.1.den N := a.2
    have hb : Nat.Coprime b.1.den N := b.2
    change redFun N (a.1 * b.1) = redFun N a.1 * redFun N b.1
    have hc : Nat.Coprime (a.1.den * b.1.den) N := Nat.Coprime.mul_left ha hb
    apply cancel_coprime hc
    rw [redFun_spec ((goodSub N).mul_mem a.2 b.2) (a.1.num * b.1.num) _ (by
      push_cast; rw [← Rat.mul_den_eq_num a.1, ← Rat.mul_den_eq_num b.1]; ring)]
    push_cast
    rw [← redFun_mul_den ha, ← redFun_mul_den hb]; ring
  map_add' a b := by
    have ha : Nat.Coprime a.1.den N := a.2
    have hb : Nat.Coprime b.1.den N := b.2
    change redFun N (a.1 + b.1) = redFun N a.1 + redFun N b.1
    have hc : Nat.Coprime (a.1.den * b.1.den) N := Nat.Coprime.mul_left ha hb
    apply cancel_coprime hc
    rw [redFun_spec ((goodSub N).add_mem a.2 b.2) (a.1.num * b.1.den + b.1.num * a.1.den) _ (by
      push_cast; rw [← Rat.mul_den_eq_num a.1, ← Rat.mul_den_eq_num b.1]; ring)]
    push_cast
    rw [← redFun_mul_den ha, ← redFun_mul_den hb]; ring

/-- Into the subring (junk value `0` off it). -/
def toGood (N : ℕ) (r : ℚ) : goodSub N := if h : Nat.Coprime r.den N then ⟨r, h⟩ else 0

/-- The total reduction map `ℚ → ZMod N` (the ring hom `redHom` on `goodSub N`). -/
def redQ (N : ℕ) (r : ℚ) : ZMod N := redHom N (toGood N r)

variable {m : ℕ}

end Reduction

/-! ## 3. `FacAtZ`: the factorization only mod `N`; integer lifts; the bridge -/

section Bridge

variable {m : ℕ}

/-- **Generalized Lemma 10 data mod `N`.** As `FacAt`, but `A = E₁ Π E₂` is only required in
`ZMod N` (this is all `shear_via_pivots_array` uses). -/
def FacAtZ (N : ℕ) (fac : IMat m → IMat m × IMat m × IMat m) (A : IMat m) : Prop :=
  (∀ i j, i < j → (fac A).1 i j = 0) ∧ (∀ i j, i < j → (fac A).2.2 i j = 0) ∧
  (∀ i, IsUnit (((fac A).1 i i : ℤ) : ZMod N)) ∧ (∀ i, IsUnit (((fac A).2.2 i i : ℤ) : ZMod N)) ∧
  (∀ i j, (fac A).2.1 i j = 0 ∨ (fac A).2.1 i j = 1) ∧
  castM N A = castM N (fac A).1 * castM N (fac A).2.1 * castM N (fac A).2.2

/-- Integer lift of the reduction mod `Q` (the representative in `[0, Q)`). -/
def liftZ (Q : ℕ) (r : ℚ) : ℤ := ((redQ Q r).val : ℤ)

/-- The `Π` of `ratFac A` as an integer `0/1` matrix. -/
noncomputable def ratPi (A : IMat m) : IMat m :=
  fun i j => if (ratFac A).2.1 i j = 1 then 1 else 0

/-- The factorization data handed to Prop. 11 at modulus `Q = q ^ b`: integer lifts of the
reductions of the rational `E`'s, and the integer `0/1` matrix `Π`. -/
noncomputable def facZ (Q : ℕ) (A : IMat m) : IMat m × IMat m × IMat m :=
  ((ratFac A).1.map (liftZ Q), ratPi A, (ratFac A).2.2.map (liftZ Q))

end Bridge

/-! ## 4. Prop. 11 from `FacAtZ` (re-proved chain, reusing `AModelProp11`'s lemmas) -/

section Frames

variable {Wr R M : Type*} [Fintype Wr] [DecidableEq Wr] [AddGroup M]

/-- The multiset of all frame changes of a run (the finite set of matrices to factor). -/
def frameMS (gates : List (NGate Wr R M)) (lab Mout : Wr → M) : Multiset M :=
  frameSum (fun a => ({a} : Multiset M)) gates lab Mout

end Frames

section Level

variable {R : Type*} [CommRing R] {m W : ℕ}

/-- The circuit hypotheses of Prop. 11 at field size `N`, with `FacAtZ` in place of `FacAt`. -/
structure OKZ (C : Circuit R m W) (N : ℕ) : Prop where
  gatesOK : ∀ g ∈ C.gates, g.OK
  route : ∀ (v : Fin W → R) w, scalarRun C.gates v (C.ρ w) = v w
  unit : ∀ w, C.Mout (C.ρ w) - C.Min w = 1
  frames : FramesOK (FacAtZ N C.fac) C.gates C.Min C.Mout

variable (C : Circuit R m W) {N : ℕ} [NeZero N] (P L' : ℕ) (A B Cs : Shape)

variable [NeZero W]

end Level

section Recursion

variable {R : Type*} [CommRing R] {m W : ℕ} [NeZero W]

end Recursion

section Cost

variable {R : Type*} [CommRing R] {m W : ℕ}

/-- The only property of the factorization used by the cost analysis: `Π` is `0/1`. -/
def PiOK (fac : IMat m → IMat m × IMat m × IMat m) (A : IMat m) : Prop :=
  ∀ i j, (fac A).2.1 i j = 0 ∨ (fac A).2.1 i j = 1

end Cost

section Assembled

variable {R : Type*} [CommRing R] {m W : ℕ} [NeZero W]

/-- The circuit with its factorization data replaced. -/
def withFac (C : Circuit R m W) (f : IMat m → IMat m × IMat m × IMat m) : Circuit R m W :=
  { C with fac := f }

end Assembled

end AModelFac

end

section
-- Source module: Solutions.AModelL12Basic
/-!
# Lemma 12 (arbitrary-width interchange), generic array-model lemmas

1. `decode_append`, `encode_append_len`; `pull_prefix` (a pull acting only after a prefix does
   not see how the prefix is split into fields), `pull_blockSwap_prefix`.
2. Embedding a small shape into a pointwise larger one (`embed`, zero/default filled) and
   extracting back (`pull T S id`): `extract_pull_embed` — for every address map preserving
   the small-shape validity, embed → pull → extract = pull on the original.
   This is both the row padding `R → R'` and the binary padding `[2^u] → [q^e]`.
3. Adjacent block swaps compose (`blockSwap_blockSwap_adj`); a list of aligned pieces
   (`runPieces`) composes to the swap of the whole chunks (`runPieces_eq`).
4. Block rotation `rotIdx` (the joining moves) on lists, validity, inverse.
-/

namespace AModelL12

open AModel

/-! ## 1. Prefix transport -/

section Prefix

end Prefix

/-! ## 2. Embedding into a larger shape (padding) and extraction -/

section Embed

variable {α : Type*}

/-- `S` is pointwise at most `T` (same number of fields). -/
def ShapeLE (S T : Shape) : Prop :=
  S.length = T.length ∧ ∀ k (h₁ : k < S.length) (h₂ : k < T.length), S[k] ≤ T[k]

open Classical in
/-- Embed an array on `S` into the larger shape `T`: an entry at a `T`-address that is a valid
`S`-address is copied, every other entry is the default `d` (zero padding). -/
noncomputable def embed (S T : Shape) (d : α) (xs : List α) : List α :=
  (List.range (vol T)).map fun j =>
    if Valid S (decode T j) then xs.getD (encode S (decode T j)) d else d

/-- Extraction (scanning away the padding): the pull `T → S` by the identity address map. -/
def extract (S T : Shape) (d : α) (ys : List α) : List α := pull T S id d ys

end Embed

/-! ## 3. Pieces: adjacent block swaps compose -/

section Pieces

/-- Run piece operations: piece `t` has offset `o` (sum of the earlier widths) and width `w t`. -/
def runPieces {α β : Type*} (step : ℕ → β → List α → List α) (w : β → ℕ) :
    ℕ → List β → List α → List α
  | _, [], xs => xs
  | o, t :: ts, xs => runPieces step w (o + w t) ts (step o t xs)

end Pieces

/-! ## 4. Block rotation (the joining moves) -/

section Rotation

/-- Index map of the rotation `A ++ B ++ C ++ D ↦ A ++ C ++ B ++ D` with `|A| = s`, `|B| = b`,
`|C| = c` (digit `k` of the result is digit `rotIdx s b c k` of the input). -/
def rotIdx (s b c k : ℕ) : ℕ :=
  if s ≤ k ∧ k < s + c then k + b else if s + c ≤ k ∧ k < s + c + b then k - c else k

end Rotation

end AModelL12

end

section
-- Source module: Solutions.AModelL12
/-!
# Lemma 12 (arbitrary-width interchange), array model: the procedure and its denotation

Radix-`q` layer. The array lives on `qShape Pre G Suf q e = Pre ++ [q]*e ++ G ++ [q]*e ++ Suf`
(prefix, chunk `h`, gap, chunk `d`, suffix; the chunks are `e` radix-`q` digits each).
The procedure `l12q` (paper, proof of Lemma 12):
1. interchange the top `ρ` digits of the two chunks one pair at a time (`highSwaps`, `ρ` primitive
   single-digit swaps);
2. join: move `h⁺` next to `d⁺` (a field rotation; the result lives on `joinShape`);
3. reshape (free): prefix ↦ one field `P = vol Pre`, `d⁺h⁺` ↦ one row field `R = q^(2ρ)`;
4. pad the row field to `R' = padRows W Kp R` (a multiple of `W ^ Kp`, zero rows appended);
5. write `n = e - ρ` in base `m` and call the Prop. 11 oracle `child j` on each corresponding pair of
   pieces of width `m ^ j` (`pieceLevels`, `runPieces`), all with the shared row field `R'`;
6. unpad (scan away the padded rows), reshape back, reverse the joining rotation.
If `ρ > e` (bounded widths) the procedure does `e` single-digit interchanges instead.

Main theorem `l12q_correct`: for every Prop.-11 oracle `child` satisfying the block-swap spec
(as `prop11`'s first conjunct gives it, when `W ^ j ∣ rows`), `l12q` is the pull by the full
chunk interchange `blockSwap a (a + e + |G|) e`.

Binary layer (`l12bin`): pad each `[2^u]` chunk numerically to `[q^e]` (`embed`, zero outside),
run `l12q` on the radix-`q` digits (reshape is free), extract. `l12bin_correct`: the result is
the pull by the single-field interchange `swapPos a (a + 1 + |G|)` on
`Pre ++ [2^u] ++ G ++ [2^u] ++ Suf` — the paper's `(p,h,g,d,z) ↦ (p,d,g,h,z)`.
-/

namespace AModelL12

open AModel

/-! ## 1. Shapes -/

section Shapes

/-- `[q]*n ++ G ++ [q]*n ++ Suf`: the two remaining chunks with the gap and suffix. -/
def restShape (G Suf : Shape) (q n : ℕ) : Shape :=
  List.replicate n q ++ (G ++ (List.replicate n q ++ Suf))

/-- The radix-`q` shape `Pre ++ [q]*e ++ G ++ [q]*e ++ Suf`. -/
def qShape (Pre G Suf : Shape) (q e : ℕ) : Shape := Pre ++ restShape G Suf q e

/-- The shape after joining: `Pre ++ (d⁺ h⁺) ++ h⁻ ++ G ++ d⁻ ++ Suf` (all digits of size `q`). -/
def joinShape (Pre G Suf : Shape) (q ρ n : ℕ) : Shape :=
  (Pre ++ List.replicate (ρ + ρ) q) ++ restShape G Suf q n

/-- The shape with a single prefix field `P`, a row field `R`, and the remaining chunks. -/
def rowShape (P R : ℕ) (G Suf : Shape) (q n : ℕ) : Shape := P :: R :: restShape G Suf q n

/-- Row count padded to the next multiple of `W ^ K`. -/
def padRows (W K R : ℕ) : ℕ := W ^ K * ((R + W ^ K - 1) / W ^ K)

end Shapes

/-! ## 2. The pieces -/

section PiecesDef

/-- The piece levels: `digit m n j` pieces of width `m ^ j` for every base-`m` digit position `j`. -/
def pieceLevels (m n : ℕ) : List ℕ :=
  (List.range (Nat.digits m n).length).flatMap fun j => List.replicate (digit m n j) j

end PiecesDef

/-! ## 3. The procedure -/

section Procedure

variable {α : Type*}

/-- The oracle type: `child j P L A B Cs` acts on `P :: L :: hdShape A B Cs (m ^ j) q`
(the signature of `prog C q`). -/
abbrev Oracle (α : Type*) := ℕ → ℕ → ℕ → Shape → Shape → Shape → List α → List α

/-- The oracle specification (`prop11`'s first conjunct): at every level `j ≤ K`, with `W ^ j`
dividing the row count, the call is the block swap of the two `m ^ j`-digit chunks. -/
def OracleSpec (child : Oracle α) (d : α) (W m q K : ℕ) : Prop :=
  ∀ j ≤ K, ∀ (P L : ℕ) (A B Cs : Shape) (ys : List α), W ^ j ∣ L →
    ys.length = vol (P :: L :: hdShape A B Cs (m ^ j) q) →
    child j P L A B Cs ys =
      pull (P :: L :: hdShape A B Cs (m ^ j) q) (P :: L :: hdShape A B Cs (m ^ j) q)
        (blockSwap (A.length + 2) (A.length + B.length + m ^ j + 2) (m ^ j)) d ys

/-- `r` primitive single-digit interchanges of positions `i + t ↔ j + t`, `t < r`. -/
def highSwaps (sh : Shape) (i j r : ℕ) (d : α) (xs : List α) : List α :=
  (List.range r).foldl (fun ys t => pull sh sh (swapPos (i + t) (j + t)) d ys) xs

/-- The call on the piece at offset `o` of level `j` (width `m ^ j`): prefix `P`, rows `R'`,
`A = [q]*o`, `B = [q]*(n-o-m^j) ++ G ++ [q]*o`, `Cs = [q]*(n-o-m^j) ++ Suf`. -/
def pieceStep (child : Oracle α) (P R' : ℕ) (G Suf : Shape) (q m n o j : ℕ) (ys : List α) :
    List α :=
  child j P R' (List.replicate o q)
    (List.replicate (n - o - m ^ j) q ++ (G ++ List.replicate o q))
    (List.replicate (n - o - m ^ j) q ++ Suf) ys

/-- **The radix-`q` procedure of Lemma 12.** -/
noncomputable def l12q (child : Oracle α) (W m q Kp ρ : ℕ) (Pre G Suf : Shape) (e : ℕ) (d : α)
    (xs : List α) : List α :=
  if ρ ≤ e then
    let a := Pre.length
    let g := G.length
    let n := e - ρ
    let sh0 := qShape Pre G Suf q e
    let sh1 := joinShape Pre G Suf q ρ n
    let S := rowShape (vol Pre) (q ^ (ρ + ρ)) G Suf q n
    let R' := padRows W Kp (q ^ (ρ + ρ))
    let T := rowShape (vol Pre) R' G Suf q n
    let xs1 := highSwaps sh0 a (a + e + g) ρ d xs
    let xs2 := pull sh0 sh1 (permIdx (rotIdx (a + ρ) ρ (n + g))) d xs1
    let xs3 := AModelL12.embed S T d xs2
    let xs4 := runPieces (pieceStep child (vol Pre) R' G Suf q m n) (fun j => m ^ j) 0
      (pieceLevels m n) xs3
    let xs5 := extract S T d xs4
    pull sh1 sh0 (permIdx (rotIdx (a + ρ) (n + g) ρ)) d xs5
  else
    highSwaps (qShape Pre G Suf q e) Pre.length (Pre.length + e + G.length) e d xs

end Procedure

/-! ## 4. Correctness -/

section Correct

variable {α : Type*}

end Correct

/-! ## 5. Binary chunks: numeric padding `[2^u] → [q^e]` -/

section Binary

variable {α : Type*}

/-- `Pre ++ [F] ++ G ++ [F] ++ Suf`: the two chunks as single fields of size `F`. -/
def binShape (Pre G Suf : Shape) (F : ℕ) : Shape := Pre ++ (F :: (G ++ (F :: Suf)))

/-- **The binary procedure of Lemma 12**: pad both `[2^u]` chunks numerically to `[q^e]` (zero
outside), run the radix-`q` procedure on the digits (reshape is free), scan away the padding. -/
noncomputable def l12bin (child : Oracle α) (W m q Kp ρ e u : ℕ) (Pre G Suf : Shape) (d : α)
    (xs : List α) : List α :=
  extract (binShape Pre G Suf (2 ^ u)) (binShape Pre G Suf (q ^ e)) d
    (l12q child W m q Kp ρ Pre G Suf e d
      (AModelL12.embed (binShape Pre G Suf (2 ^ u)) (binShape Pre G Suf (q ^ e)) d xs))

end Binary

end AModelL12

end

section
-- Source module: Solutions.AModelL12Cost
/-!
# Lemma 12 (arbitrary-width interchange), array model: cost

Cost model (analogous to `progCost`): every primitive pass over the array is charged `c · V`
(`V` its current volume): each of the `ρ` single-digit swaps, each of the `ρ` single-field moves
of the joining and of the separating rotation, the row padding and unpadding (on the padded
volume `V'`). Each Prop. 11 call on a piece of level `j` costs `Tc j V'` (oracle cost at the
padded volume). The fallback (`ρ > e`) costs `e` passes.

* `l12qCost_le`: if `Tc j V ≤ K₀ V (m^j)^τ`, `W ^ Kp ≤ q^(2ρ)` (so `V' < 2V`) and `ρ ≤ D e^τ`,
  then `l12qCost ≤ K'' V e^τ` with `K''` depending only on `m, τ, K₀, c, D`.
* `clog_le_rpow`: `⌈log_m e⌉ ≤ C e^τ`, so `ρ = ρ₀ ⌈log_m e⌉` satisfies the hypothesis.
* `l12binCost`: adds the binary padding / extraction passes.
-/

namespace AModelL12

open AModel Finset

section Cost

/-- **Cost model of `l12q`.** `V = vol (qShape …)`, `V' = vol` of the padded row shape. -/
noncomputable def l12qCost (Tc : ℕ → ℝ → ℝ) (c : ℝ) (W m q Kp ρ : ℕ) (Pre G Suf : Shape)
    (e : ℕ) : ℝ :=
  if ρ ≤ e then
    -- `ρ` single-digit swaps, `ρ` joining moves
    c * (vol (qShape Pre G Suf q e) : ℝ) * ρ + c * (vol (qShape Pre G Suf q e) : ℝ) * ρ +
    -- padding
    c * (vol (rowShape (vol Pre) (padRows W Kp (q ^ (ρ + ρ))) G Suf q (e - ρ)) : ℝ) +
    -- the Prop. 11 calls on the pieces
    ((pieceLevels m (e - ρ)).map fun j =>
      Tc j (vol (rowShape (vol Pre) (padRows W Kp (q ^ (ρ + ρ))) G Suf q (e - ρ)) : ℝ)).sum +
    -- unpadding, `ρ` separating moves
    c * (vol (rowShape (vol Pre) (padRows W Kp (q ^ (ρ + ρ))) G Suf q (e - ρ)) : ℝ) +
    c * (vol (qShape Pre G Suf q e) : ℝ) * ρ
  else c * (vol (qShape Pre G Suf q e) : ℝ) * e

/-- **Cost model of the binary layer**: the padding and extraction passes (on the padded
volume `vol (binShape … (q^e)) = vol (qShape … e)`) plus the radix-`q` procedure. -/
noncomputable def l12binCost (Tc : ℕ → ℝ → ℝ) (c : ℝ) (W m q Kp ρ e : ℕ)
    (Pre G Suf : Shape) : ℝ :=
  c * (vol (binShape Pre G Suf (q ^ e)) : ℝ) + l12qCost Tc c W m q Kp ρ Pre G Suf e +
    c * (vol (binShape Pre G Suf (q ^ e)) : ℝ)

end Cost

end AModelL12

end

section
-- Source module: Solutions.AModelRat
/-!
# Proposition 11 for circuits with rational labels

The paper's circuit (Prop. 7) has rational labels (projections `P_U` with denominators). This
file proves Prop. 11 for a circuit `QCircuit R m W` whose labels are `Matrix (Fin m) (Fin m) ℚ`.

Route (scaling, with a scaled frame action): let `D` be a common denominator of all labels.
The integer circuit `C.toCircuit Q` has labels `D • M` (exact integers, `intM`), so every integer
frame change is exactly `D • A_e` (`A_e` the rational frame change) and the endpoint differences
are exactly `I0 = D • 1` over `ℤ`. The frame action at field size `N` is `φ_N X = D⁻¹ • (X mod N)`
(`phiD`), with `φ_N I0 = 1` and `φ_N (D • A) = A mod N`. The Lemma 10 data of an integer frame
change `X` is the lift mod `Q = q ^ m ^ K` of the rational Bruhat factorization of
`D⁻¹ • X = A_e` (`facD`). The correctness chain of `AModelFac` is re-proved with a general
additive frame action `φ` (`FacAtφ`, `OKφ`, `implΦ_eqφ`, ..., `prog_correctφ`); the cost chain
is reused unchanged (`prog_costZ`, `levelCost_eqZ`).

1. Rational Lemma 9 for rational matrices (`ratFacQ`, `facQ`), the bridge `facAtφ_facD`.
2. Pair-level frame bookkeeping (`FramesOK2`, `frameSum2`) and label maps.
3. The chain with a general frame action.
4. `QCircuit`, `QCircuit.toCircuit`, and `prop11_ratLabels`.
-/

namespace AModelRat

open Matrix AModel AModelFac

/-! ## 1. Rational factorizations of rational matrices; the bridge -/

section RatFac

variable {m : ℕ}

lemma exists_ratFacQ (A : QMat m) : ∃ F, RatFacSpec A F := by
  obtain ⟨E₁, P, E₂, h⟩ := bruhat_lower A
  exact ⟨(E₁, P, E₂), h⟩

/-- The rational Lemma 9 factorization of a rational matrix, chosen once. -/
noncomputable def ratFacQ (A : QMat m) : QMat m × QMat m × QMat m := (exists_ratFacQ A).choose

/-- `q` avoids the denominators of `A` and is good for `ratFacQ A`. -/
def GoodAtQ (q : ℕ) (A : QMat m) : Prop :=
  (∀ i j, Nat.Coprime (A i j).den q) ∧ GoodAt q (ratFacQ A)

/-- A bound beyond which every prime is good for `A`. -/
noncomputable def facBoundQ (A : QMat m) : ℕ :=
  (∑ i, ∑ j, (((ratFacQ A).1 i j).den + ((ratFacQ A).2.2 i j).den + (A i j).den)) +
  ∑ i, (((ratFacQ A).1 i i).num.natAbs + ((ratFacQ A).2.2 i i).num.natAbs)

/-- The `Π` of `ratFacQ A` as an integer `0/1` matrix. -/
noncomputable def piQ (A : QMat m) : IMat m :=
  fun i j => if (ratFacQ A).2.1 i j = 1 then 1 else 0

/-- Lemma 10 data at modulus `Q` for the rational matrix `A`. -/
noncomputable def facQ (Q : ℕ) (A : QMat m) : IMat m × IMat m × IMat m :=
  ((ratFacQ A).1.map (liftZ Q), piQ A, (ratFacQ A).2.2.map (liftZ Q))

/-- Lemma 10 data of an integer frame change `X = D • A`: that of `A = D⁻¹ • X`. -/
noncomputable def facD (Q D : ℕ) (X : IMat m) : IMat m × IMat m × IMat m :=
  facQ Q ((D : ℚ)⁻¹ • ratM X)

/-- The scaled frame action `X ↦ D⁻¹ • (X mod N)`. -/
noncomputable def phiD (N D : ℕ) : IMat m →+ Matrix (Fin m) (Fin m) (ZMod N) where
  toFun X := ((D : ZMod N)⁻¹) • castM N X
  map_zero' := by ext i j; simp [castM]
  map_add' X Y := by ext i j; simp [castM, mul_add]

/-- Lemma 10 data mod `N` for a general additive frame action `φ`: as `FacAtZ`, with the
identity `φ A = E₁ Π E₂` in `ZMod N`. -/
def FacAtφ (N : ℕ) (φ : IMat m →+ Matrix (Fin m) (Fin m) (ZMod N))
    (fac : IMat m → IMat m × IMat m × IMat m) (A : IMat m) : Prop :=
  (∀ i j, i < j → (fac A).1 i j = 0) ∧ (∀ i j, i < j → (fac A).2.2 i j = 0) ∧
  (∀ i, IsUnit (((fac A).1 i i : ℤ) : ZMod N)) ∧ (∀ i, IsUnit (((fac A).2.2 i i : ℤ) : ZMod N)) ∧
  (∀ i j, (fac A).2.1 i j = 0 ∨ (fac A).2.1 i j = 1) ∧
  φ A = castM N (fac A).1 * castM N (fac A).2.1 * castM N (fac A).2.2

end RatFac

/-! ## 2. Pair-level frame bookkeeping -/

section Pairs

variable {Wr R M M' : Type*} [Fintype Wr] [DecidableEq Wr]

/-- `FramesOK` on (new label, old label) pairs. -/
def FramesOK2 (Pr : M → M → Prop) : List (NGate Wr R M) → (Wr → M) → (Wr → M) → Prop
  | [], lab, Mout => ∀ w, Pr (Mout w) (lab w)
  | g :: gs, lab, Mout => (∀ w ∈ g.2.1, Pr g.2.2 (lab w)) ∧ FramesOK2 Pr gs (labStep lab g) Mout

/-- `frameSum` on pairs. -/
def frameSum2 {β : Type*} [AddCommMonoid β] (F : M → M → β) :
    List (NGate Wr R M) → (Wr → M) → (Wr → M) → β
  | [], lab, Mout => ∑ w, F (Mout w) (lab w)
  | g :: gs, lab, Mout => (∑ w ∈ g.2.1, F g.2.2 (lab w)) + frameSum2 F gs (labStep lab g) Mout

/-- Relabel the gates by `f`. -/
def mapG (f : M → M') (gates : List (NGate Wr R M)) : List (NGate Wr R M') :=
  gates.map fun g => (g.1, g.2.1, f g.2.2)

omit [Fintype Wr] in
lemma framesOK2_mono {P P' : M → M → Prop} (h : ∀ a b, P a b → P' a b) :
    ∀ (gates : List (NGate Wr R M)) (lab Mout : Wr → M),
      FramesOK2 P gates lab Mout → FramesOK2 P' gates lab Mout := by
  intro gates
  induction gates with
  | nil => intro lab Mout hF w; exact h _ _ (hF w)
  | cons g gs ih => intro lab Mout hF; exact ⟨fun w hw => h _ _ (hF.1 w hw), ih _ _ hF.2⟩

end Pairs

/-! ## 3. The correctness chain with a general frame action -/

section Level

variable {R : Type*} [CommRing R] {m W : ℕ}

/-- The circuit hypotheses at field size `N` with frame action `φ` and endpoint difference `I0`
(`φ I0 = 1`). -/
structure OKφ (C : Circuit R m W) (N : ℕ) (φ : IMat m →+ Matrix (Fin m) (Fin m) (ZMod N))
    (I0 : IMat m) : Prop where
  gatesOK : ∀ g ∈ C.gates, g.OK
  route : ∀ (v : Fin W → R) w, scalarRun C.gates v (C.ρ w) = v w
  unit : ∀ w, C.Mout (C.ρ w) - C.Min w = I0
  one : φ I0 = 1
  frames : FramesOK (FacAtφ N φ C.fac) C.gates C.Min C.Mout

variable (C : Circuit R m W) {N : ℕ} [NeZero N] (P L' : ℕ) (A B Cs : Shape)

variable [NeZero W]

end Level

section Recursion

variable {R : Type*} [CommRing R] {m W : ℕ} [NeZero W]

end Recursion

/-! ## 4. Rational-label circuits and Proposition 11 -/

section QCirc

variable {m : ℕ}

/-- Product of the denominators of the entries of `M`. -/
def matDen (M : QMat m) : ℕ := ∏ i, ∏ j, (M i j).den

/-- All denominators of `M` divide `D`. -/
def LabOK (D : ℕ) (M : QMat m) : Prop := ∀ i j, (M i j).den ∣ D

/-- The integer matrix `D • M` (exact when `LabOK D M`). -/
def intM (D : ℕ) (M : QMat m) : IMat m := fun i j => ((D : ℚ) * M i j).num

end QCirc

/-- A circuit with **rational** matrix labels (the paper's Prop. 7 circuit): gates on `W` wires
with labels in `Matrix (Fin m) (Fin m) ℚ`, input/output labels, and the routing `ρ`. -/
structure QCircuit (R : Type*) [CommRing R] (m W : ℕ) where
  gates : List (NGate (Fin W) R (QMat m))
  Min : Fin W → QMat m
  Mout : Fin W → QMat m
  ρ : Fin W → Fin W

namespace QCircuit

variable {R : Type*} [CommRing R] {m W : ℕ}

/-- A common denominator of all labels. -/
def den (C : QCircuit R m W) : ℕ :=
  (C.gates.map fun g => matDen g.2.2).prod * (∏ w, matDen (C.Min w)) * ∏ w, matDen (C.Mout w)

/-- **The integer circuit run by the program at modulus `Q`**: labels `D • M` (`D = C.den`),
Lemma 10 data `facD Q D` (lifts mod `Q` of the rational Bruhat factors of `D⁻¹ • X`). -/
noncomputable def toCircuit (C : QCircuit R m W) (Q : ℕ) : Circuit R m W where
  gates := mapG (intM C.den) C.gates
  Min w := intM C.den (C.Min w)
  Mout w := intM C.den (C.Mout w)
  ρ := C.ρ
  fac := facD Q C.den

end QCircuit

section Main

variable {R : Type*} [CommRing R] {m W : ℕ} [NeZero W]

end Main

end AModelRat

end

section
-- Source module: Solutions.AModelL12Main
/-!
# Lemma 12 (arbitrary-width interchange), array model: final statements

* `lemma12`: parametric in a Prop.-11 oracle family `child K` (level bound `K`) satisfying the
  block-swap spec (`OracleSpec`) with cost `Tc K j V ≤ K₀ V (m^j)^τ`. With the parameters of the
  paper — `e = ⌈log_q 2^u⌉` (`eOf`), `Kp = ⌈log_m e⌉` (`kOf`), `ρ = ρ₀ Kp` where
  `W ≤ q^(2ρ₀)` (so `q^(2ρ) ≥ W^Kp`) — the procedure `lemma12Prog` is the interchange
  `(…, h, g, d, …) ↦ (…, d, g, h, …)` of the two `[2^u]` fields of `Pre ++ [2^u] ++ G ++ [2^u] ++ Suf`,
  for every `u ≥ 1` and all shapes, and its cost is `≤ K V u^τ`.
* `lemma12_paper`: the layout `[P] × [2^u] × [G] × [2^u] × [B]`, permutation `swapPos 1 3`.
* `lemma12_rat`: the oracle instantiated with the Prop. 11 program of `prop11_ratLabels`
  (`prog (C.toCircuit (q ^ m ^ K)) q`, cost `progCost`), for every large prime `q`, `ρ₀ = W`.
-/

namespace AModelL12

open AModel

section Final

variable {α : Type*}

/-- `e = ⌈log_q 2^u⌉`: the least `e` with `2^u ≤ q^e`. -/
def eOf (q u : ℕ) : ℕ := Nat.clog q (2 ^ u)

/-- `Kp = ⌈log_m e⌉`: the least `k` with `e ≤ m^k`. -/
def kOf (m q u : ℕ) : ℕ := Nat.clog m (eOf q u)

/-- **The Lemma 12 procedure** with the paper's parameters (`ρ = ρ₀ Kp`). -/
noncomputable def lemma12Prog (child : ℕ → Oracle α) (W m q ρ₀ : ℕ) (Pre G Suf : Shape) (u : ℕ)
    (d : α) (xs : List α) : List α :=
  l12bin (child (kOf m q u)) W m q (kOf m q u) (ρ₀ * kOf m q u) (eOf q u) u Pre G Suf d xs

/-- Its cost under the model of `l12binCost` with oracle costs `Tc (kOf m q u)`. -/
noncomputable def lemma12Cost (Tc : ℕ → ℕ → ℝ → ℝ) (c : ℝ) (W m q ρ₀ : ℕ) (Pre G Suf : Shape)
    (u : ℕ) : ℝ :=
  l12binCost (Tc (kOf m q u)) c W m q (kOf m q u) (ρ₀ * kOf m q u) (eOf q u) Pre G Suf

end Final

/-! ## Instantiation with the Prop. 11 program (`prop11_ratLabels`) -/

section Rat

open AModelRat AModelFac

end Rat

end AModelL12

end

section
-- Source module: Solutions.ButterflyTensor
/-!
# §5.1 Phase-frame identities and tensor products on selected bits

The `2 × 2` matrices of §5.1 of the OpenAI preprint "Integer multiplication below n log n":
`H0 = ½[[1,1],[1,-1]]`, `S = diag(1,i)`, `Z = diag(1,-1)`, `X = [[0,1],[1,0]]`,
`C = HSH = aI + bX` with `a = (1+i)/2`, `b = (1-i)/2`, and the identities (13), (14), (17), (18).

Tensor products over an arbitrary finite set `ι` of selected bits:
`tensF A x y = ∏ i, A i (x i) (y i)` for a family `A : ι → M₂`.
-/

namespace Butterfly

open Matrix Complex Finset

/-- `2 × 2` complex matrices. -/
abbrev M2 := Matrix (Fin 2) (Fin 2) ℂ

/-- `H0 = ½ [[1,1],[1,-1]]`, the normalized butterfly `(u,v) ↦ ((u+v)/2, (u-v)/2)`. -/
noncomputable def H0 : M2 := !![1 / 2, 1 / 2; 1 / 2, -1 / 2]
/-- `S = diag(1, i)`. -/
def Sm : M2 := !![1, 0; 0, I]
/-- `Z = diag(1, -1)`. -/
def Zm : M2 := !![1, 0; 0, -1]
/-- The bit flip `X`. -/
def Xm : M2 := !![0, 1; 1, 0]
/-- `a = (1+i)/2`. -/
noncomputable def ca : ℂ := (1 + I) / 2
/-- `b = (1-i)/2`. -/
noncomputable def cb : ℂ := (1 - I) / 2
/-- The auxiliary two-point kernel `C = aI + bX` (13). -/
noncomputable def Cm : M2 := Butterfly.ca • (1 : M2) + Butterfly.cb • Xm
/-- Its inverse `C⁻¹ = bI + aX` (14). -/
noncomputable def Cinv : M2 := Butterfly.cb • (1 : M2) + Butterfly.ca • Xm

lemma Cm_eq : Butterfly.Cm = !![Butterfly.ca, Butterfly.cb; Butterfly.cb, Butterfly.ca] := by
  ext i j; fin_cases i <;> fin_cases j <;> simp [Butterfly.Cm, Xm]

lemma Cinv_eq : Cinv = !![Butterfly.cb, Butterfly.ca; Butterfly.ca, Butterfly.cb] := by
  ext i j; fin_cases i <;> fin_cases j <;> simp [Cinv, Xm]

/-- (14): `C⁻¹ = -i Z C Z`. -/
theorem Cinv_eq_ZCZ : Cinv = (-I) • (Zm * Butterfly.Cm * Zm) := by
  rw [Cm_eq, Cinv_eq]
  ext i j; fin_cases i <;> fin_cases j <;>
    simp [Zm, Matrix.mul_apply, Fin.sum_univ_two, Butterfly.ca, Butterfly.cb] <;> ring_nf <;> simp [I_sq] <;> ring

/-! ## Tensor products on selected bits -/

section Tensor

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

/-- Tensor product of a family of `2 × 2` matrices over the selected bits `ι`. -/
def tensF (A : ι → M2) : Matrix (ι → Fin 2) (ι → Fin 2) ℂ := fun x y => ∏ i, A i (x i) (y i)

/-- Tensor power of one `2 × 2` matrix. -/
def tens (A : M2) : Matrix (ι → Fin 2) (ι → Fin 2) ℂ := tensF (fun _ => A)

/-- The tensor product is multiplicative. -/
theorem tensF_mul (A B : ι → M2) : tensF (fun i => A i * B i) = tensF A * tensF B := by
  ext x z
  simp only [tensF, Matrix.mul_apply]
  rw [Fintype.prod_sum]
  simp only [← Finset.prod_mul_distrib]

theorem tensF_smul (c : ι → ℂ) (A : ι → M2) :
    tensF (fun i => c i • A i) = (∏ i, c i) • tensF A := by
  ext x y
  simp [tensF, Finset.prod_mul_distrib]

theorem tens_mul (A B : M2) : (Butterfly.tens (A * B) : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) =
    Butterfly.tens A * Butterfly.tens B := tensF_mul _ _

theorem tens_smul (c : ℂ) (A : M2) :
    (Butterfly.tens (c • A) : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) = c ^ Fintype.card ι • Butterfly.tens A := by
  rw [Butterfly.tens, tensF_smul (fun _ => c), Finset.prod_const, Finset.card_univ]; rfl

/-- (17): `(C⁻¹)^{⊗f} = (-i)^f Z^{⊗f} C^{⊗f} Z^{⊗f}`. -/
theorem tens_Cinv : (Butterfly.tens Cinv : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) =
    (-I) ^ Fintype.card ι • (Butterfly.tens Zm * Butterfly.tens Butterfly.Cm * Butterfly.tens Zm) := by
  rw [Cinv_eq_ZCZ, tens_smul, tens_mul, tens_mul]

/-- Number of selected bits equal to one. -/
def wt (x : ι → Fin 2) : ℕ := ∑ i, (x i : ℕ)

/-- The phase diagonal `S_D`: `(S_D z)(x) = i^{x_1+⋯+x_D} z(x)`. -/
noncomputable def SD : Matrix (ι → Fin 2) (ι → Fin 2) ℂ := diagonal fun x => I ^ Butterfly.wt x

/-- The sign diagonal `Z^{⊗}`: `(-1)^{x_1+⋯+x_D}`. -/
noncomputable def ZD : Matrix (ι → Fin 2) (ι → Fin 2) ℂ := diagonal fun x => (-1) ^ Butterfly.wt x

lemma diag2_tens (u v : ℂ) : (Butterfly.tens !![u, 0; 0, v] : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) =
    diagonal fun x => ∏ i, (if x i = 0 then u else v) := by
  ext x y
  simp only [Butterfly.tens, tensF, diagonal_apply]
  by_cases h : x = y
  · subst h; simp only [if_true]
    refine Finset.prod_congr rfl fun i _ => ?_
    rcases Fin.exists_fin_two.1 ⟨x i, rfl⟩ with h0 | h1 <;> simp [*]
  · simp only [h, if_false]
    obtain ⟨i, hi⟩ : ∃ i, x i ≠ y i := by
      by_contra hc; push Not at hc; exact h (funext hc)
    apply Finset.prod_eq_zero (mem_univ i)
    rcases Fin.exists_fin_two.1 ⟨x i, rfl⟩ with h0 | h0 <;>
    rcases Fin.exists_fin_two.1 ⟨y i, rfl⟩ with h1 | h1 <;> simp_all

lemma prod_phase (c : ℂ) (x : ι → Fin 2) :
    ∏ i, (if x i = 0 then (1 : ℂ) else c) = c ^ Butterfly.wt x := by
  rw [Butterfly.wt, ← Finset.prod_pow_eq_pow_sum]
  congr 1; ext i
  rcases Fin.exists_fin_two.1 ⟨x i, rfl⟩ with h | h <;> simp [h]

theorem tens_Zm : (Butterfly.tens Zm : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) = ZD := by
  rw [Zm, diag2_tens, ZD]; congr 1; funext x; exact prod_phase (-1) x

/-- (17) with the `S_D`/`Z_D` diagonals named. -/
theorem tens_Cinv' : (Butterfly.tens Cinv : Matrix (ι → Fin 2) (ι → Fin 2) ℂ) =
    (-I) ^ Fintype.card ι • (ZD * Butterfly.tens Butterfly.Cm * ZD) := by
  rw [tens_Cinv, tens_Zm]

end Tensor

/-! ## (15): address permutations conjugate translations -/

section Translate

variable {K V : Type*} [Ring K] [AddCommGroup V] [Module K V] {β : Type*}

end Translate

/-! ## Elementary row additions generate `GL(m, F₂)` -/

end Butterfly

end

section
-- Source module: Solutions.ButterflyLayer
/-!
# §5 The simultaneous butterfly layer as an array denotation

The array has complete address set `[P] × [2^K]^D × [S]` (shape `layerShape P K D S`), in
lexicographic order, one record per address. Records live in a complex vector space `M`
(a record of `r` coefficients is `Fin r → ℂ`; its sup norm is the maximum coefficient norm, and
every operator below acts coefficientwise because it is `ℂ`-linear on records).

On each of the `D` chunks, bit `ρ` (from the least significant end) is selected.
A matrix `T` on the selected bits acts by
`(selOp T xs)(a) = ∑_y T(bits a, y) • xs(a with selected bits y)`,
fixing every unselected bit and both spectator fields.

* `selOp_mul`, `selOp_one`, `selOp_list_prod`: composition = matrix product.
* `selOp_diagonal`: diagonals are pointwise maps (the phase `S_D`, the signs `Z^{⊗}`).
* `kernel_pair`: an individual selected-bit kernel is the paired map of (25) on the two
  aligned streams split at that bit.
* `layer_eq_kernels`: the layer `A^{⊗D}` is the composition of the `D` individual kernels.
* `butterflyLayer_eq` (18 on arrays): `H0`-layer `= b^D • S_D ∘ C-layer ∘ S_D`.
* `butterflyLayer_norm_le`: disk arrays go to disk arrays.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace Butterfly

open AModel Matrix Finset

/-! ## 1. Abstract selected-bit structure on a shape -/

/-- Selected bits of a shape: read them (`bits`) and overwrite them (`set`). -/
structure SelBits (sh : Shape) (ι : Type*) where
  bits : List ℕ → ι → Fin 2
  set : List ℕ → (ι → Fin 2) → List ℕ
  valid_set : ∀ a y, Valid sh a → Valid sh (set a y)
  bits_set : ∀ a y, Valid sh a → bits (set a y) = y
  set_set : ∀ a y z, Valid sh a → set (set a y) z = set a z
  set_bits : ∀ a, Valid sh a → set a (bits a) = a

variable {sh : Shape} {ι : Type*} [Fintype ι] [DecidableEq ι]
variable {M : Type*} [AddCommMonoid M] [Module ℂ M]

/-- A matrix on the selected bits acting on the array (coefficientwise on records). -/
noncomputable def selOp (Sb : SelBits sh ι) (T : Matrix (ι → Fin 2) (ι → Fin 2) ℂ)
    (xs : List M) : List M :=
  (List.range (vol sh)).map fun j =>
    ∑ y, T (Sb.bits (decode sh j)) y • xs.getD (encode sh (Sb.set (decode sh j) y)) 0

variable (Sb : SelBits sh ι)

/-! ## 2. Individual kernels and the full layer -/

/-- The individual kernel of `A` on selected bit `i`: `A` on bit `i`, identity elsewhere. -/
def kerAt (i : ι) (A : M2) : ι → M2 := Function.update (fun _ => 1) i A

/-- The normalized simultaneous butterfly layer `H0^{⊗D}`. -/
noncomputable def butterflyLayer (xs : List M) : List M := selOp Sb (Butterfly.tens H0) xs

/-- The `C`-layer `C^{⊗D}`. -/
noncomputable def cLayer (xs : List M) : List M := selOp Sb (Butterfly.tens Butterfly.Cm) xs

/-- The phase layer `S_D` (pointwise). -/
noncomputable def phaseLayer (xs : List M) : List M := selOp Sb SD xs

/-! ## 3. The concrete layout `[P] × [2^K]^D × [S]` with bit `ρ` selected in every chunk -/

/-- The shape `[P] × [2^K]^D × [S]`. -/
def layerShape (P K D S : ℕ) : Shape := P :: List.replicate D (2 ^ K) ++ [S]

/-- Bit `ρ` of `v` (from the least significant end). -/
def getBit (ρ v : ℕ) : Fin 2 := ⟨(v / 2 ^ ρ) % 2, Nat.mod_lt _ (by norm_num)⟩

/-- Overwrite bit `ρ` of `v` by `b`. -/
def setBit (ρ v : ℕ) (b : Fin 2) : ℕ := v / 2 ^ (ρ + 1) * 2 ^ (ρ + 1) + (b : ℕ) * 2 ^ ρ + v % 2 ^ ρ

section BitArith

lemma setBit_eq (ρ v : ℕ) (b : Fin 2) :
    setBit ρ v b = (2 * (v / 2 ^ (ρ + 1)) + b) * 2 ^ ρ + v % 2 ^ ρ := by
  simp only [setBit, pow_succ]; ring

lemma setBit_div (ρ v : ℕ) (b : Fin 2) : setBit ρ v b / 2 ^ ρ = 2 * (v / 2 ^ (ρ + 1)) + b := by
  rw [setBit_eq, Nat.add_comm, Nat.add_mul_div_right _ _ (by positivity),
    Nat.div_eq_of_lt (Nat.mod_lt _ (by positivity)), zero_add]

lemma setBit_mod (ρ v : ℕ) (b : Fin 2) : setBit ρ v b % 2 ^ ρ = v % 2 ^ ρ := by
  rw [setBit_eq, Nat.add_comm, Nat.add_mul_mod_self_right, Nat.mod_mod]

lemma getBit_setBit (ρ v : ℕ) (b : Fin 2) : getBit ρ (setBit ρ v b) = b := by
  ext; simp only [getBit, setBit_div]
  have := b.isLt; omega

lemma div_two_pow_succ (ρ v : ℕ) : v / 2 ^ (ρ + 1) = v / 2 ^ ρ / 2 := by
  rw [pow_succ, Nat.div_div_eq_div_mul]

lemma setBit_setBit (ρ v : ℕ) (b c : Fin 2) : setBit ρ (setBit ρ v b) c = setBit ρ v c := by
  have h1 : setBit ρ v b / 2 ^ (ρ + 1) = v / 2 ^ (ρ + 1) := by
    rw [div_two_pow_succ, setBit_div]; have := b.isLt; omega
  rw [setBit_eq ρ (setBit ρ v b) c, h1, setBit_mod, ← setBit_eq]

lemma setBit_getBit (ρ v : ℕ) : setBit ρ v (getBit ρ v) = v := by
  have h := Nat.div_add_mod v (2 ^ ρ)
  have h2 := Nat.div_add_mod (v / 2 ^ ρ) 2
  rw [setBit_eq, div_two_pow_succ]
  simp only [getBit]
  calc (2 * (v / 2 ^ ρ / 2) + (v / 2 ^ ρ) % 2) * 2 ^ ρ + v % 2 ^ ρ
      = (v / 2 ^ ρ) * 2 ^ ρ + v % 2 ^ ρ := by rw [h2]
    _ = v := by rw [Nat.mul_comm]; exact h

lemma setBit_lt {ρ K v : ℕ} (hρ : ρ < K) (hv : v < 2 ^ K) (b : Fin 2) : setBit ρ v b < 2 ^ K := by
  have hP : 0 < 2 ^ ρ := by positivity
  have hK : 2 ^ K = 2 ^ (ρ + 1) * 2 ^ (K - (ρ + 1)) := by rw [← pow_add]; congr 1; omega
  have hq : v / 2 ^ (ρ + 1) < 2 ^ (K - (ρ + 1)) := by
    rw [Nat.div_lt_iff_lt_mul (by positivity), Nat.mul_comm, ← hK]; exact hv
  have hr := Nat.mod_lt v hP
  rw [setBit_eq, hK, pow_succ]
  have hb := b.isLt
  have : (2 * (v / 2 ^ (ρ + 1)) + b) * 2 ^ ρ + v % 2 ^ ρ < (2 * (v / 2 ^ (ρ + 1)) + 2) * 2 ^ ρ := by
    nlinarith
  calc _ < (2 * (v / 2 ^ (ρ + 1)) + 2) * 2 ^ ρ := this
    _ = 2 ^ ρ * 2 * (v / 2 ^ (ρ + 1) + 1) := by ring
    _ ≤ 2 ^ ρ * 2 * 2 ^ (K - (ρ + 1)) := Nat.mul_le_mul_left _ hq

end BitArith

section Layout

variable (P K D S ρ : ℕ)

/-- Address normal form `p :: c₀ … c_{D-1} :: [s]`. -/
def mkAddr (p : ℕ) (c : Fin D → ℕ) (s : ℕ) : List ℕ := p :: (List.ofFn c ++ [s])

lemma getD_mkAddr_chunk (p : ℕ) (c : Fin D → ℕ) (s : ℕ) (i : Fin D) :
    (mkAddr D p c s).getD (i + 1) 0 = c i := by
  simp [mkAddr, List.getD_eq_getElem?_getD, List.getElem?_append_left (by simp : (i : ℕ) < (List.ofFn c).length)]

lemma getD_mkAddr_last (p : ℕ) (c : Fin D → ℕ) (s : ℕ) :
    (mkAddr D p c s).getD (D + 1) 0 = s := by
  simp [mkAddr, List.getD_eq_getElem?_getD, List.getElem?_append_right]

lemma headD_mkAddr (p : ℕ) (c : Fin D → ℕ) (s : ℕ) : (mkAddr D p c s).headD 0 = p := rfl

lemma eq_mkAddr {a : List ℕ} (h : a.length = D + 2) :
    a = mkAddr D (a.headD 0) (fun i => a.getD (i + 1) 0) (a.getD (D + 1) 0) := by
  apply List.ext_getElem (by simp [mkAddr, h])
  intro k h₁ h₂
  rcases k with _ | k
  · cases a with
    | nil => simp at h
    | cons x l => simp [mkAddr]
  · simp only [mkAddr, List.getElem_cons_succ]
    by_cases hk : k < D
    · rw [List.getElem_append_left (by simpa using hk)]
      simp [List.getD_eq_getElem?_getD, List.getElem?_eq_getElem h₁]
    · have hk' : k = D := by rw [h] at h₁; omega
      subst hk'
      rw [List.getElem_append_right (by simp)]
      simp [List.getD_eq_getElem?_getD, List.getElem?_eq_getElem h₁]

lemma valid_layer_iff {a : List ℕ} :
    Valid (layerShape P K D S) a ↔ a.length = D + 2 ∧ a.headD 0 < P ∧
      (∀ i : Fin D, a.getD (i + 1) 0 < 2 ^ K) ∧ a.getD (D + 1) 0 < S := by
  constructor
  · intro hv
    have hl : a.length = D + 2 := by rw [hv.length_eq]; simp [layerShape]
    rw [eq_mkAddr D hl] at hv
    refine ⟨hl, ?_, ?_, ?_⟩
    · have := (valid_iff _ _).1 hv
      have h0 := this.2 0 (by simp [mkAddr]) (by simp [layerShape])
      simpa [mkAddr, layerShape] using h0
    · intro i
      have := ((valid_iff _ _).1 hv).2 (i + 1) (by simp [mkAddr]) (by simp [layerShape])
      rw [← getD_mkAddr_chunk D (a.headD 0) (fun i => a.getD (i + 1) 0) (a.getD (D + 1) 0) i]
      rw [List.getD_eq_getElem _ _ (by simp [mkAddr])]
      simpa [layerShape, List.getElem_append_left, mkAddr] using this
    · have := ((valid_iff _ _).1 hv).2 (D + 1) (by simp [mkAddr]) (by simp [layerShape])
      rw [← getD_mkAddr_last D (a.headD 0) (fun i => a.getD (i + 1) 0) (a.getD (D + 1) 0)]
      rw [List.getD_eq_getElem _ _ (by simp [mkAddr])]
      simpa [layerShape, List.getElem_append_right, mkAddr] using this
  · rintro ⟨hl, hp, hc, hs⟩
    rw [valid_iff]
    refine ⟨by simp [hl, layerShape], ?_⟩
    intro k h₁ h₂
    rcases k with _ | k
    · cases a with
      | nil => simp at hl
      | cons x l => simpa [layerShape] using hp
    · simp only [layerShape, List.getElem_cons_succ]
      by_cases hk : k < D
      · rw [List.getElem_append_left (by simpa using hk)]
        have := hc ⟨k, hk⟩
        simpa [List.getD_eq_getElem?_getD, List.getElem?_eq_getElem h₁] using this
      · have hk' : k = D := by rw [hl] at h₁; omega
        subst hk'
        rw [List.getElem_append_right (by simp)]
        simpa [List.getD_eq_getElem?_getD, List.getElem?_eq_getElem h₁] using hs

/-- Selected bits of an address: bit `ρ` of each chunk. -/
def layerBits (a : List ℕ) (i : Fin D) : Fin 2 := getBit ρ (a.getD (i + 1) 0)

/-- Overwrite the selected bits. -/
def layerSet (a : List ℕ) (y : Fin D → Fin 2) : List ℕ :=
  mkAddr D (a.headD 0) (fun i => setBit ρ (a.getD (i + 1) 0) (y i)) (a.getD (D + 1) 0)

/-- The selected-bit structure of `[P] × [2^K]^D × [S]` with offset `ρ < K`. -/
def layerSel (hρ : ρ < K) : SelBits (layerShape P K D S) (Fin D) where
  bits := layerBits D ρ
  set := layerSet D ρ
  valid_set a y ha := by
    obtain ⟨hl, hp, hc, hs⟩ := (valid_layer_iff P K D S).1 ha
    refine (valid_layer_iff P K D S).2 ⟨by simp [layerSet, mkAddr], hp, fun i => ?_, ?_⟩
    · rw [layerSet, getD_mkAddr_chunk]; exact setBit_lt hρ (hc i) _
    · rw [layerSet, getD_mkAddr_last]; exact hs
  bits_set a y _ := by
    funext i; simp only [layerBits, layerSet, getD_mkAddr_chunk, getBit_setBit]
  set_set a y z _ := by
    simp only [layerSet, getD_mkAddr_chunk, getD_mkAddr_last, headD_mkAddr, setBit_setBit]
  set_bits a ha := by
    obtain ⟨hl, -⟩ := (valid_layer_iff P K D S).1 ha
    conv_rhs => rw [eq_mkAddr D hl]
    simp only [layerSet, layerBits, setBit_getBit]

/-- **The simultaneous normalized butterfly layer** of §5 on `[P] × [2^K]^D × [S]`:
`H0^{⊗D}` on bit `ρ` of each of the `D` chunks, coefficientwise, fixing all other bits. -/
noncomputable def layerH (hρ : ρ < K) (xs : List M) : List M :=
  butterflyLayer (layerSel P K D S ρ hρ) xs

end Layout

end Butterfly

end

section
-- Source module: Solutions.ButterflyGuard
/-!
# §5.4 An explicit guard-width bound (and the numeric parts of Proposition 14)

1. The constants `E = 64(W+m+1)³`, `B = s+E`, `ν = min{j ≥ 1 : m^j ≥ B}`, `C1 = max{20, ν+3}`,
   `C0 = 128 m B²` for the complex network `m = 10⁶`, `W_c`, `s_c`: `m^10 < B < m^11`, `ν = 11`,
   `C1 = 20`.
2. The depth recurrence `A(e) ≤ s A(e/m) + E`, `A ≤ 8e` at leaves, gives `A ≤ 8 d B^{k+1}
   ≤ 8 B² d^{ν+1}`, and the final chain `8mB²d^{ν+2}+8d+t+9 ≤ … < 128mB²d^{ν+3} ≤ C0 d^{C1}`.
3. Dyadic growth: a value at arithmetic depth `ℓ` (add, sub, halve, negate, multiply by `i`)
   lies in `2^{-(p+ℓ)} ℤ[i]` with modulus `≤ 2^ℓ`; hence `p+Δ` fractional and `Δ+3` integer and
   sign bits represent it exactly.
4. The final truncation `Q_p` toward zero: `|Q_p z| ≤ |z|`, `|Q_p z - z| < √2 · 2^{-p}`.
5. The width bookkeeping of the proof of Proposition 14.
-/

namespace Butterfly

open Complex

/-! ## 1. The constants -/

/-! ## 2. Depth recurrence and the chain -/

/-! ## 3. Dyadic depth invariant -/

/-! ## 4. Final truncation toward zero -/

/-- Truncation of a real toward zero to the grid `2^{-p} ℤ`, as an integer numerator. -/
noncomputable def truncNum (p : ℕ) (x : ℝ) : ℤ := if 0 ≤ x then ⌊x * 2 ^ p⌋ else ⌈x * 2 ^ p⌉

/-- `Q_p` on a real component. -/
noncomputable def truncR (p : ℕ) (x : ℝ) : ℝ := truncNum p x / 2 ^ p

/-- `Q_p` on a complex coefficient: truncate both components toward zero. -/
noncomputable def Qp (p : ℕ) (z : ℂ) : ℂ := ⟨truncR p z.re, truncR p z.im⟩

/-! ## 5. Width bookkeeping in the proof of Proposition 14 -/

end Butterfly

end