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

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

import Definitions.Def_IntMul_MultitapeModel
import Mathlib

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.AModel
/-!
# Array model: mixed-radix arrays, pull maps, and linear address maps via Bruhat factors

Arrays are lists in row-major lexicographic order over fields (first field most significant).
All denotations are plain `List` functions.
-/

namespace AModel

open Matrix

/-! ## 1. Mixed-radix basics -/

/-- A shape is the list of field sizes. -/
abbrev Shape := List ℕ

/-- Volume (number of entries) of a shape. -/
def vol (sh : Shape) : ℕ := sh.prod

/-- Mixed-radix decoding, first field most significant. -/
def decode : Shape → ℕ → List ℕ
  | [], _ => []
  | q :: sh, n => (n / vol sh) % q :: decode sh (n % vol sh)

/-- Mixed-radix encoding, first field most significant. -/
def encode : Shape → List ℕ → ℕ
  | [], _ => 0
  | _ :: sh, a => a.headD 0 * vol sh + encode sh a.tail

/-- A valid address: same length as the shape, each digit below its field size. -/
def Valid (sh : Shape) (a : List ℕ) : Prop := List.Forall₂ (· < ·) a sh

@[simp] lemma vol_nil : vol [] = 1 := rfl
@[simp] lemma vol_cons (q : ℕ) (sh : Shape) : vol (q :: sh) = q * vol sh := List.prod_cons

lemma valid_iff (sh : Shape) (a : List ℕ) :
    Valid sh a ↔ a.length = sh.length ∧
      ∀ (k : ℕ) (h₁ : k < a.length) (h₂ : k < sh.length), a[k] < sh[k] := by
  unfold Valid; rw [List.forall₂_iff_get]; rfl

lemma Valid.length_eq {sh : Shape} {a : List ℕ} (h : Valid sh a) : a.length = sh.length :=
  ((valid_iff sh a).1 h).1

lemma encode_lt {sh : Shape} {a : List ℕ} (h : Valid sh a) : encode sh a < vol sh := by
  induction h with
  | nil => simp [encode]
  | @cons x q as sh hx _ ih =>
    simp only [encode, List.headD_cons, List.tail_cons, vol_cons]
    have : (x + 1) * vol sh ≤ q * vol sh := Nat.mul_le_mul_right _ hx
    nlinarith

lemma valid_decode {sh : Shape} {m : ℕ} (h : m < vol sh) : Valid sh (decode sh m) := by
  induction sh generalizing m with
  | nil => exact List.Forall₂.nil
  | cons q sh ih =>
    simp only [vol_cons] at h
    have hV : 0 < vol sh := Nat.pos_of_ne_zero (by rintro h'; simp [h'] at h)
    have hq : 0 < q := Nat.pos_of_ne_zero (by rintro h'; simp [h'] at h)
    exact List.Forall₂.cons (Nat.mod_lt _ hq) (ih (Nat.mod_lt m hV))

/-! ## 2. Pull maps -/

/-- Array permutation by its pull map: output address `y` reads input address `g y`. -/
def pull {α : Type*} (sh sh' : Shape) (g : List ℕ → List ℕ) (d : α) (xs : List α) : List α :=
  (List.range (vol sh')).map (fun j => xs.getD (encode sh (g (decode sh' j))) d)

section Pull
variable {α : Type*}

end Pull

/-! ## 3. Active fields -/

/-- `n` active fields of size `Q` inside the shape `sh`, at strictly increasing positions. -/
structure Layout (sh : Shape) (n Q : ℕ) where
  pos : Fin n → Fin sh.length
  mono : StrictMono pos
  size : ∀ i, sh.get (pos i) = Q

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q)

/-- The active field (if any) at list index `k`. -/
def Layout.actIdx (k : ℕ) : Option (Fin n) :=
  (List.finRange n).find? (fun i => (Λ.pos i : ℕ) == k)

/-- Active vector of an address. -/
def act (a : List ℕ) (i : Fin n) : ZMod Q := ((a.getD (Λ.pos i) 0 : ℕ) : ZMod Q)

/-- Replace the active digits of `a` by `(v i).val`. -/
def setAct (a : List ℕ) (v : Fin n → ZMod Q) : List ℕ :=
  a.mapIdx fun k x => match Λ.actIdx k with
    | some i => (v i).val
    | none => x

/-- Address map induced by a map on active vectors. -/
def vmap (F : (Fin n → ZMod Q) → (Fin n → ZMod Q)) (a : List ℕ) : List ℕ :=
  setAct Λ a (F (act Λ a))

/-- Address map of a matrix acting on the active vector. -/
def amap (L : Matrix (Fin n) (Fin n) (ZMod Q)) (a : List ℕ) : List ℕ :=
  setAct Λ a (L *ᵥ act Λ a)

namespace Layout

end Layout

/-! ## 4. Controlled shift -/

/-- Active-vector form of the controlled shift. -/
def shiftVec (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q) (x : Fin n → ZMod Q) :
    Fin n → ZMod Q :=
  Function.update x i (u * x i + ∑ j, c j * x j)

/-- Inverse of `shiftVec` (when `u` is a unit and `c j = 0` for `j ≥ i`). -/
def shiftInvVec (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q) (y : Fin n → ZMod Q) :
    Fin n → ZMod Q :=
  Function.update y i (u⁻¹ * (y i - ∑ j, c j * y j))

/-- Controlled shift address map: `x_i ↦ u x_i + Σ_j c_j x_j`. -/
def shiftMap (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q) (a : List ℕ) : List ℕ :=
  setAct Λ a (Function.update (act Λ a) i (u * act Λ a i + ∑ j, c j * act Λ a j))

/-- Its explicit inverse: `y_i ↦ u⁻¹ (y_i - Σ_j c_j y_j)`. -/
def shiftInv (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q) (a : List ℕ) : List ℕ :=
  setAct Λ a (Function.update (act Λ a) i (u⁻¹ * (act Λ a i - ∑ j, c j * act Λ a j)))

/-! ## 5. Field transposition -/

/-- Exchange the digits at `pos i` and `pos j`. -/
def swapMap (i j : Fin n) (a : List ℕ) : List ℕ :=
  (a.set (Λ.pos i) (a.getD (Λ.pos j) 0)).set (Λ.pos j) (a.getD (Λ.pos i) 0)

section Shift

end Shift

section Swap

end Swap

/-! ## Primitive steps and their composition -/

/-- A primitive step on the active fields. -/
inductive Step (n Q : ℕ) where
  /-- Controlled shift (d): `x_i ↦ u x_i + Σ_j c_j x_j`. -/
  | shift (i : Fin n) (u : ZMod Q) (c : Fin n → ZMod Q)
  /-- Field transposition (e): exchange fields `i` and `j`. -/
  | swap (i j : Fin n)

namespace Step

/-- Well-formedness: shifts have unit coefficient and only read earlier fields. -/
def WF : Step n Q → Prop
  | shift i u c => IsUnit u ∧ ∀ j, i ≤ j → c j = 0
  | swap _ _ => True

/-- Action on active vectors. -/
def vec : Step n Q → (Fin n → ZMod Q) → (Fin n → ZMod Q)
  | shift i u c => shiftVec i u c
  | swap i j => fun x => x ∘ Equiv.swap i j

/-- Inverse action on active vectors. -/
def invVec : Step n Q → (Fin n → ZMod Q) → (Fin n → ZMod Q)
  | shift i u c => shiftInvVec i u c
  | swap i j => fun x => x ∘ Equiv.swap i j

/-- Address map of a step. -/
def addr (s : Step n Q) : List ℕ → List ℕ :=
  match s with
  | shift i u c => shiftMap Λ i u c
  | swap i j => swapMap Λ i j

/-- Explicit inverse address map of a step. -/
def inv (s : Step n Q) : List ℕ → List ℕ :=
  match s with
  | shift i u c => shiftInv Λ i u c
  | swap i j => swapMap Λ i j

/-- Array action of a step: `pull` with the inverse address map. -/
def arr {α : Type*} (d : α) (s : Step n Q) (xs : List α) : List α :=
  pull sh sh (s.inv Λ) d xs

variable [NeZero Q]

end Step

/-- Run a list of steps on an address (first step first). -/
def run (steps : List (Step n Q)) (a : List ℕ) : List ℕ :=
  steps.foldl (fun b s => s.addr Λ b) a

/-- Run a list of steps on an array (first step first). -/
def runArr {α : Type*} (d : α) (steps : List (Step n Q)) (xs : List α) : List α :=
  steps.foldl (fun ys s => s.arr Λ d ys) xs

/-- Run a list of steps on an active vector (first step first). -/
def runVec (steps : List (Step n Q)) (x : Fin n → ZMod Q) : Fin n → ZMod Q :=
  steps.foldl (fun y s => s.vec y) x

/-- Composition of the inverse address maps: `inv s₁ ∘ inv s₂ ∘ ⋯ ∘ inv sₖ`. -/
def invChain (steps : List (Step n Q)) : List ℕ → List ℕ :=
  steps.foldr (fun s g => s.inv Λ ∘ g) id

section Run
variable [NeZero Q]

end Run

/-! ## 6. Linear address maps through Bruhat factors -/

section Bruhat

/-- The controlled shift realizing row `i` of a lower triangular `E`:
`x_i ↦ E i i * x_i + Σ_{j < i} E i j * x_j`. -/
def shiftStep (E : Matrix (Fin n) (Fin n) (ZMod Q)) (i : Fin n) : Step n Q :=
  .shift i (E i i) (fun j => if j < i then E i j else 0)

/-- Shifts for `E`, in descending target order `n-1, …, 0`. -/
def shiftSteps (E : Matrix (Fin n) (Fin n) (ZMod Q)) : List (Step n Q) :=
  (List.finRange n).reverse.map (shiftStep E)

/-- Field transpositions, one per pair. -/
def swapSteps (ts : List (Fin n × Fin n)) : List (Step n Q) :=
  ts.map fun p => .swap p.1 p.2

/-- The full pipeline for `E₁ * w * E₂`: shifts for `E₂`, then swaps, then shifts for `E₁`. -/
def bruhatSteps (E₁ : Matrix (Fin n) (Fin n) (ZMod Q)) (ts : List (Fin n × Fin n))
    (E₂ : Matrix (Fin n) (Fin n) (ZMod Q)) : List (Step n Q) :=
  shiftSteps E₂ ++ swapSteps ts ++ shiftSteps E₁

variable {E : Matrix (Fin n) (Fin n) (ZMod Q)}

/-- The product of the transpositions of `ts`, leftmost first. -/
def swapProd (ts : List (Fin n × Fin n)) : Equiv.Perm (Fin n) :=
  (ts.map fun p => Equiv.swap p.1 p.2).prod

variable [NeZero Q]

end Bruhat

end AModel

end

section
-- Source module: Solutions.AModelCost
/-!
# Cost recurrences for the A-model (§4, Prop. 11 eq. (12) and Lemma 12)

* `power_width_recurrence`: `F (k+1) ≤ (s/W) F k + C` with `s/W < m^τ` gives
  `F k ≤ K (m^k)^τ`.
* `power_width_time`: the volume-scaled ("time") version
  `T (k+1) V ≤ C V + s T k (V/W)` gives `T k V ≤ K V (m^k)^τ`.
* `arbitrary_width_sum`: digit decomposition `Σ_j a_j (m^j)^τ ≤ K' e^τ`.
* `log_two_mul_le_rpow`: `log (2e) ≤ c e^τ`.
* `arbitrary_width_cost`: combination, total cost `≤ K'' V e^τ`.
-/

namespace AModel

open Finset

/-- The coefficient sequence `G 0 = C`, `G (k+1) = (s/W) G k + C`. -/
noncomputable def coeffSeq (s W : ℕ) (C : ℝ) : ℕ → ℝ
  | 0 => C
  | k + 1 => ((s : ℝ) / W) * coeffSeq s W C k + C

/-- The `j`-th base-`m` digit of `e` (zero beyond the length). -/
def digit (m e j : ℕ) : ℕ := (Nat.digits m e).getD j 0

end AModel

end

section
-- Source module: Solutions.AModelRuns
/-!
# Array model, Stage 2: runs of pivots as single interchanges of merged fields

Builds on `Solutions/AModel.lean` (whose definitions are unchanged).

* §1 Reshape (`mergeShape`, `mergeAddr`, `splitAddr`): merging adjacent fields is free —
  `pull_merge`, `pull_split`, general `pull_reshape`.
* §2 Block swap (`blockSwap`): it is the composition of the single position swaps
  (`chain_runPairs`), and after merging both runs it is one swap of two merged fields
  (`merge2_blockSwap`, `pull_blockSwap`).
* §3 A run of `r` consecutive pivots on contiguous active fields acts on arrays as that block swap
  (`runArr_pivotRun`), hence as one merged-field interchange (`runArr_pivotRun_merged`).
* §4 A pivot list splits into `numRuns` maximal runs (`groupRuns`, `flatMap_groupRuns`), and
  `runArr` over its swaps is the composition of one merged-field swap per run
  (`runArr_swapSteps_runs`).
-/

namespace AModel

/-! ## 1. Reshape -/

section Reshape

/-- Merge the fields `[i, i+r)` of `sh` into one field. -/
def mergeShape (sh : Shape) (i r : ℕ) : Shape :=
  sh.take i ++ [vol ((sh.drop i).take r)] ++ sh.drop (i + r)

/-- Merge the digits `[i, i+r)` of an address into one digit. -/
def mergeAddr (sh : Shape) (i r : ℕ) (a : List ℕ) : List ℕ :=
  a.take i ++ [encode ((sh.drop i).take r) ((a.drop i).take r)] ++ a.drop (i + r)

/-- Split digit `i` of a merged address back into `r` digits. -/
def splitAddr (sh : Shape) (i r : ℕ) (b : List ℕ) : List ℕ :=
  b.take i ++ decode ((sh.drop i).take r) (b.getD i 0) ++ b.drop (i + 1)

end Reshape

/-! ## 2. Block swaps -/

section BlockSwap

/-- Gather digits by an index map: digit `k` of the result is digit `f k` of `a`. -/
def permIdx (f : ℕ → ℕ) (a : List ℕ) : List ℕ :=
  (List.range a.length).map fun k => a.getD (f k) 0

/-- Exchange the digits at list positions `p` and `q` (same formula as `swapMap`). -/
def swapPos (p q : ℕ) (a : List ℕ) : List ℕ :=
  (a.set p (a.getD q 0)).set q (a.getD p 0)

/-- Index map of `swapPos`. -/
def swapIdx (p q k : ℕ) : ℕ := if k = q then p else if k = p then q else k

/-- Index map of `blockSwap`. -/
def bsIdx (i j r k : ℕ) : ℕ :=
  if i ≤ k ∧ k < i + r then k + j - i else if j ≤ k ∧ k < j + r then k + i - j else k

/-- Exchange the digit runs `[i, i+r)` and `[j, j+r)`. -/
def blockSwap (i j r : ℕ) (a : List ℕ) : List ℕ := permIdx (bsIdx i j r) a

/-! ### Segments of a block-swapped address -/

section Segments
variable {i j r : ℕ} {a : List ℕ} (hij : i + r ≤ j) (hj : j + r ≤ a.length)
include hij hj

end Segments

/-! ### Merging both runs -/

/-- Both runs merged (the later run first). -/
def merge2 (sh : Shape) (i j r : ℕ) (a : List ℕ) : List ℕ :=
  mergeAddr (mergeShape sh j r) i r (mergeAddr sh j r a)

/-- The shape with both runs merged. -/
def mergeShape2 (sh : Shape) (i j r : ℕ) : Shape := mergeShape (mergeShape sh j r) i r

/-- Inverse of `merge2`. -/
def split2 (sh : Shape) (i j r : ℕ) (b : List ℕ) : List ℕ :=
  splitAddr sh j r (splitAddr (mergeShape sh j r) i r b)

end BlockSwap

/-! ## 3. A run of consecutive pivots is one merged-field interchange -/

section PivotRun

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q)

variable [NeZero Q]

end PivotRun

/-! ## 4. Counting runs -/

section Runs

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) [NeZero Q]

end Runs

end AModel

end

section
-- Source module: Solutions.AModelShear
/-!
# Shears and interchanges in the array model (Lemma 10 / end of Prop. 11)

Active fields are indexed by `Fin (m + m)`: the `H` fields are `Fin.castAdd m i` (the first `m`),
the `D` fields are `Fin.natAdd m j` (the last `m`). With a `Layout` (strictly increasing
positions) every `H` field therefore precedes every `D` field.
-/

namespace AModel

open Matrix

section Blocks

variable {m Q : ℕ}

/-- `H` part of an active vector. -/
def hv {R : Type*} (x : Fin (m + m) → R) : Fin m → R := fun i => x (Fin.castAdd m i)

/-- `D` part of an active vector. -/
def dv {R : Type*} (x : Fin (m + m) → R) : Fin m → R := fun j => x (Fin.natAdd m j)

/-- A `2 × 2` block matrix over `Fin m ⊕ Fin m`, reindexed to `Fin (m + m)`. -/
def blk (P B C D : Matrix (Fin m) (Fin m) (ZMod Q)) : Matrix (Fin (m + m)) (Fin (m + m)) (ZMod Q) :=
  reindex finSumFinEquiv finSumFinEquiv (fromBlocks P B C D)

/-- The shear `(H, D) ↦ (H + A D, D)`. -/
def shearMat (A : Matrix (Fin m) (Fin m) (ZMod Q)) : Matrix (Fin (m + m)) (Fin (m + m)) (ZMod Q) :=
  AModel.blk 1 A 0 1

/-- The block interchange `(H, D) ↦ (D, H)`. -/
def swapMat : Matrix (Fin (m + m)) (Fin (m + m)) (ZMod Q) := AModel.blk 0 1 1 0

section AmapShear

variable {sh : Shape} {m Q : ℕ} (Λ : Layout sh (m + m) Q) [NeZero Q]

end AmapShear

end Blocks

/-! ## Inverse steps -/

section Rev

variable {n Q : ℕ}

/-- The explicit inverse of a step, as a step: a shift by `u⁻¹` with controls `-u⁻¹ c`. -/
def Step.rev : Step n Q → Step n Q
  | .shift i u c => .shift i u⁻¹ (fun j => -(u⁻¹ * c j))
  | .swap i j => .swap i j

/-- Is the step a field transposition? -/
def Step.isSwap : Step n Q → Bool
  | .shift _ _ _ => false
  | .swap _ _ => true

/-- Inverse of a step list: inverted steps in reverse order. -/
def revSteps (l : List (Step n Q)) : List (Step n Q) := (l.map Step.rev).reverse

end Rev

/-! ## Lower triangular transforms of one block -/

section Tri

variable {m Q : ℕ}

/-- `H` field indices, descending. -/
def hIdx (m : ℕ) : List (Fin (m + m)) := (List.finRange m).reverse.map (Fin.castAdd m)

/-- `D` field indices, descending. -/
def dIdx (m : ℕ) : List (Fin (m + m)) := (List.finRange m).reverse.map (Fin.natAdd m)

/-- Shifts realizing `H ← E H` (one per `H` field, descending targets). -/
def hSteps (E : Matrix (Fin m) (Fin m) (ZMod Q)) : List (Step (m + m) Q) :=
  (hIdx m).map (shiftStep (AModel.blk E 0 0 1))

/-- Shifts realizing `D ← E D` (one per `D` field, descending targets). -/
def dSteps (E : Matrix (Fin m) (Fin m) (ZMod Q)) : List (Step (m + m) Q) :=
  (dIdx m).map (shiftStep (AModel.blk 1 0 0 E))

end Tri

/-! ## The pivot gadget: `H_i ← H_i + D_j` with one interchange -/

section Pivot

variable {m Q : ℕ}

/-- `D_j ← D_j + H_i`, `(H_i, D_j) ← (D_j, H_i)`, `D_j ← H_i − D_j`. -/
def pivotSteps (i j : Fin m) : List (Step (m + m) Q) :=
  [.shift (Fin.natAdd m j) 1 (Pi.single (Fin.castAdd m i) 1),
   .swap (Fin.castAdd m i) (Fin.natAdd m j),
   .shift (Fin.natAdd m j) (-1) (Pi.single (Fin.castAdd m i) 1)]

/-- One gadget per pivot, in list order. -/
def addSteps (ps : List (Fin m × Fin m)) : List (Step (m + m) Q) :=
  ps.flatMap fun p => pivotSteps p.1 p.2

/-- The pivots (positions of the ones) of `W`, in lexicographic order. -/
def pivots (W : Matrix (Fin m) (Fin m) (ZMod Q)) : List (Fin m × Fin m) :=
  (List.finRange m ×ˢ List.finRange m).filter fun p => W p.1 p.2 = 1

end Pivot

/-! ## Lemma 10: the shear through Bruhat-type factors -/

section Shear

variable {m Q : ℕ}

/-- Steps for the shear by `A = E₁ W E₂`: `D ← E₂ D`, `H ← E₁⁻¹ H` (the inverse shifts of
`H ← E₁ H`), `H ← H + W D` (one pivot gadget per one of `W`), `H ← E₁ H`, `D ← E₂⁻¹ D`. -/
def shearSteps (E₁ W E₂ : Matrix (Fin m) (Fin m) (ZMod Q)) : List (Step (m + m) Q) :=
  dSteps E₂ ++ revSteps (hSteps E₁) ++ addSteps (pivots W) ++ hSteps E₁ ++ revSteps (dSteps E₂)

variable {E₁ W E₂ : Matrix (Fin m) (Fin m) (ZMod Q)}

variable {sh : Shape} (Λ : Layout sh (m + m) Q) [NeZero Q]

end Shear

/-! ## Interchange from a shear (end of Prop. 11) -/

section Interchange

variable {m Q : ℕ}

/-- The shift `D_j ← u D_j + c H_j`. -/
def dhStep (u c : ZMod Q) (j : Fin m) : Step (m + m) Q :=
  .shift (Fin.natAdd m j) u (Pi.single (Fin.castAdd m j) c)

/-- `D ← D − H`, one shift per `D` field. -/
def subSteps (m : ℕ) : List (Step (m + m) Q) := (List.finRange m).map (dhStep 1 (-1))

/-- `D ← H − D`, one shift per `D` field. -/
def flipSteps (m : ℕ) : List (Step (m + m) Q) := (List.finRange m).map (dhStep (-1) 1)

variable {sh : Shape} (Λ : Layout sh (m + m) Q) [NeZero Q]

end Interchange

end AModel

end

section
-- Source module: Solutions.AModelNet
/-!
# Row streams and the gate network (Stage 3 of the interchange procedure)

1. Cyclic row split / merge (primitive (b)) on shapes `P :: (W * L) :: rest`:
   row `r = W * g + w` of the big array becomes row `g` of stream `w`.
2. The abstract network theorem: a circuit of pointwise linear gates, executed physically with
   frame changes `Φ (M_g - label)` on every wire, realizes `Φ I0` on every wire, routed by `ρ`.
3. The combination: split, run the network with `Φ_M = pull` by `π_M` on the `rest` digits, merge
   back with `ρ` undone; this equals the `pull` of the whole array by `π_{I0}` on the `rest` digits.
-/

namespace AModel

open Matrix

/-! ## 1. Cyclic row split and merge -/

section RowSplit

variable {α : Type*}

/-- Stream address `(p, g, x)` ↦ big address `(p, W g + w, x)`. -/
def rowSplitAddr (W w : ℕ) (a : List ℕ) : List ℕ :=
  a.headD 0 :: (W * a.tail.headD 0 + w) :: a.tail.tail

/-- Big address `(p, r, x)` ↦ stream address `(p, r / W, x)`. -/
def rowMergeAddr (W : ℕ) (a : List ℕ) : List ℕ :=
  a.headD 0 :: a.tail.headD 0 / W :: a.tail.tail

/-- The stream (role) of a big address `(p, r, x)`: `r % W`. -/
def rowIdx (W : ℕ) [NeZero W] (a : List ℕ) : Fin W :=
  ⟨a.tail.headD 0 % W, Nat.mod_lt _ (NeZero.pos W)⟩

/-- Lift an address map on the `rest` digits to `(p, r, x) ↦ (p, r, g x)`. -/
def rowLift (g : List ℕ → List ℕ) (a : List ℕ) : List ℕ :=
  a.headD 0 :: a.tail.headD 0 :: g a.tail.tail

/-- Stream `w` of the cyclic row split: rows `r` with `r % W = w`, shape `P :: L :: rest`. -/
def rowSplit (P W L : ℕ) (rest : Shape) (w : Fin W) (d : α) (xs : List α) : List α :=
  pull (P :: W * L :: rest) (P :: L :: rest) (rowSplitAddr W w) d xs

/-- Merge `W` streams of shape `P :: L :: rest` back into shape `P :: (W * L) :: rest`. -/
def rowMerge (P W L : ℕ) [NeZero W] (rest : Shape) (d : α) (s : Fin W → List α) : List α :=
  (List.range (vol (P :: W * L :: rest))).map fun j =>
    (s (rowIdx W (decode (P :: W * L :: rest) j))).getD
      (encode (P :: L :: rest) (rowMergeAddr W (decode (P :: W * L :: rest) j))) d

lemma valid_cons2 {P Q : ℕ} {sh : Shape} {a : List ℕ} (h : Valid (P :: Q :: sh) a) :
    ∃ p q x, a = p :: q :: x ∧ p < P ∧ q < Q ∧ Valid sh x := by
  unfold Valid at h
  rcases h with _ | ⟨hp, h'⟩
  rcases h' with _ | ⟨hq, hx⟩
  exact ⟨_, _, _, rfl, hp, hq, hx⟩

lemma valid_cons2_mk {P Q : ℕ} {sh : Shape} {p q : ℕ} {x : List ℕ} (hp : p < P) (hq : q < Q)
    (hx : Valid sh x) : Valid (P :: Q :: sh) (p :: q :: x) :=
  List.Forall₂.cons hp (List.Forall₂.cons hq hx)

lemma valid_rowLift {P L : ℕ} {rest : Shape} {g : List ℕ → List ℕ}
    (hg : ∀ x, Valid rest x → Valid rest (g x)) {a : List ℕ}
    (ha : Valid (P :: L :: rest) a) : Valid (P :: L :: rest) (rowLift g a) := by
  obtain ⟨p, q, x, rfl, hp, hq, hx⟩ := valid_cons2 ha
  exact valid_cons2_mk hp hq (hg x hx)

end RowSplit

/-! ## 2. The abstract network theorem -/

section Network

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

/-- A gate: matrix over `R` on the wires, the set of touched wires, and the gate's label. -/
abbrev NGate (Wr R M : Type*) := Matrix Wr Wr R × Finset Wr × M

/-- A gate acts as the identity off its touched set and touched rows only read touched wires. -/
def NGate.OK (g : NGate Wr R M) : Prop :=
  ∀ w w', (w ∉ g.2.1 → g.1 w w' = if w = w' then 1 else 0) ∧
    (w ∈ g.2.1 → w' ∉ g.2.1 → g.1 w w' = 0)

/-- Physical execution of one gate: on each touched wire apply the frame change
`Φ (M_g - current label)`, then the gate pointwise; touched wires get label `M_g`. -/
def physStep (Φ : M → (Ω → R) →ₗ[R] (Ω → R)) (st : (Wr → M) × (Wr → Ω → R))
    (g : NGate Wr R M) : (Wr → M) × (Wr → Ω → R) :=
  (fun w => if w ∈ g.2.1 then g.2.2 else st.1 w,
   fun w => if w ∈ g.2.1 then ∑ w' ∈ g.2.1, g.1 w w' • Φ (g.2.2 - st.1 w') (st.2 w')
     else st.2 w)

/-- The physical network: start with labels `Min`, run all gates, finish with `Φ (Mout - label)`. -/
def netRun (Φ : M → (Ω → R) →ₗ[R] (Ω → R)) (gates : List (NGate Wr R M)) (Min Mout : Wr → M)
    (f : Wr → Ω → R) : Wr → Ω → R :=
  fun w => Φ (Mout w - (gates.foldl (physStep Φ) (Min, f)).1 w)
    ((gates.foldl (physStep Φ) (Min, f)).2 w)

/-- Logical (frame-free) gate on streams. -/
def logStep (h : Wr → Ω → R) (g : NGate Wr R M) : Wr → Ω → R :=
  fun w => ∑ w', g.1 w w' • h w'

/-- The scalar circuit. -/
def scalarRun (gates : List (NGate Wr R M)) (v : Wr → R) : Wr → R :=
  gates.foldl (fun v g => g.1 *ᵥ v) v

lemma physStep_inv (Φ : M → (Ω → R) →ₗ[R] (Ω → R))
    (hadd : ∀ a b x, Φ (a + b) x = Φ a (Φ b x))
    (st : (Wr → M) × (Wr → Ω → R)) (h : Wr → Ω → R) (hst : ∀ w, st.2 w = Φ (st.1 w) (h w))
    (g : NGate Wr R M) (hg : g.OK) :
    ∀ w, (physStep Φ st g).2 w = Φ ((physStep Φ st g).1 w) (logStep h g w) := by
  intro w
  simp only [physStep, logStep]
  by_cases hw : w ∈ g.2.1
  · simp only [hw, if_true, map_sum, map_smul]
    rw [← Finset.sum_subset (Finset.subset_univ g.2.1)]
    · refine Finset.sum_congr rfl fun w' _ => ?_
      rw [hst w', ← hadd, sub_add_cancel]
    · intro w' _ hw'
      simp [(hg w w').2 hw hw']
  · simp only [hw, if_false]
    rw [Finset.sum_eq_single w, hst w]
    · simp [(hg w w).1 hw]
    · intro b _ hb; simp [(hg w b).1 hw, Ne.symm hb]
    · simp

lemma physRun_inv (Φ : M → (Ω → R) →ₗ[R] (Ω → R))
    (hadd : ∀ a b x, Φ (a + b) x = Φ a (Φ b x))
    (gates : List (NGate Wr R M)) (hg : ∀ g ∈ gates, g.OK) :
    ∀ (st : (Wr → M) × (Wr → Ω → R)) (h : Wr → Ω → R), (∀ w, st.2 w = Φ (st.1 w) (h w)) →
      ∀ w, (gates.foldl (physStep Φ) st).2 w =
        Φ ((gates.foldl (physStep Φ) st).1 w) (gates.foldl logStep h w) := by
  induction gates with
  | nil => intro st h hst; simpa using hst
  | cons g gs ih =>
    intro st h hst
    simp only [List.foldl_cons]
    exact ih (fun g' hg' => hg g' (List.mem_cons_of_mem _ hg')) _ _
      (physStep_inv Φ hadd st h hst g (hg g List.mem_cons_self))

omit [DecidableEq Wr] [AddGroup M] in
lemma logRun_apply (gates : List (NGate Wr R M)) :
    ∀ (h : Wr → Ω → R) (w : Wr) (ω : Ω),
      gates.foldl logStep h w ω = scalarRun gates (fun w => h w ω) w := by
  induction gates with
  | nil => intro h w ω; rfl
  | cons g gs ih =>
    intro h w ω
    simp only [List.foldl_cons, scalarRun] at ih ⊢
    rw [ih]
    have : (fun w => logStep h g w ω) = g.1 *ᵥ fun w => h w ω := by
      funext w'
      simp [logStep, Matrix.mulVec, dotProduct, Finset.sum_apply]
    rw [this]

end Network

/-! ## 3. The full procedure: split, network with pull frames, merge with `ρ` undone -/

section Combine

/-- The address bijection candidate on `Fin (vol sh)` induced by a valid-preserving address map. -/
def addrPerm (sh : Shape) (G : List ℕ → List ℕ) (hG : ∀ a, Valid sh a → Valid sh (G a)) :
    Fin (vol sh) → Fin (vol sh) :=
  fun j => ⟨encode sh (G (decode sh j)), encode_lt (hG _ (valid_decode j.isLt))⟩

/-- The stream frame action: `Φ_m` = pull by `π m` on the `rest` digits, as a linear map on
streams `Fin (vol (P :: L :: rest)) → R`. -/
def streamΦ {R M : Type*} [CommRing R] (P L : ℕ) (rest : Shape) (π : M → List ℕ → List ℕ)
    (hπv : ∀ m x, Valid rest x → Valid rest (π m x)) (m : M) :
    (Fin (vol (P :: L :: rest)) → R) →ₗ[R] (Fin (vol (P :: L :: rest)) → R) :=
  LinearMap.funLeft R R (addrPerm (P :: L :: rest) (rowLift (π m))
    (fun _ ha => valid_rowLift (hπv m) ha))

end Combine

end AModel

end

section
-- Source module: Solutions.AModelProp11
/-!
# Proposition 11 (power-width interchange with rows) in the array model

1. `Layout.cons2`: a layout on `rest` lifted to `P :: L :: rest`; `amap_cons2`.
2. The concrete frame action `shearπ` (shear address maps on `rest`), its hypotheses, and
   `shear_network`: the row-split network realizes `H ← H + D` on every row.
3. `netRunF`: the network with implemented frame changes; `netRunF_congr` (agreement on the frame
   changes that actually occur, `FramesOK`, suffices).
4. `runC`: Lemma 10 steps with every field interchange delegated to a child.
5. Chunk shapes `hdShape` (chunks split into `m` fields) / `swShape` (chunks merged).
6. `amap_swapMat_eq_blockSwap`.
7. One level (`levelR`) with an oracle child: `levelR_correct_amap`, `levelR_correct`.
8. The recursive program `prog` and `prog_correct` (induction on the level).
9. Cost: `levelCost_eq` (`T_{k+1}(V) = C₀ V + s T_k(V/W)`), `prog_cost`.
10. `prop11`: denotation (block swap of the two `m^k`-digit chunks) and cost `O(V (m^k)^τ)`.
-/

namespace AModel

open Matrix

/-! ## 1. Lifting a layout past the prefix and row fields -/

section Cons2

variable {sh : Shape} {n Q : ℕ}

/-- A layout on `sh` viewed on `P :: L :: sh` (positions shifted by two). -/
def Layout.cons2 (Λ : Layout sh n Q) (P L : ℕ) : Layout (P :: L :: sh) n Q where
  pos i := ⟨Λ.pos i + 2, by simp⟩
  mono a b h := by
    have := Λ.mono h
    rw [Fin.lt_def] at this ⊢
    simp only
    omega
  size i := by simpa using Λ.size i

variable (Λ : Layout sh n Q)

@[simp] lemma Layout.cons2_pos (P L : ℕ) (i : Fin n) :
    ((Λ.cons2 P L).pos i : ℕ) = Λ.pos i + 2 := rfl

end Cons2

/-! ## 2. The concrete frame action: shears on the `rest` digits -/

section ShearAction

variable {rest : Shape} {m Q : ℕ} {M : Type*} [AddCommGroup M]

/-- The frame action of a label `a : M` (through an additive map `φ` to `m × m` matrices over
`ZMod Q`): the address map of the shear by `-φ a` on the `rest` digits. Pull by it is the array
permutation of the address map `(H, D) ↦ (H + φ a · D, D)`. -/
def shearπ (Λ : Layout rest (m + m) Q) (φ : M →+ Matrix (Fin m) (Fin m) (ZMod Q)) (a : M) :
    List ℕ → List ℕ :=
  amap Λ (shearMat (-(φ a)))

variable (Λ : Layout rest (m + m) Q) (φ : M →+ Matrix (Fin m) (Fin m) (ZMod Q))

end ShearAction

/-! ## 3. Networks with implemented frame changes -/

section Frames

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

/-- `physStep` with a plain function family (an implementation of the frame changes). -/
def physStepF (Φ : M → (Ω → R) → (Ω → R)) (st : (Wr → M) × (Wr → Ω → R))
    (g : NGate Wr R M) : (Wr → M) × (Wr → Ω → R) :=
  (fun w => if w ∈ g.2.1 then g.2.2 else st.1 w,
   fun w => if w ∈ g.2.1 then ∑ w' ∈ g.2.1, g.1 w w' • Φ (g.2.2 - st.1 w') (st.2 w')
     else st.2 w)

/-- `netRun` with a plain function family. -/
def netRunF (Φ : M → (Ω → R) → (Ω → R)) (gates : List (NGate Wr R M)) (Min Mout : Wr → M)
    (f : Wr → Ω → R) : Wr → Ω → R :=
  fun w => Φ (Mout w - (gates.foldl (physStepF Φ) (Min, f)).1 w)
    ((gates.foldl (physStepF Φ) (Min, f)).2 w)

/-- Labels after a gate. -/
def labStep (lab : Wr → M) (g : NGate Wr R M) : Wr → M :=
  fun w => if w ∈ g.2.1 then g.2.2 else lab w

/-- Every frame change of the run (each touched wire of each gate, and the final frames)
satisfies `Pr`. -/
def FramesOK (Pr : 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)) ∧ FramesOK Pr gs (labStep lab g) Mout

end Frames

/-! ## 4. Running steps with the interchanges delegated to a child procedure -/

section RunChild

variable {sh : Shape} {n Q : ℕ} (Λ : Layout sh n Q) {α : Type*}

/-- One step, with field interchanges delegated to `child`. -/
def stepC (child : Fin n → Fin n → List α → List α) (d : α) (ys : List α) :
    Step n Q → List α
  | .swap a b => child a b ys
  | .shift i u c => (Step.shift i u c).arr Λ d ys

/-- Run steps, delegating every field interchange to `child`. -/
def runC (child : Fin n → Fin n → List α → List α) (d : α) (steps : List (Step n Q))
    (xs : List α) : List α :=
  steps.foldl (stepC Λ child d) xs

end RunChild

/-! ## 5. Chunk shapes: two chunks split into `m` fields, and merged -/

section Chunks

/-- `rest` with two chunks, each split into `m` fields of size `N`:
`A ++ [N]*m ++ B ++ [N]*m ++ C`. -/
def hdShape (A B C : Shape) (m N : ℕ) : Shape :=
  A ++ (List.replicate m N ++ (B ++ (List.replicate m N ++ C)))

/-- `rest` with two single fields of size `F`: `A ++ [F] ++ B ++ [F] ++ C`. -/
def swShape (A B C : Shape) (F : ℕ) : Shape := A ++ (F :: (B ++ (F :: C)))

/-- Position of active field `t` (`H_t` for `t < m`, `D_{t-m}` otherwise) in `hdShape`. -/
def hdPos (A B : Shape) (m t : ℕ) : ℕ := if t < m then A.length + t else A.length + B.length + t

lemma length_hdShape (A B C : Shape) (m N : ℕ) :
    (hdShape A B C m N).length = A.length + m + B.length + m + C.length := by
  simp [hdShape]; omega

lemma hdShape_getD (A B C : Shape) (m N t : ℕ) (ht : t < m + m) :
    (hdShape A B C m N).getD (hdPos A B m t) 0 = N := by
  unfold hdShape hdPos
  split_ifs with h
  · rw [List.getD_append_right _ _ _ _ (by omega), Nat.add_sub_cancel_left,
      List.getD_append _ _ _ _ (by simp; omega)]
    simp [h]
  · rw [List.getD_append_right _ _ _ _ (by omega),
      show A.length + B.length + t - A.length = m + (B.length + (t - m)) by omega,
      List.getD_append_right _ _ _ _ (by simp), List.length_replicate,
      Nat.add_sub_cancel_left, List.getD_append_right _ _ _ _ (by omega),
      Nat.add_sub_cancel_left, List.getD_append _ _ _ _ (by simp; omega),
      List.getD_eq_getElem _ _ (by simp; omega), List.getElem_replicate]

lemma hdPos_lt (A B C : Shape) (m N t : ℕ) (ht : t < m + m) :
    hdPos A B m t < (hdShape A B C m N).length := by
  rw [length_hdShape]; unfold hdPos; split_ifs <;> omega

/-- The layout of the `2m` active fields of `hdShape`. -/
def hdLayout (A B C : Shape) (m N : ℕ) : Layout (hdShape A B C m N) (m + m) N where
  pos t := ⟨hdPos A B m t, hdPos_lt A B C m N t t.2⟩
  mono a b h := by
    rw [Fin.lt_def] at h ⊢
    simp only [hdPos]
    split_ifs <;> omega
  size t := by
    have := hdShape_getD A B C m N t t.2
    rw [List.getD_eq_getElem _ _ (hdPos_lt A B C m N t t.2)] at this
    simpa using this

@[simp] lemma hdLayout_pos (A B C : Shape) (m N : ℕ) (t : Fin (m + m)) :
    ((hdLayout A B C m N).pos t : ℕ) = hdPos A B m t := rfl

end Chunks

/-! ## 6. The chunk interchange `amap swapMat` is a block swap -/

section SwapBlock

variable {sh : Shape} {m Q : ℕ} [NeZero Q] (Λ : Layout sh (m + m) Q)

end SwapBlock

/-! ## 7. One level of Prop. 11, with the children as an oracle -/

section Level

/-- Integer `m × m` matrices: the labels of the circuit. -/
abbrev IMat (m : ℕ) := Matrix (Fin m) (Fin m) ℤ

/-- Reduction of an integer matrix mod `N`. -/
def castM {m : ℕ} (N : ℕ) (X : IMat m) : Matrix (Fin m) (Fin m) (ZMod N) :=
  X.map (Int.castRingHom (ZMod N))

/-- Reduction mod `N` as an additive map (the `φ` of `shear_network`). -/
def castHom (m N : ℕ) : IMat m →+ Matrix (Fin m) (Fin m) (ZMod N) :=
  (Int.castRingHom (ZMod N)).mapMatrix.toAddMonoidHom

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

/-- The data of the circuit: gates on `W` wires with integer matrix labels, input/output labels,
the routing `ρ`, and for every label difference `A` a factorization `A = E₁ Π E₂` (Lemma 10). -/
structure Circuit (R : Type*) [CommRing R] (m W : ℕ) where
  gates : List (NGate (Fin W) R (IMat m))
  Min : Fin W → IMat m
  Mout : Fin W → IMat m
  ρ : Fin W → Fin W
  fac : IMat m → IMat m × IMat m × IMat m

/-- `fac A = (E₁, Π, E₂)` is a Lemma 10 factorization of `A` usable mod `N`: `E₁, E₂` lower
triangular with diagonal units mod `N`, `Π` a `0/1` matrix, `A = E₁ Π E₂` over `ℤ`. -/
def FacAt (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 = (fac A).1 * (fac A).2.1 * (fac A).2.2

/-- The Lemma 10 steps of the frame change `Φ_A` (fields of size `N`). -/
def edgeSteps (C : Circuit R m W) (N : ℕ) (A : IMat m) : List (Step (m + m) N) :=
  shearSteps (castM N (C.fac A).1) (castM N (C.fac A).2.1) (castM N (C.fac A).2.2)

/-- The circuit hypotheses of Prop. 11 at field size `N`. -/
structure Circuit.OK (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 (FacAt N C.fac) C.gates C.Min C.Mout

/-- The implemented frame change `Φ_A` on a stream: the Lemma 10 steps of `A`, every field
interchange delegated to `child`. -/
def implΦ (C : Circuit R m W) (N P L' : ℕ) (A B Cs : Shape)
    (child : Fin (m + m) → Fin (m + m) → List R → List R) (X : IMat m)
    (f : Fin (vol (P :: L' :: hdShape A B Cs m N)) → R) :
    Fin (vol (P :: L' :: hdShape A B Cs m N)) → R :=
  fun i => (runC ((hdLayout A B Cs m N).cons2 P L') child 0 (AModel.edgeSteps C N X) (List.ofFn f)).getD
    (i : ℕ) 0

/-- **One level of the procedure** on `P :: (W * L') :: hdShape A B Cs m N` (two chunks of `m`
fields of size `N`): `D ← D − H`; row split into `W` streams; the network (gates pointwise,
frame changes `Φ_A` by `implΦ`); merge undoing `ρ`; `D ← H − D`. -/
def levelR [NeZero W] (C : Circuit R m W) (N P L' : ℕ) (A B Cs : Shape)
    (child : Fin (m + m) → Fin (m + m) → List R → List R) (xs : List R) : List R :=
  AModel.runArr ((hdLayout A B Cs m N).cons2 P (W * L')) 0 (flipSteps m)
    (rowMerge P W L' (hdShape A B Cs m N) 0 (fun w => List.ofFn
      (netRunF (implΦ C N P L' A B Cs child) C.gates C.Min C.Mout
        (fun w' i => (rowSplit P W L' (hdShape A B Cs m N) w' 0
          (AModel.runArr ((hdLayout A B Cs m N).cons2 P (W * L')) 0 (subSteps m) xs)).getD (i : ℕ) 0)
        (C.ρ w))))

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

/-- The oracle hypothesis on the children: `child H_i D_j` interchanges the two size-`N` fields
`H_i`, `D_j` of a stream (single fields, i.e. a Prop. 11 instance one level down). -/
def ChildOK (child : Fin (m + m) → Fin (m + m) → List R → List R) : Prop :=
  ∀ (i j : Fin m) (ys : List R), ys.length = vol (P :: L' :: hdShape A B Cs m N) →
    child (Fin.castAdd m i) (Fin.natAdd m j) ys =
      pull (P :: L' :: hdShape A B Cs m N) (P :: L' :: hdShape A B Cs m N)
        (swapPos (A.length + i + 2) (A.length + B.length + m + j + 2)) 0 ys

end Level

/-! ## 8. The recursive procedure and Prop. 11 by induction on the level -/

section Recursion

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

/-- **The interchange program.** `prog C q k P L A B Cs` interchanges the two fields of size
`q ^ m ^ k` (chunks of `m ^ k` radix-`q` digits) of `P :: L :: swShape A B Cs (q ^ m ^ k)`.
Level `0`: a single digit swap (primitive). Level `k + 1`: view each chunk as `m` fields of size
`N = q ^ m ^ k` (reshape, free) and run `levelR` with `L / W` stream rows, every field
interchange `H_a ↔ D_b` of every edge's Lemma 10 steps being a level-`k` call on the stream. -/
def prog (C : Circuit R m W) (q : ℕ) : ℕ → ℕ → ℕ → Shape → Shape → Shape → List R → List R
  | 0, P, L, A, B, Cs, xs =>
    pull (P :: L :: swShape A B Cs q) (P :: L :: swShape A B Cs q)
      (swapPos (A.length + 2) (A.length + B.length + 3)) 0 xs
  | k + 1, P, L, A, B, Cs, xs =>
    levelR C (q ^ m ^ k) P (L / W) A B Cs
      (fun a b => prog C q k P (L / W) (A ++ List.replicate a (q ^ m ^ k))
        (List.replicate (m - 1 - a) (q ^ m ^ k) ++ (B ++ List.replicate (b - m) (q ^ m ^ k)))
        (List.replicate (m + m - 1 - b) (q ^ m ^ k) ++ Cs)) xs

end Recursion

/-! ## 9. Cost -/

section CostDefs

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

/-- Sum of `f` over all frame changes of the run (as in `FramesOK`). -/
def frameSum {β : Type*} [AddCommMonoid β] (f : 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)) + frameSum f gs (labStep lab g) Mout

/-- Total number of (gate, touched wire) pairs: the pointwise gate passes on streams. -/
def touched (gates : List (NGate Wr R M)) : ℕ := (gates.map fun g => g.2.1.card).sum

/-- Cost of the network on streams of volume `v`: each frame change costs `fc` of it, each
touched wire of each gate one pass `v` for the pointwise gate. -/
noncomputable def netCost (fc : M → ℝ) (v : ℝ) : List (NGate Wr R M) → (Wr → M) → (Wr → M) → ℝ
  | [], lab, Mout => ∑ w, fc (Mout w - lab w)
  | g :: gs, lab, Mout => (∑ w ∈ g.2.1, (fc (g.2.2 - lab w) + v)) + netCost fc v gs (labStep lab g) Mout

end CostDefs

section StepCost

variable {n Q : ℕ}

/-- Cost of running steps on an array of volume `v`: a shift is one pass (`v`), a field
interchange is a child call (`T v`). -/
noncomputable def runCost (T : ℝ → ℝ) (v : ℝ) (steps : List (Step n Q)) : ℝ :=
  (steps.map fun s => if s.isSwap then T v else v).sum

variable {m : ℕ}

end StepCost

section ProgCost

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

/-- `r_A`: the number of ones of the `Π` of the factorization of `A` (its rank). -/
def pivRank (C : Circuit R m W) (A : IMat m) : ℕ :=
  (Finset.univ.filter fun p : Fin m × Fin m => (C.fac A).2.1 p.1 p.2 = 1).card

/-- `s = Σ_e r_e`, summed over all frame changes of the network: the number of child calls. -/
def rankSum (C : Circuit R m W) : ℕ := frameSum (pivRank C) C.gates C.Min C.Mout

/-- The constant `C₀` of the linear part: `2m + 2` full passes (`D ← D − H`, split, merge,
`D ← H − D`) plus, per stream (volume `V / W`), the shift passes `4m + 2 r_e` of each frame
change and one pass per touched wire of each gate. -/
noncomputable def costConst (C : Circuit R m W) : ℝ :=
  2 * m + 2 + (frameSum (fun A => (4 * m + 2 * pivRank C A : ℝ)) C.gates C.Min C.Mout +
    touched C.gates) / W

/-- Cost of one level with child cost `T`, on input volume `V` (per the structure of `levelR`). -/
noncomputable def levelCost (C : Circuit R m W) (N : ℕ) (T : ℝ → ℝ) (V : ℝ) : ℝ :=
  ((subSteps (Q := N) m).length : ℝ) * V + V +
    netCost (fun X => runCost T (V / W) (AModel.edgeSteps C N X)) (V / W) C.gates C.Min C.Mout + V +
    ((flipSteps (Q := N) m).length : ℝ) * V

/-- Cost of `prog` at level `k` on volume `V` (level `0`: one digit-swap pass). -/
noncomputable def progCost (C : Circuit R m W) (q : ℕ) : ℕ → ℝ → ℝ
  | 0, V => V
  | k + 1, V => levelCost C (q ^ m ^ k) (progCost C q k) V

end ProgCost

/-! ## 10. Proposition 11, assembled -/

section Assembled

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

end Assembled

end AModel

end