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

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

import Definitions.Def_IntMul_Kappa_K_Defs2

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.ButterflyProp14
/-!
# §5.1/§5.3/Proposition 14, conditional on the Proposition 8 interface

The phase-frame network of Proposition 8 is taken as an explicit hypothesis `Prop8Iface W m s`:
for every column index type `κ`, a list of network steps acting on `W` roles, each role
holding a complex function of the `m·|κ|` selected bits of one group (`Fin m × κ`: slot `h`,
column `j`). A step is either a non-recursive linear step (binary row additions `P_M`, scalar
gates, source/sink/bank corrections — any matrix) or a child call on slot `h` of role `w`,
forward (`(C)^{⊗κ}`) or inverse (implemented by (17) as `(-i)^{|κ|} Z C^{⊗κ} Z`, one forward call).
The interface asserts: exactly `s` child steps, and with exact children the composite is
`C^{⊗(Fin m × κ)}` on every role (the "exact recursive contract" of §5.1).

From it we prove:
* `recOp_eq`: the recursive program (leaves = individual kernels (25)) computes `C^{⊗}` on
  `u·m^k` selected bits at every depth `k`, on every role including scratch roles.
* `recCost_le`: its cost obeys the §5.3 recurrence and is `≤ K₁ V ((u m^k)^λ + (u m^k)^σ u^{1-σ})`.
* `tens_blocks`, `embed_tens`: pieces on disjoint blocks of selected bits multiply to `C^{⊗D}`.
* `prop14_output`: the layer procedure `Q_p(b^D S_D C^{⊗D} S_D z)` equals `Q_p(H0^{⊗D} z)`, stays
  in the unit disk and is within `√2·2^{-p}` of `H0^{⊗D} z` in every coefficient.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace Butterfly

open Matrix Finset Complex AModel

/-- Operators on the selected bits `κ`. -/
abbrev Op (κ : Type) := Matrix (κ → Fin 2) (κ → Fin 2) ℂ

/-- Operators on `W` roles of functions of the selected bits `ι`. -/
abbrev NOp (W : ℕ) (ι : Type) := Matrix (Fin W × (ι → Fin 2)) (Fin W × (ι → Fin 2)) ℂ

section Network

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

/-- Bits of slot `h` of a group address. -/
def slotBits (x : Fin m × κ → Fin 2) (h : Fin m) : κ → Fin 2 := fun j => x (h, j)

/-- A child call: `T` on slot `h` of role `w`; identity on all other slots and roles. -/
def childMat (w : Fin W) (h : Fin m) (T : Op κ) : NOp W (Fin m × κ) := fun p q =>
  if p.1 = q.1 then
    (if p.1 = w then
      (if ∀ h', h' ≠ h → ∀ j, p.2 (h', j) = q.2 (h', j) then T (slotBits p.2 h) (slotBits q.2 h)
        else 0)
     else if p.2 = q.2 then 1 else 0)
  else 0

/-- The same operator on every role. -/
def roleDiag {ι : Type} (T : Op ι) : NOp W ι := fun p q => if p.1 = q.1 then T p.2 q.2 else 0

/-- One role's block. -/
def roleBlock {ι : Type} (N : NOp W ι) (w : Fin W) : Op ι := fun x y => N (w, x) (w, y)

/-- A step of one network invocation on a group of `m` slots of columns `κ`. -/
inductive NStep (W m : ℕ) (κ : Type) where
  | loc (A : NOp W (Fin m × κ))
  | child (w : Fin W) (h : Fin m) (fwd : Bool)

/-- Whether a step is a (recursive) child call. -/
def NStep.isChild : NStep W m κ → Bool
  | .loc _ => false
  | .child .. => true

/-- Denotation of a step, given the operator `T` computed by a forward child call.
An inverse child is the forward call wrapped by (17). -/
noncomputable def NStep.den (T : Op κ) : NStep W m κ → NOp W (Fin m × κ)
  | .loc A => A
  | .child w h true => childMat w h T
  | .child w h false => childMat w h ((-I) ^ Fintype.card κ • (ZD * T * ZD))

/-- Denotation of a step list (first step applied first). -/
noncomputable def denList (T : Op κ) (l : List (NStep W m κ)) : NOp W (Fin m × κ) :=
  l.foldl (fun acc st => st.den T * acc) 1

end Network

/-- **The Proposition 8 interface** (plus the §5.1 exact recursive contract), for a complex
network with `W` wires, label dimension `m` and `s` residual basis vectors in total. -/
structure Prop8Iface (W m s : ℕ) where
  net : (κ : Type) → [Fintype κ] → [DecidableEq κ] → List (NStep W m κ)
  count : ∀ (κ : Type) [Fintype κ] [DecidableEq κ], (net κ).countP NStep.isChild = s
  correct : ∀ (κ : Type) [Fintype κ] [DecidableEq κ],
    denList (Butterfly.tens Butterfly.Cm) (net κ) = (roleDiag (Butterfly.tens Butterfly.Cm) : NOp W (Fin m × κ))

/-! ## The recursion -/

/-- Selected bits of a node `k` levels above leaves of width `u`. -/
def Lvl (u m : ℕ) : ℕ → Type
  | 0 => Fin u
  | k + 1 => Fin m × Lvl u m k

instance instFintypeLvl (u m : ℕ) : ∀ k, Fintype (Lvl u m k)
  | 0 => inferInstanceAs (Fintype (Fin u))
  | k + 1 => @instFintypeProd (Fin m) (Lvl u m k) _ (instFintypeLvl u m k)

instance instDecEqLvl (u m : ℕ) : ∀ k, DecidableEq (Lvl u m k)
  | 0 => inferInstanceAs (DecidableEq (Fin u))
  | k + 1 => @instDecidableEqProd (Fin m) (Lvl u m k) _ (instDecEqLvl u m k)

theorem card_Lvl (u m : ℕ) : ∀ k, Fintype.card (Lvl u m k) = u * m ^ k
  | 0 => by simp only [pow_zero, mul_one]; exact Fintype.card_fin u
  | k + 1 => by
    rw [show Fintype.card (Lvl u m (k + 1)) = Fintype.card (Fin m × Lvl u m k) from rfl,
      Fintype.card_prod, card_Lvl u m k, Fintype.card_fin, pow_succ]; ring

/-- Leaves: `u` individual kernels (25) in sequence. -/
noncomputable def leafOp (u : ℕ) : Op (Fin u) :=
  ((List.finRange u).map fun i => tensF (kerAt i Butterfly.Cm)).prod

variable {W m s : ℕ} [NeZero W] (I : Prop8Iface W m s) (u : ℕ)

/-- The recursive program: at depth `k+1`, one network invocation whose child calls run the
depth-`k` program; the parent stream is read off role `0` (all roles agree, see `recOp_roles`). -/
noncomputable def recOp : (k : ℕ) → Op (Lvl u m k)
  | 0 => leafOp u
  | k + 1 => roleBlock (denList (recOp k) (I.net (Lvl u m k))) 0

/-- Cost of the recursive program: local work `Cloc (k+1) V` at a depth-`k+1` node on volume `V`
plus one depth-`k` call on volume `V/W` per child step; leaves cost `Cleaf V`. -/
noncomputable def recCost (Cloc : ℕ → ℝ → ℝ) (Cleaf : ℝ → ℝ) : ℕ → ℝ → ℝ
  | 0, V => Cleaf V
  | k + 1, V => Cloc (k + 1) V +
      ((I.net (Lvl u m k)).countP NStep.isChild : ℝ) * recCost Cloc Cleaf k (V / W)

/-! ## Pieces on disjoint blocks of selected bits -/

section Blocks

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

/-- The block family: `A` on the block `S`, identity elsewhere. -/
def blockFam (S : Finset ι) (A : M2) : ι → M2 := fun i => if i ∈ S then A else 1

/-- A piece computed on its own index type `κ ≃ S`, embedded into the full selected-bit set. -/
def embed {κ : Type} [Fintype κ] (S : Finset ι) (e : κ ≃ S) (T : Op κ) : Op ι := fun x y =>
  T (fun j => x (e j)) (fun j => y (e j)) * (if ∀ i, i ∉ S → x i = y i then 1 else 0)

end Blocks

/-! ## Proposition 14: the normalized layer from the C-layer -/

section Prop14

variable {P K D S ρ r : ℕ}

/-- `Q_p` applied to every coefficient of every record. -/
noncomputable def QpArr (p : ℕ) (xs : List (Fin r → ℂ)) : List (Fin r → ℂ) :=
  xs.map fun v j => Qp p (v j)

/-- The layer procedure: phase `S_D`, the C-layer `Cimpl`, phase `S_D`, the scalar
`b^D = (-i/2)^{⌊D/2⌋} b^{D mod 2}` (a shift and a fourth-root phase, at most one factor `b`),
and the final truncation `Q_p`. -/
noncomputable def prop14Proc (hρ : ρ < K) (p : ℕ) (Cimpl : Op (Fin D)) (xs : List (Fin r → ℂ)) :
    List (Fin r → ℂ) :=
  QpArr p ((selOp (layerSel P K D S ρ hρ) SD (selOp (layerSel P K D S ρ hρ) Cimpl
    (selOp (layerSel P K D S ρ hρ) SD xs))).map
      (((-Complex.I / 2) ^ (D / 2) * Butterfly.cb ^ (D % 2)) • ·))

end Prop14

end Butterfly

end

section
-- Source module: Solutions.SynthDFT
/-!
# §6.1 Synthetic roots, character cancellation, and the normalized transforms

Records live in a commutative `ℂ`-algebra `R` containing a unit `y` with `y ^ r = -1`,
`r = 2^ℓ` (the model is `R_r = ℂ[y]/(y^r+1)`, see `SynthNorm.lean`). For an axis of length
`t = 2^a` with `a ≤ ℓ + 1` the root is `ω_t = y^{2r/t} = y^{2^{ℓ+1-a}}`.

* `yp y e = y^e` for `e : ℤ` (`yp_add`, `yp_two_r_mul`, `yp_r_mul_odd`).
* **Lemma 15** `char_sum`: `∑_{k<t} ω_t^{jk} = t` if `t ∣ j`, else `0` (any integer `j`);
  proved by the shift argument, using that `R` has no two-torsion.
* Multidimensional transforms `Fm` (`F^-`) and `Fp` (`F^+`) on arrays `(Fin D → ℕ) → R`
  supported on the box `∏ [0, t_i)`; `char_sum_box`, and the inversion formulas
  `Fp_Fm : F^+ F^- = M⁻¹ I`, `Fm_Fp : F^- F^+ = M⁻¹ I` on the box (Lemma 15, last part).
-/

set_option linter.unusedSectionVars false

namespace Synth

open Finset

variable {R : Type*} [CommRing R] [Algebra ℂ R]

/-! ## 1. Powers of `y` -/

/-- `y^e` for an integer exponent `e`. -/
def yp (y : Rˣ) (e : ℤ) : R := ((y ^ e : Rˣ) : R)

variable (y : Rˣ)

variable {r : ℕ} (hy : y ^ r = -1)
include hy

omit hy

/-! ## 2. Lemma 15 (synthetic character cancellation) -/

/-- Exponent constant `2r/t = 2^{ℓ+1-a}`. -/
def cexp (ℓ a : ℕ) : ℤ := 2 ^ (ℓ + 1 - a)

/-! ## 3. The box, and multidimensional transforms -/

variable {D : ℕ}

/-- The box `∏ [0, 2^{a_i})` of axis coordinates. -/
def box (a : Fin D → ℕ) : Finset (Fin D → ℕ) := Fintype.piFinset fun i => range (2 ^ a i)

/-- `M = ∏ t_i`. -/
def Mv (a : Fin D → ℕ) : ℕ := ∏ i, 2 ^ a i

lemma mem_box {a : Fin D → ℕ} {x : Fin D → ℕ} : x ∈ box a ↔ ∀ i, x i < 2 ^ a i := by
  simp [box, Fintype.mem_piFinset]

/-- The exponent `∑_i (2r/t_i) j_i k_i`. -/
def pair (ℓ : ℕ) (a : Fin D → ℕ) (j k : Fin D → ℕ) : ℤ := ∑ i, cexp ℓ (a i) * j i * k i

/-- `F^-` on arrays over `R`, axes of lengths `2^{a_i}`:
`(F^- z)_j = M⁻¹ ∑_{k ∈ box} y^{-∑ (2r/t_i) j_i k_i} z_k`. -/
noncomputable def Fm (ℓ : ℕ) (a : Fin D → ℕ) (z : (Fin D → ℕ) → R) (j : Fin D → ℕ) : R :=
  ((Mv a : ℂ)⁻¹) • ∑ k ∈ box a, yp y (-pair ℓ a j k) * z k

/-- `F^+`: `(F^+ z)_j = M⁻¹ ∑_{k ∈ box} y^{∑ (2r/t_i) j_i k_i} z_k`. -/
noncomputable def Fp (ℓ : ℕ) (a : Fin D → ℕ) (z : (Fin D → ℕ) → R) (j : Fin D → ℕ) : R :=
  ((Mv a : ℂ)⁻¹) • ∑ k ∈ box a, yp y (pair ℓ a j k) * z k

end Synth

end

section
-- Source module: Solutions.SynthLevels
/-!
# §6.1 Butterflies, decimation in frequency, and the stored frequency order

One level of the decimation-in-frequency procedure on an axis of length `t = 2^a`, at named slot
`h` (address bit `h`): the butterfly `H0` on bit `h`, then the twiddle `ω_{2^{h+1}}^{-b k}` where
`b` is the new branch bit and `k = x mod 2^h` is encoded by the unprocessed lower slots
(`lev`). After `s` levels (slots `a-1, …, a-s`) the array is the partial transform with kernel
`Ws s` (`levCore`): the processed slots hold the frequency bits `0, …, s-1` (from the most
significant slot), which is the inductive description of §6.1.

Multidimensional arrays `(Fin D → ℕ) → R`: `axisK i W` applies a one-dimensional kernel along
axis `i`, `axisL i h` a level along axis `i`. Levels on different axes commute with each other's
kernels, so *any* schedule that runs every axis's slots from high to low gives the same result
(`runF_sched`). With `rev` the bit reversal (`Bt`), the result is `(B F^-) z` (`synth_forward`):
the entry with frequency `j` is stored at `rev j`. The paper's round schedule (round `h` on all
axes with `a_i > h`, for `h = ℓ-1, …, 0`) is such a schedule (`paperSched_ok`).
-/

set_option linter.unusedSectionVars false

namespace Synth

open Finset Butterfly

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ)

/-! ## 1. Bits -/

lemma mod_pow_succ_eq (h v : ℕ) : v % 2 ^ (h + 1) = (getBit h v : ℕ) * 2 ^ h + v % 2 ^ h := by
  simp only [getBit]
  rw [pow_succ, Nat.mod_mul, Nat.add_comm, Nat.mul_comm]

/-! ## 2. Frequency bits and the one-dimensional partial transforms -/

/-- The frequency written in the `s` most significant slots of an axis of length `2^a`:
slot `a-1-c` holds frequency bit `c`. -/
def fq (a s x : ℕ) : ℕ := ∑ c ∈ range s, (getBit (a - 1 - c) x : ℕ) * 2 ^ c

/-- The record permutation `B_t`: bit reversal of `a` binary digits. -/
def rev (a x : ℕ) : ℕ := fq a a x

/-- Kernel of the partial transform after `s` levels:
`Ws s x c = 2^{-s} ω_t^{-f c} [c ≡ x mod 2^{a-s}]` with `f = fq a s x`. -/
noncomputable def Ws (ℓ a s x c : ℕ) : R :=
  ((2 : ℂ) ^ s)⁻¹ • (yp y (-(cexp ℓ a * fq a s x * c)) *
    if c % 2 ^ (a - s) = x % 2 ^ (a - s) then 1 else 0)

/-- The twiddle exponent of a level at slot `h`: `-(2r/2^{h+1}) · b · k`, `b` = bit `h`,
`k = x mod 2^h`. -/
def twE (ℓ h x : ℕ) : ℤ := -(2 ^ (ℓ - h) * (getBit h x : ℕ) * (x % 2 ^ h) : ℕ)

/-- One forward level at slot `h` (one-dimensional): butterfly then twiddle. -/
noncomputable def lev (ℓ h : ℕ) (g : ℕ → R) (x : ℕ) : R :=
  yp y (twE ℓ h x) * ∑ b : Fin 2, H0 (getBit h x) b • g (setBit h x b)

/-! ## 3. Endpoints of the partial transforms -/

/-- One-dimensional kernel operator `(kop t W g)(x) = ∑_{c<t} W x c · g c`. -/
def kop (t : ℕ) (W : ℕ → ℕ → R) (g : ℕ → R) (x : ℕ) : R := ∑ c ∈ range t, W x c * g c

/-! ## 4. Multidimensional arrays: operators along one axis -/

variable {D : ℕ}

/-- Arrays of records indexed by axis coordinates. -/
abbrev Arr (R : Type*) (D : ℕ) := (Fin D → ℕ) → R

/-- A one-dimensional kernel applied along axis `i`. -/
def axisK (i : Fin D) (t : ℕ) (W : ℕ → ℕ → R) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  ∑ c ∈ range t, W (x i) c * z (Function.update x i c)

/-- A forward level at slot `h` along axis `i`. -/
noncomputable def axisL (ℓ : ℕ) (i : Fin D) (h : ℕ) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  lev y ℓ h (fun c => z (Function.update x i c)) (x i)

/-! ## 5. The state after any number of levels per axis -/

/-- The array after `s i` levels on each axis `i`: the product of the axis partial transforms. -/
noncomputable def State (ℓ : ℕ) (a s : Fin D → ℕ) (z : Synth.Arr R D) : Synth.Arr R D :=
  (List.finRange D).foldr (fun i acc => axisK i (2 ^ a i) (Ws y ℓ (a i) (s i)) acc) z

/-! ## 6. Schedules -/

/-- Run a list of levels `(axis, slot)`, first element first. -/
noncomputable def runF (ℓ : ℕ) : List (Fin D × ℕ) → Synth.Arr R D → Synth.Arr R D
  | [], z => z
  | st :: l, z => runF ℓ l (axisL y ℓ st.1 st.2 z)

/-- A schedule from per-axis counts `s`: every axis `j` runs its remaining slots
`a_j - 1 - s_j, …, 0` from high to low (in any interleaving with the other axes). -/
def SchedOK (a s : Fin D → ℕ) (l : List (Fin D × ℕ)) : Prop :=
  ∀ j, (l.filter (fun st => st.1 = j)).map Prod.snd = (List.range (a j - s j)).reverse

end Synth

end

section
-- Source module: Solutions.SynthSched
/-!
# §6.1 The forward and opposite synthetic procedures: `B F^-` and `F^+ B^{-1}`

* `foldr_axisK_eq`: a product of axis kernels is the separable kernel `∏_i W_i(x_i, m_i)`.
* `State_zero`, `State_full`: before any level the state is the input (on the box); after all
  levels it is `F^-` read at the bit-reversed address.
* `synth_forward` (**(28), forward**): for every schedule running each axis's slots from high to
  low, `runF l z x = (F^- z)(rev x)` on the box, i.e. the procedure implements `B_t F^-`.
* `levInv`, `runI`: the opposite procedure (inverse monomial first, then `H0`, schedule
  reversed); `runI_runF`: it composes with the forward procedure to `2^{-#levels} = M⁻¹`;
  `synth_opposite` (**(28), opposite**): it implements `F^+ B^{-1}`.
* `rev_rev`, `rev_lt`: bit reversal is an involution on `[0, 2^a)`.
* `paperSched`: the round schedule of §6.2 (rounds `h = ℓ-1, …, 0`, each on the axes with
  `a_i > h`); `paperSched_ok`.
-/

set_option linter.unusedSectionVars false

namespace Synth

open Finset Butterfly

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ) {D : ℕ}

/-! ## 1. Separable kernels -/

/-- The partial box: full range on the axes of `l`, the single value `x j` elsewhere. -/
def pbox (t : Fin D → ℕ) (l : List (Fin D)) (x : Fin D → ℕ) : Finset (Fin D → ℕ) :=
  Fintype.piFinset fun j => if j ∈ l then range (t j) else {x j}

/-! ## 2. Endpoints -/

/-- Bit reversal on every axis. -/
def revV (a : Fin D → ℕ) (x : Fin D → ℕ) : Fin D → ℕ := fun i => Synth.rev (a i) (x i)

/-! ## 3. Locality: levels read only inside the box -/

/-- Two arrays agreeing on the box. -/
def EqOn (a : Fin D → ℕ) (z z' : Synth.Arr R D) : Prop := ∀ x ∈ box a, z x = z' x

/-! ## 4. Bit reversal is an involution -/

lemma sum_bits_eq_mod (a x : ℕ) : ∑ c ∈ range a, (getBit c x : ℕ) * 2 ^ c = x % 2 ^ a := by
  induction a with
  | zero => simp [Nat.mod_one]
  | succ a ih => rw [sum_range_succ, ih, mod_pow_succ_eq]; ring

/-! ## 5. The opposite procedure -/

/-- One opposite level at slot `h`: the inverse monomial first, then the butterfly `H0`. -/
noncomputable def levInv (ℓ h : ℕ) (g : ℕ → R) (x : ℕ) : R :=
  ∑ b : Fin 2, H0 (getBit h x) b • (yp y (-twE ℓ h (setBit h x b)) * g (setBit h x b))

/-- An opposite level along axis `i`. -/
noncomputable def axisLInv (ℓ : ℕ) (i : Fin D) (h : ℕ) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  levInv y ℓ h (fun c => z (Function.update x i c)) (x i)

/-- The opposite procedure: the schedule traversed in reverse, with opposite levels. -/
noncomputable def runI (ℓ : ℕ) : List (Fin D × ℕ) → Synth.Arr R D → Synth.Arr R D
  | [], z => z
  | st :: l, z => axisLInv y ℓ st.1 st.2 (runI ℓ l z)

/-! ## 6. The round schedule of §6.2 -/

end Synth

end

section
-- Source module: Solutions.SynthList
/-!
# §6.1 on the tape: grouped rounds as `selOp`-layers followed by the monomial `J_E`

A *coordinate layout* identifies the valid physical addresses of a shape `sh` (lists of digits,
`AModel`) with the box of axis coordinates. A round at slot `h` acts with the simultaneous
butterfly layer `H = selOp Sb (tens H0)` of §5 (`ButterflyLayer`), where the selected bits
`Sb` are the physical positions of the named slots `(i, h)` for the participating axes `P`
(`SelCompat`), and then with the recordwise monomial `J_E`, `E = ∑_{i ∈ P} -(2r/2^{h+1}) b_i k_i`
(29) (`roundL`).

* `fn_roundL` (grouping, (29)): on the box, `J_E ∘ H` is the composition of the individual axis
  levels `axisL i h`, `i ∈ P` (the twiddle of axis `i` commutes with the other butterflies).
* `synth_forward_list` (**§6.1 on lists**): running the rounds `h = ℓ-1, …, 0`, each on the axes
  with `a_i > h`, the record at physical address `v` is `(F^- z)(rev (coord v))`: the procedure
  implements `B F^-` in the given layout.
-/

set_option linter.unusedSectionVars false

namespace Synth

open Finset Butterfly AModel

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ) {D : ℕ}

/-! ## 1. Butterflies and twiddles as separate array operators -/

/-- The butterfly `H0` on slot `h` of axis `i`. -/
noncomputable def bfly (h : ℕ) (i : Fin D) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  ∑ b : Fin 2, H0 (getBit h (x i)) b • z (Function.update x i (setBit h (x i) b))

/-- The twiddle of axis `i` at slot `h`: multiply by `y^{twE}`. -/
noncomputable def twid (ℓ h : ℕ) (i : Fin D) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  yp y (twE ℓ h (x i)) * z x

/-- All twiddles of a list of axes at once: `J_E` with `E = ∑ twE`. -/
noncomputable def twidAll (ℓ h : ℕ) (l : List (Fin D)) (z : Synth.Arr R D) (x : Fin D → ℕ) : R :=
  yp y ((l.map fun i => twE ℓ h (x i)).sum) * z x

/-! ## 2. Coordinate layouts and compatible selected bits -/

/-- Valid physical addresses of `sh` ↔ axis coordinates in the box. -/
structure CoordLayout (sh : Shape) (a : Fin D → ℕ) where
  addr : (Fin D → ℕ) → List ℕ
  coord : List ℕ → (Fin D → ℕ)
  valid_addr : ∀ x ∈ box a, Valid sh (addr x)
  coord_addr : ∀ x ∈ box a, coord (addr x) = x
  addr_coord : ∀ v, Valid sh v → addr (coord v) = v
  coord_mem : ∀ v, Valid sh v → coord v ∈ box a

variable {sh : Shape} {a : Fin D → ℕ}

/-- The array of records (in physical order) read in axis coordinates. -/
def fn (L : CoordLayout sh a) (xs : List R) : Synth.Arr R D := fun x => xs.getD (encode sh (L.addr x)) 0

/-- Selected bits `Sb` read and write exactly the named slots `(i, h)`, `i ∈ P`. -/
structure SelCompat (L : CoordLayout sh a) (h : ℕ) (P : Finset (Fin D))
    (Sb : SelBits sh {i // i ∈ P}) : Prop where
  bits_eq : ∀ v, Valid sh v → ∀ i, Sb.bits v i = getBit h (L.coord v i.1)
  set_eq : ∀ v, Valid sh v → ∀ yb, L.coord (Sb.set v yb) =
    fun j => if hj : j ∈ P then setBit h (L.coord v j) (yb ⟨j, hj⟩) else L.coord v j

variable (L : CoordLayout sh a)

/-! ## 3. Rounds -/

/-- A round on the tape: the layer `H`, then the recordwise monomial `J_E` (29). -/
noncomputable def roundL (ℓ h : ℕ) (P : Finset (Fin D)) (Sb : SelBits sh {i // i ∈ P})
    (xs : List R) : List R :=
  (List.range (vol sh)).map fun j =>
    yp y ((((Finset.univ : Finset {i // i ∈ P}).toList.map Subtype.val).map
      fun i => twE ℓ h (L.coord (decode sh j) i)).sum) * (selOp Sb (Butterfly.tens H0) xs).getD j 0

/-- The axis list of a round. -/
noncomputable def roundAxes (P : Finset (Fin D)) : List (Fin D) :=
  (Finset.univ : Finset {i // i ∈ P}).toList.map Subtype.val

/-! ## 4. The full forward procedure on the tape -/

/-- Rounds `n-1, …, 0` (round `h` on the axes `P h` with selected bits `Sb h`). -/
noncomputable def procL (ℓ : ℕ) (P : ℕ → Finset (Fin D)) (Sb : (h : ℕ) → SelBits sh {i // i ∈ P h}) :
    ℕ → List R → List R
  | 0, xs => xs
  | n + 1, xs => procL ℓ P Sb n (roundL y L ℓ n (P n) (Sb n) xs)

/-- The level schedule of `procL`. -/
noncomputable def schedL (P : ℕ → Finset (Fin D)) : ℕ → List (Fin D × ℕ)
  | 0 => []
  | n + 1 => ((roundAxes (P n)).reverse.map fun i => (i, n)) ++ schedL P n

/-- The participating axes of round `h`: those with `a_i > h`. -/
def partAxes (a : Fin D → ℕ) (h : ℕ) : Finset (Fin D) := Finset.univ.filter fun i => h < a i

end Synth

end

section
-- Source module: Solutions.SynthLayout
/-!
# §6.2 Lemma 16: the positional layout permutation `L`

Record addresses are lists of address bits (shape `replicate n 2`); every bit position holds a
*named slot* (a natural-number name, `nm i h` for slot `h` of axis `i`). Two kinds of primitive
moves act on positions:

* `LOp.mv c p` — move the bit at position `p ≥ c` to position `c` (the single-field move of
  Lemma 3, `rotIdx` of `AModelL12Basic`);
* `LOp.xch i j K` — exchange two width-`K` address chunks (`blockSwap`; by `pull_blockSwap` this is
  the chunk interchange of Lemma 12, `lemma12_bit`, on the merged shape).

`applyN ops arr` is the resulting arrangement of slot names, `runArr ops xs` the array
denotation (each primitive is the `pull` by the inverse position map).

* `runArr_eq_pull` (**denotation**): `runArr ops xs = pull (relabel init (applyN ops init))`:
  the composite moves the record whose slot values are `v` from the address with slots in the
  initial order to the address with slots in the final order (`relabel`).
* `prefix_moves`: listed slots can be moved one by one to a prefix with `|Lst|` single moves,
  leaving the others in order.
* `chunk_sort`: any rearrangement of `n` distinct `K`-chunks by at most `n - 1` exchanges.
-/

set_option linter.unusedSectionVars false

namespace Synth

open AModel AModelL12

/-! ## 1. Primitive moves -/

/-- A primitive positional move. -/
inductive LOp where
  /-- move position `p` to position `c` (`c ≤ p`) -/
  | mv (c p : ℕ)
  /-- exchange the chunks `[i, i+K)` and `[j, j+K)` (`i + K ≤ j`) -/
  | xch (i j K : ℕ)

namespace LOp

/-- Position map of the arrangement: new position `k` holds old position `idx k`. -/
def idx : LOp → ℕ → ℕ
  | mv c p => rotIdx c (p - c) 1
  | xch i j K => bsIdx i j K

/-- Inverse position map (the address `pull` of the primitive). -/
def inv : LOp → ℕ → ℕ
  | mv c p => rotIdx c 1 (p - c)
  | xch i j K => bsIdx i j K

/-- Well-formed on `n` positions. -/
def WF (n : ℕ) : LOp → Prop
  | mv c p => c ≤ p ∧ p < n
  | xch i j K => i + K ≤ j ∧ j + K ≤ n

def isMv : LOp → Bool
  | mv _ _ => true
  | xch _ _ _ => false

end LOp

/-- The arrangement of slot names after a list of moves. -/
def applyN (ops : List LOp) (arr : List ℕ) : List ℕ :=
  ops.foldl (fun acc o => permIdx o.idx acc) arr

/-- The binary address shape. -/
def binSh (n : ℕ) : Shape := List.replicate n 2

/-- Array denotation: each primitive is the `pull` by its inverse position map. -/
def runArr {α : Type*} (n : ℕ) (d : α) (ops : List LOp) (xs : List α) : List α :=
  ops.foldl (fun acc o => pull (binSh n) (binSh n) (permIdx o.inv) d acc) xs

/-- The address map of a slot relabelling: the output address `y` (slots in order `tgt`) reads
the input address listing, for each slot of `init`, its value in `y`. -/
def relabel (init tgt : List ℕ) (y : List ℕ) : List ℕ := init.map fun nm => y.getD (tgt.idxOf nm) 0

/-! ## 2. Prefix moves -/

lemma getD_take_eq (l : List ℕ) (m k : ℕ) (hk : k < m) : (l.take m).getD k 0 = l.getD k 0 := by
  rw [List.getD_eq_getElem?_getD, List.getElem?_take_of_lt hk, List.getD_eq_getElem?_getD]

lemma getD_drop_eq (l : List ℕ) (j k : ℕ) : (l.drop j).getD k 0 = l.getD (j + k) 0 := by
  rw [List.getD_eq_getElem?_getD, List.getElem?_drop, List.getD_eq_getElem?_getD]

/-! ## 3. Chunk exchanges -/

end Synth

end

section
-- Source module: Solutions.SynthLayoutPaper
/-!
# §6.2 Lemma 16: the paper's layout `L`

Axes `i < D` have `a_i = ℓ` (long) or `ℓ - 1` (short) address bits; slot `h` of axis `i` is named
`nm i h = h·D + i`. With `b = (ℓ-1) mod K` and `q = ⌊(ℓ-1)/K⌋`:

* the **prefix list** `Lst`: the top slots `(i, ℓ-1)` of the long axes (increasing `i`), then
  `(i, h)` for `h = b-1, …, 0` and, within each `h`, increasing `i`;
* after the prefix moves every axis keeps the `qK` slots `h = ℓ-2, …, b` (`restA`), which are
  `q` consecutive `K`-bit chunks (`chunk i j`, numbered from the top);
* the target order `tgt = Lst ++ (chunks in chunk-major order)`.

`paper_layout`: `L` is reached by `|Lst| ≤ D(b+1)` single moves and at most `Dq - 1` chunk
exchanges, and its array denotation is the `pull` by `relabel init tgt`.
`tgt_slot`: in the target, slot `(i, h)` (`b ≤ h ≤ ℓ-2`) lies in chunk `i` of chunk group
`j = ⌊(qK-1-(h-b))/K⌋`, at offset `(h - b) mod K` from the least significant end — the common
offset used by the simultaneous layer (Proposition 14).
-/

set_option linter.unusedSectionVars false

namespace Synth

open AModel AModelL12

variable (D ℓ K : ℕ) (long : Fin D → Bool)

/-- Number of address bits of axis `i`. -/
def al (i : Fin D) : ℕ := if long i then ℓ else ℓ - 1

/-- `b = (ℓ-1) mod K`. -/
def bb : ℕ := (ℓ - 1) % K
/-- `q = ⌊(ℓ-1)/K⌋`. -/
def qq : ℕ := (ℓ - 1) / K

/-- Slot names. -/
def nm (i : Fin D) (h : ℕ) : ℕ := h * D + i

/-- Initial order: axis order, and within each axis decreasing `h`. -/
def initL : List ℕ := (List.finRange D).flatMap fun i => (List.range (al D ℓ long i)).reverse.map (nm D i)

/-- The prefix list. -/
def lstL : List ℕ :=
  ((List.finRange D).filter fun i => long i).map (fun i => nm D i (ℓ - 1)) ++
    (List.range (bb ℓ K)).reverse.flatMap fun h => (List.finRange D).map fun i => nm D i h

/-- The `qK` remaining slots `h = b + qK - 1, …, b` of axis `i`. -/
def restA (i : Fin D) : List ℕ :=
  (List.range (qq ℓ K * K)).reverse.map fun t => nm D i (bb ℓ K + t)

/-- Chunk `j` (from the top) of axis `i`. -/
def chunk (i : Fin D) (j : ℕ) : List ℕ := ((restA D ℓ K i).drop (j * K)).take K

/-- Axis-major chunk order (after the prefix moves). -/
def axisMajor : List (List ℕ) := (List.finRange D).flatMap fun i => (List.range (qq ℓ K)).map (chunk D ℓ K i)

/-- Chunk-major chunk order (the target). -/
def chunkMajor : List (List ℕ) := (List.range (qq ℓ K)).flatMap fun j => (List.finRange D).map fun i => chunk D ℓ K i j

/-- The target slot order `L`. -/
def tgtL : List ℕ := lstL D ℓ K long ++ (chunkMajor D ℓ K).flatten

/-! ## 1. Names -/

lemma nm_inj {i i' : Fin D} {h h' : ℕ} (he : nm D i h = nm D i' h') : i = i' ∧ h = h' := by
  unfold nm at he
  have hi := i.isLt; have hi' := i'.isLt
  have h1 : (h * D + i) % D = (h' * D + i') % D := by rw [he]
  rw [Nat.mul_comm, Nat.mul_add_mod, Nat.mod_eq_of_lt hi, Nat.mul_comm, Nat.mul_add_mod,
    Nat.mod_eq_of_lt hi'] at h1
  have hii : i = i' := Fin.ext h1
  subst hii
  refine ⟨rfl, ?_⟩
  have : h * D = h' * D := by omega
  exact Nat.eq_of_mul_eq_mul_right (by omega) this

lemma mem_lstL {i : Fin D} {h : ℕ} :
    nm D i h ∈ lstL D ℓ K long ↔ (long i = true ∧ h = ℓ - 1) ∨ h < bb ℓ K := by
  unfold lstL
  simp only [List.mem_append, List.mem_map, List.mem_filter, List.mem_finRange, true_and,
    List.mem_flatMap, List.mem_reverse, List.mem_range]
  constructor
  · rintro (⟨i', hl, he⟩ | ⟨h', hh', i', he⟩)
    · obtain ⟨h1, h2⟩ := nm_inj D he; subst h1; exact Or.inl ⟨hl, h2.symm⟩
    · obtain ⟨h1, h2⟩ := nm_inj D he; subst h1; subst h2; exact Or.inr hh'
  · rintro (⟨hl, rfl⟩ | hh)
    · exact Or.inl ⟨i, hl, rfl⟩
    · exact Or.inr ⟨h, hh, i, rfl⟩

/-! ## 2. Counting -/

lemma length_restA (i : Fin D) : (restA D ℓ K i).length = qq ℓ K * K := by simp [restA]

/-! ## 3. Flattening uniform blocks -/

lemma flatten_chunks (l : List ℕ) : ∀ q, l.length = q * K →
    ((List.range q).map fun j => (l.drop (j * K)).take K).flatten = l
  | 0, h => by simp at h ⊢; exact List.eq_nil_of_length_eq_zero (by simpa using h)
  | q + 1, h => by
    rw [List.range_succ_eq_map, List.map_cons, List.flatten_cons, List.map_map]
    have ih := flatten_chunks (l.drop K) q (by rw [List.length_drop, h, Nat.succ_mul]; omega)
    have e : ((fun j => (l.drop (j * K)).take K) ∘ Nat.succ) =
        fun j => ((l.drop K).drop (j * K)).take K := by
      funext j; simp only [Function.comp, List.drop_drop]; congr 2; rw [Nat.succ_mul]; ring
    rw [e, ih]; simp

lemma flatten_flatMap' {α β : Type*} (f : α → List (List β)) :
    ∀ l : List α, (l.flatMap f).flatten = l.flatMap fun a => (f a).flatten
  | [] => rfl
  | a :: l => by simp [List.flatMap_cons, List.flatten_append, flatten_flatMap' f l]

/-! ## 4. After the prefix moves -/

lemma range_filter_window (b m n : ℕ) (h : b + m ≤ n) :
    (List.range n).filter (fun x => decide (b ≤ x ∧ x < b + m)) = (List.range m).map (b + ·) := by
  obtain ⟨k, rfl⟩ : ∃ k, n = b + m + k := ⟨n - (b + m), by omega⟩
  rw [List.range_add, List.range_add, List.filter_append, List.filter_append]
  rw [List.filter_eq_nil_iff.2 (fun x hx => by simp at hx ⊢; omega),
    List.filter_eq_self.2 (fun x hx => by simp at hx ⊢; omega),
    List.filter_eq_nil_iff.2 (fun x hx => by simp at hx ⊢; omega)]
  simp [List.map_map, Function.comp_def]

theorem initL_filter (hK : 0 < K) (hℓ : 1 ≤ ℓ) :
    (initL D ℓ long).filter (· ∉ lstL D ℓ K long) = (axisMajor D ℓ K).flatten := by
  have hbq : bb ℓ K + qq ℓ K * K = ℓ - 1 := by
    unfold bb qq; rw [Nat.mul_comm]; exact Nat.mod_add_div _ _
  unfold initL axisMajor
  rw [List.filter_flatMap, flatten_flatMap']
  apply List.flatMap_congr; intro i _
  rw [show ((List.range (qq ℓ K)).map (chunk D ℓ K i)).flatten = restA D ℓ K i from
    flatten_chunks K (restA D ℓ K i) (qq ℓ K) (length_restA D ℓ K i)]
  rw [List.filter_map, List.filter_reverse]
  have : ((List.range (al D ℓ long i)).filter
      ((fun x => decide (x ∉ lstL D ℓ K long)) ∘ nm D i)) =
      (List.range (al D ℓ long i)).filter (fun x => decide (bb ℓ K ≤ x ∧ x < bb ℓ K + qq ℓ K * K)) := by
    apply List.filter_congr; intro x hx
    simp only [Function.comp_apply, mem_lstL, List.mem_range, al] at hx ⊢
    by_cases hl : long i
    · simp [hl] at hx ⊢
      by_cases h1 : x = ℓ - 1 <;> by_cases h2 : bb ℓ K ≤ x <;> simp [h1, h2] <;> omega
    · simp [hl] at hx ⊢; intro; omega
  rw [this, range_filter_window _ _ _ (by unfold al; split_ifs <;> omega)]
  simp [restA, List.map_reverse, List.map_map, Function.comp_def]

/-! ## 5. Nodup and the chunk permutation -/

lemma initL_nodup : (initL D ℓ long).Nodup := by
  unfold initL
  rw [List.nodup_flatMap]
  refine ⟨fun i _ => ?_, ?_⟩
  · exact (List.nodup_reverse.2 (List.nodup_range)).map (fun h h' e => (nm_inj D e).2)
  · refine (List.nodup_finRange D).pairwise_of_forall_ne fun i _ i' _ hii' => ?_
    rw [Function.onFun, List.disjoint_left]
    intro x hx hx'
    simp only [List.mem_map] at hx hx'
    obtain ⟨h, -, rfl⟩ := hx; obtain ⟨h', -, he⟩ := hx'
    exact hii' (nm_inj D he).1.symm

lemma lstL_nodup : (lstL D ℓ K long).Nodup := by
  unfold lstL
  rw [List.nodup_append]
  refine ⟨?_, ?_, ?_⟩
  · exact ((List.nodup_finRange D).filter _).map_on (fun i _ i' _ e => (nm_inj D e).1)
  · rw [List.nodup_flatMap]
    refine ⟨fun h _ => (List.nodup_finRange D).map_on (fun i _ i' _ e => (nm_inj D e).1), ?_⟩
    refine (List.nodup_reverse.2 List.nodup_range).pairwise_of_forall_ne fun h _ h' _ hh' => ?_
    rw [Function.onFun, List.disjoint_left]
    intro x hx hx'
    simp only [List.mem_map] at hx hx'
    obtain ⟨i, -, rfl⟩ := hx; obtain ⟨i', -, he⟩ := hx'
    exact hh' (nm_inj D he).2.symm
  · intro x hx y hy hxy
    subst hxy
    simp only [List.mem_map, List.mem_filter, List.mem_flatMap, List.mem_reverse,
      List.mem_range] at hx hy
    obtain ⟨i, -, rfl⟩ := hx
    obtain ⟨h, hh, i', -, he⟩ := hy
    have := (nm_inj D he).2
    have : bb ℓ K ≤ ℓ - 1 := Nat.mod_le _ _
    omega

lemma lstL_sub_initL (hℓ : 1 ≤ ℓ) : ∀ x ∈ lstL D ℓ K long, x ∈ initL D ℓ long := by
  intro x hx
  unfold lstL at hx
  simp only [List.mem_append, List.mem_map, List.mem_filter, List.mem_finRange, true_and,
    List.mem_flatMap, List.mem_reverse, List.mem_range] at hx
  unfold initL
  simp only [List.mem_flatMap, List.mem_finRange, true_and, List.mem_map, List.mem_reverse,
    List.mem_range]
  rcases hx with ⟨i, hl, rfl⟩ | ⟨h, hh, i, -, rfl⟩
  · exact ⟨i, ℓ - 1, by unfold al; simp [hl]; omega, rfl⟩
  · have : bb ℓ K ≤ ℓ - 1 := Nat.mod_le _ _
    exact ⟨i, h, by unfold al; split_ifs <;> omega, rfl⟩

/-- First slot of a chunk (determines the chunk). -/
lemma chunk_head (hK : 0 < K) (i : Fin D) {j : ℕ} (hj : j < qq ℓ K) :
    (chunk D ℓ K i j).getD 0 0 = nm D i (bb ℓ K + (qq ℓ K * K - 1 - j * K)) := by
  unfold chunk
  have hlt : j * K < qq ℓ K * K := Nat.mul_lt_mul_of_pos_right hj hK
  rw [getD_take_eq _ _ _ hK, getD_drop_eq, add_zero, List.getD_eq_getElem _ _
    (by rw [length_restA]; exact hlt)]
  unfold restA
  rw [List.getElem_map, List.getElem_reverse, List.getElem_range]
  simp

lemma chunk_inj (hK : 0 < K) {i i' : Fin D} {j j' : ℕ} (hj : j < qq ℓ K) (hj' : j' < qq ℓ K)
    (he : chunk D ℓ K i j = chunk D ℓ K i' j') : i = i' ∧ j = j' := by
  have h1 := chunk_head D ℓ K hK i hj
  rw [he, chunk_head D ℓ K hK i' hj'] at h1
  obtain ⟨rfl, h2⟩ := nm_inj D h1
  refine ⟨rfl, ?_⟩
  have a1 : j * K < qq ℓ K * K := Nat.mul_lt_mul_of_pos_right hj hK
  have a2 : j' * K < qq ℓ K * K := Nat.mul_lt_mul_of_pos_right hj' hK
  have : j * K = j' * K := by omega
  exact Nat.eq_of_mul_eq_mul_right hK this

lemma mem_axisMajor {c : List ℕ} :
    c ∈ axisMajor D ℓ K ↔ ∃ i j, j < qq ℓ K ∧ c = chunk D ℓ K i j := by
  simp [axisMajor, eq_comm]

lemma mem_chunkMajor {c : List ℕ} :
    c ∈ chunkMajor D ℓ K ↔ ∃ i j, j < qq ℓ K ∧ c = chunk D ℓ K i j := by
  simp [chunkMajor, eq_comm]
  constructor
  · rintro ⟨j, hj, i, h⟩; exact ⟨i, j, hj, h⟩
  · rintro ⟨i, j, hj, h⟩; exact ⟨j, hj, i, h⟩

lemma axisMajor_nodup (hK : 0 < K) : (axisMajor D ℓ K).Nodup := by
  unfold axisMajor
  rw [List.nodup_flatMap]
  refine ⟨fun i _ => List.nodup_range.map_on (fun j hj j' hj' e => by
    simp at hj hj'; exact (chunk_inj D ℓ K hK hj hj' e).2), ?_⟩
  refine (List.nodup_finRange D).pairwise_of_forall_ne fun i _ i' _ hii' => ?_
  rw [Function.onFun, List.disjoint_left]
  intro x hx hx'
  simp only [List.mem_map, List.mem_range] at hx hx'
  obtain ⟨j, hj, rfl⟩ := hx; obtain ⟨j', hj', he⟩ := hx'
  exact hii' (chunk_inj D ℓ K hK hj' hj he).1.symm

lemma chunkMajor_nodup (hK : 0 < K) : (chunkMajor D ℓ K).Nodup := by
  unfold chunkMajor
  rw [List.nodup_flatMap]
  refine ⟨fun j hj => (List.nodup_finRange D).map_on (fun i _ i' _ e => by
    simp at hj; exact (chunk_inj D ℓ K hK hj hj e).1), ?_⟩
  refine List.nodup_range.pairwise_of_forall_ne fun j hj j' hj' hjj' => ?_
  rw [Function.onFun, List.disjoint_left]
  intro x hx hx'
  simp only [List.mem_map] at hx hx'
  obtain ⟨i, -, rfl⟩ := hx; obtain ⟨i', -, he⟩ := hx'
  simp at hj hj'
  exact hjj' (chunk_inj D ℓ K hK hj' hj he).2.symm

lemma chunk_perm (hK : 0 < K) : (chunkMajor D ℓ K).Perm (axisMajor D ℓ K) := by
  rw [List.perm_ext_iff_of_nodup (chunkMajor_nodup D ℓ K hK) (axisMajor_nodup D ℓ K hK)]
  intro c; rw [mem_axisMajor, mem_chunkMajor]

/-! ## 6. Lemma 16, layout part -/

/-! ## 7. Round geometry -/

end Synth

end

section
-- Source module: Solutions.GlueLayout
/-!
# Glue 2a: the paper's layout as a `CoordLayout`, its selected bits, and the bridge to `layerSel`

* §0 bit arithmetic: `getBit_bitsum`, `bitsum_update` (binary digit sums), `getD_decode`
  (a digit of `decode`), `getD_decode_binSh` (bits of a binary address).
* §1 `PosData`/`posSel`: selected bits of a binary address given by bit positions;
  `selOp_transport`: two selected-bit structures related by an address map preserving the
  linear index give the same `selOp` for every tensor kernel.
* §2 `slotLayout`: a slot arrangement `arr` (a list of slot names `nm i h` covering every slot
  `h < a i` exactly once) is a `CoordLayout (binSh |arr|) a`; `slotSel` (round `h`, axes `P`)
  is `SelCompat` with it (`slot_compat`).
* §3 `chunk_bridge`: if the selected positions are `A + iK + (K-1-ρ)` in a binary address of
  length `A + DK + B`, the selected-bit operator *is* `layerSel (2^A) K D (2^B) ρ` of §5.
* §4 the paper's layout `tgtL` (`SynthLayoutPaper`): `paperLayout`, `paperSel`,
  **`synth_forward_paper`** (§6.1 on the concrete layout: the rounds implement `B F^-`), and
  **`paper_round_layer`** (for every slot `b ≤ h < b + qK` the round's simultaneous layer is
  exactly Proposition 14's `layerH (2^A) K D (2^B) ρ` with `ρ = (h-b) mod K`, via `tgt_slot`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace GlueLayout

open Finset Butterfly AModel Synth

/-! ## 0. Bit arithmetic -/

lemma getBit_mod_of_lt {t K : ℕ} (ht : t < K) (v : ℕ) : getBit t (v % 2 ^ K) = getBit t v := by
  ext; simp only [getBit]
  have hK : 2 ^ K = 2 ^ t * 2 ^ (K - t) := by rw [← pow_add]; congr 1; omega
  rw [hK, Nat.mod_mul_right_div_self]
  apply Nat.mod_mod_of_dvd
  exact dvd_pow_self 2 (by omega)

lemma bitsum_lt (c : ℕ → ℕ) (hc : ∀ t, c t < 2) :
    ∀ a, ∑ t ∈ range a, c t * 2 ^ t < 2 ^ a
  | 0 => by simp
  | a + 1 => by
    rw [sum_range_succ, pow_succ]
    have h1 := bitsum_lt c hc a
    have h2 : c a * 2 ^ a ≤ 1 * 2 ^ a := Nat.mul_le_mul_right _ (by have := hc a; omega)
    omega

/-- The bits of a binary digit sum are its digits. -/
lemma getBit_bitsum (c : ℕ → ℕ) (hc : ∀ t, c t < 2) :
    ∀ a, ∀ h < a, (getBit h (∑ t ∈ range a, c t * 2 ^ t) : ℕ) = c h
  | 0, h, hh => by omega
  | a + 1, h, hh => by
    rw [sum_range_succ]
    have hlt := bitsum_lt c hc a
    rcases Nat.lt_or_ge h a with ha | ha
    · rw [← getBit_mod_of_lt ha, Nat.add_mul_mod_self_right, Nat.mod_eq_of_lt hlt]
      exact getBit_bitsum c hc a h ha
    · have hha : h = a := by omega
      subst hha
      simp only [getBit]
      rw [Nat.add_mul_div_right _ _ (by positivity), Nat.div_eq_of_lt hlt, zero_add,
        Nat.mod_eq_of_lt (hc h)]

lemma valid_binSh_iff {n : ℕ} {v : List ℕ} :
    Valid (binSh n) v ↔ v.length = n ∧ ∀ k < n, v.getD k 0 < 2 := by
  rw [valid_iff]
  simp only [binSh, List.length_replicate, List.getElem_replicate]
  constructor
  · rintro ⟨hl, h⟩
    refine ⟨hl, fun k hk => ?_⟩
    rw [List.getD_eq_getElem _ _ (by omega)]; exact h k (by omega) hk
  · rintro ⟨hl, h⟩
    refine ⟨hl, fun k h1 h2 => ?_⟩
    have := h k h2; rwa [List.getD_eq_getElem _ _ h1] at this

lemma valid_getD_lt_two {n : ℕ} {v : List ℕ} (hv : Valid (binSh n) v) (k : ℕ) : v.getD k 0 < 2 := by
  obtain ⟨hl, h⟩ := valid_binSh_iff.1 hv
  by_cases hk : k < n
  · exact h k hk
  · rw [List.getD_eq_default _ _ (by omega)]; norm_num

/-! ## 1. Selected bits given by positions; transport -/

/-- Selected bits of a binary address of length `n` at positions `pos i`, with the partial
inverse `inv`. -/
structure PosData (n : ℕ) (ι : Type*) where
  pos : ι → ℕ
  inv : ℕ → Option ι
  pos_lt : ∀ i, pos i < n
  inv_iff : ∀ k i, inv k = some i ↔ pos i = k

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

/-- The new value at position `k`. -/
def PosData.val (P : PosData n ι) (v : List ℕ) (y : ι → Fin 2) (k : ℕ) : ℕ :=
  match P.inv k with
  | some i => (y i : ℕ)
  | none => v.getD k 0

/-- Overwrite the selected positions. -/
def PosData.setL (P : PosData n ι) (v : List ℕ) (y : ι → Fin 2) : List ℕ :=
  (List.range v.length).map (P.val v y)

lemma PosData.length_setL (P : PosData n ι) (v : List ℕ) (y : ι → Fin 2) :
    (P.setL v y).length = v.length := by simp [PosData.setL]

lemma PosData.getD_setL (P : PosData n ι) (v : List ℕ) (y : ι → Fin 2) {k : ℕ}
    (hk : k < v.length) : (P.setL v y).getD k 0 = P.val v y k := by
  rw [List.getD_eq_getElem _ _ (by simpa [PosData.setL] using hk)]; simp [PosData.setL]

lemma PosData.val_pos (P : PosData n ι) (v : List ℕ) (y : ι → Fin 2) (i : ι) :
    P.val v y (P.pos i) = y i := by
  have := (P.inv_iff (P.pos i) i).2 rfl
  simp [PosData.val, this]

lemma PosData.val_lt_two (P : PosData n ι) {v : List ℕ} (hv : Valid (binSh n) v)
    (y : ι → Fin 2) (k : ℕ) : P.val v y k < 2 := by
  unfold PosData.val
  split
  · exact (y _).isLt
  · exact valid_getD_lt_two hv k

/-- The selected-bit structure of `PosData`. -/
def posSel (P : PosData n ι) : SelBits (binSh n) ι where
  bits v i := ⟨v.getD (P.pos i) 0 % 2, Nat.mod_lt _ (by norm_num)⟩
  set := P.setL
  valid_set v y hv := by
    rw [valid_binSh_iff] at hv ⊢
    refine ⟨by rw [P.length_setL, hv.1], fun k hk => ?_⟩
    rw [P.getD_setL v y (by omega)]
    exact P.val_lt_two (valid_binSh_iff.2 hv) y k
  bits_set v y hv := by
    funext i
    have hl := (valid_binSh_iff.1 hv).1
    apply Fin.ext
    simp only
    rw [P.getD_setL v y (by rw [hl]; exact P.pos_lt i), P.val_pos, Nat.mod_eq_of_lt (y i).isLt]
  set_set v y z hv := by
    have hl := (valid_binSh_iff.1 hv).1
    apply List.ext_getElem (by simp [P.length_setL])
    intro k h1 h2
    simp only [PosData.setL, List.getElem_map, List.getElem_range, List.length_map,
      List.length_range] at h1 h2 ⊢
    unfold PosData.val
    split
    · rfl
    · rw [List.getD_eq_getElem _ _ (by simpa using h1)]
      simp only [List.getElem_map, List.getElem_range, PosData.val]
      split
      · next i hi _ j hj => rw [hi] at hj; cases hj
      · rfl
  set_bits v hv := by
    have hl := (valid_binSh_iff.1 hv).1
    apply List.ext_getElem (by simp [P.length_setL])
    intro k h1 h2
    simp only [PosData.setL, List.getElem_map, List.getElem_range]
    unfold PosData.val
    split
    · next i hi =>
      have hk := (P.inv_iff k i).1 hi
      subst hk
      simp only
      rw [Nat.mod_eq_of_lt (valid_getD_lt_two hv _), List.getD_eq_getElem _ _ h2]
    · rw [List.getD_eq_getElem _ _ h2]

/-! ## 2. Slot arrangements are coordinate layouts -/

section Slots

variable (D : ℕ) [NeZero D]

/-- The axis of a slot name `nm i h = hD + i`. -/
def axOf (s : ℕ) : Fin D := ⟨s % D, Nat.mod_lt _ (NeZero.pos D)⟩

lemma axOf_nm (i : Fin D) (h : ℕ) : axOf D (nm D i h) = i := by
  ext; simp only [axOf, nm]
  rw [Nat.add_comm, Nat.add_mul_mod_self_right, Nat.mod_eq_of_lt i.isLt]

lemma nm_div (i : Fin D) (h : ℕ) : nm D i h / D = h := by
  unfold nm
  rw [Nat.add_comm, Nat.add_mul_div_right _ _ (NeZero.pos D), Nat.div_eq_of_lt i.isLt, zero_add]

lemma nm_axOf (s : ℕ) : nm D (axOf D s) (s / D) = s := by
  unfold nm axOf; simp only; exact Nat.div_add_mod' s D

variable (a : Fin D → ℕ) (arr : List ℕ)

/-- A slot arrangement: every slot `(i, h)`, `h < a i`, named `nm i h`, exactly once. -/
structure SlotArr : Prop where
  nodup : arr.Nodup
  mem_iff : ∀ s, s ∈ arr ↔ ∃ i : Fin D, ∃ h, h < a i ∧ s = nm D i h

/-- The bit position of slot `(i, h)`. -/
def spos (i : Fin D) (h : ℕ) : ℕ := arr.idxOf (nm D i h)

variable {D a arr}

lemma SlotArr.mem (hS : SlotArr D a arr) {i : Fin D} {h : ℕ} (hh : h < a i) : nm D i h ∈ arr :=
  (hS.mem_iff _).2 ⟨i, h, hh, rfl⟩

lemma SlotArr.spos_lt (hS : SlotArr D a arr) {i : Fin D} {h : ℕ} (hh : h < a i) :
    spos D arr i h < arr.length :=
  List.idxOf_lt_length_of_mem (hS.mem hh)

lemma SlotArr.getD_spos (hS : SlotArr D a arr) {i : Fin D} {h : ℕ} (hh : h < a i) :
    arr.getD (spos D arr i h) 0 = nm D i h := by
  rw [List.getD_eq_getElem _ _ (hS.spos_lt hh)]
  exact List.getElem_idxOf (hS.spos_lt hh)

lemma SlotArr.slot_lt (hS : SlotArr D a arr) {k : ℕ} (hk : k < arr.length) :
    arr.getD k 0 / D < a (axOf D (arr.getD k 0)) := by
  have hm : arr.getD k 0 ∈ arr := by
    rw [List.getD_eq_getElem _ _ hk]; exact List.getElem_mem hk
  obtain ⟨i, h, hh, he⟩ := (hS.mem_iff _).1 hm
  rw [he, axOf_nm, nm_div]; exact hh

lemma SlotArr.spos_getD (hS : SlotArr D a arr) {k : ℕ} (hk : k < arr.length) :
    spos D arr (axOf D (arr.getD k 0)) (arr.getD k 0 / D) = k := by
  unfold spos
  rw [nm_axOf, List.getD_eq_getElem _ _ hk]
  exact hS.nodup.idxOf_getElem k hk

variable (D a arr)

/-- Axis coordinates of a binary address. -/
def coordS (v : List ℕ) (i : Fin D) : ℕ := ∑ t ∈ range (a i), v.getD (spos D arr i t) 0 * 2 ^ t

/-- The binary address of a box point. -/
def addrS (x : Fin D → ℕ) : List ℕ := arr.map fun s => (getBit (s / D) (x (axOf D s)) : ℕ)

variable {D a arr}

lemma getD_addrS (x : Fin D → ℕ) {k : ℕ} (hk : k < arr.length) :
    (addrS D arr x).getD k 0 = (getBit (arr.getD k 0 / D) (x (axOf D (arr.getD k 0))) : ℕ) := by
  rw [List.getD_eq_getElem _ _ (by simpa [addrS] using hk), List.getD_eq_getElem _ _ hk]
  simp [addrS]

lemma coordS_lt {v : List ℕ} (hv : Valid (binSh arr.length) v) (i : Fin D) :
    coordS D a arr v i < 2 ^ a i :=
  bitsum_lt _ (fun t => valid_getD_lt_two hv _) _

/-- **A slot arrangement is a coordinate layout** of the binary shape. -/
def slotLayout (hS : SlotArr D a arr) : CoordLayout (binSh arr.length) a where
  addr := addrS D arr
  coord := coordS D a arr
  valid_addr x _ := valid_binSh_iff.2 ⟨by simp [addrS], fun k hk => by
    rw [getD_addrS x hk]; exact (getBit _ _).isLt⟩
  coord_addr x hx := by
    funext i
    unfold coordS
    have hxi := (mem_box.1 hx) i
    rw [← Nat.mod_eq_of_lt hxi, ← sum_bits_eq_mod]
    apply sum_congr rfl
    intro t ht
    have ht' := mem_range.1 ht
    rw [getD_addrS x (hS.spos_lt ht'), hS.getD_spos ht', nm_div, axOf_nm]
  addr_coord v hv := by
    have hl := (valid_binSh_iff.1 hv).1
    apply List.ext_getElem (by simp [addrS, hl])
    intro k h1 h2
    have hk : k < arr.length := by simpa [addrS] using h1
    have e := getD_addrS (coordS D a arr v) hk
    rw [List.getD_eq_getElem _ _ h1] at e
    rw [e]
    unfold coordS
    rw [getBit_bitsum _ (fun t => valid_getD_lt_two hv _) _ _ (hS.slot_lt hk), hS.spos_getD hk,
      List.getD_eq_getElem _ _ h2]
  coord_mem v hv := mem_box.2 fun i => coordS_lt hv i

/-- The selected positions of round `h` on the axes `P` (each with `h < a i`). -/
def slotPos (hS : SlotArr D a arr) (h : ℕ) (P : Finset (Fin D)) (hP : ∀ i ∈ P, h < a i) :
    PosData arr.length {i // i ∈ P} where
  pos i := spos D arr i.1 h
  inv k := if hk : k < arr.length ∧ arr.getD k 0 / D = h ∧ axOf D (arr.getD k 0) ∈ P then
      some ⟨axOf D (arr.getD k 0), hk.2.2⟩ else none
  pos_lt i := hS.spos_lt (hP i.1 i.2)
  inv_iff k i := by
    constructor
    · intro hk
      split_ifs at hk with hc
      · cases hk
        obtain ⟨hk1, hk2, -⟩ := hc
        conv_rhs => rw [← hS.spos_getD hk1]
        simp only [hk2]
    · intro hk
      have hlt := hS.spos_lt (hP i.1 i.2)
      have hg := hS.getD_spos (hP i.1 i.2)
      rw [hk] at hlt hg
      have hc : k < arr.length ∧ arr.getD k 0 / D = h ∧ axOf D (arr.getD k 0) ∈ P := by
        rw [hg, nm_div, axOf_nm]; exact ⟨hlt, rfl, i.2⟩
      rw [dif_pos hc]
      congr 1; apply Subtype.ext; simp only; rw [hg, axOf_nm]

/-- The selected bits of round `h` on the axes `P` in a slot arrangement. -/
def slotSel (hS : SlotArr D a arr) (h : ℕ) (P : Finset (Fin D)) (hP : ∀ i ∈ P, h < a i) :
    SelBits (binSh arr.length) {i // i ∈ P} :=
  posSel (slotPos hS h P hP)

lemma partAxes_lt {h : ℕ} : ∀ i ∈ partAxes a h, h < a i := fun _ hi => (Finset.mem_filter.1 hi).2

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ)

end Slots

/-! ## 3. Chunk positions are Proposition 14's `layerSel` -/

section Chunk

variable {A K D B ρ : ℕ}

end Chunk

/-! ## 4. The paper's layout `tgtL` -/

section Paper

variable (D ℓ K : ℕ) [NeZero D] (long : Fin D → Bool)

lemma mem_initL {s : ℕ} :
    s ∈ initL D ℓ long ↔ ∃ i : Fin D, ∃ h, h < al D ℓ long i ∧ s = nm D i h := by
  unfold initL
  simp only [List.mem_flatMap, List.mem_finRange, true_and, List.mem_map, List.mem_reverse,
    List.mem_range]
  constructor
  · rintro ⟨i, h, hh, rfl⟩; exact ⟨i, h, hh, rfl⟩
  · rintro ⟨i, h, hh, rfl⟩; exact ⟨i, h, hh, rfl⟩

/-- The target order is a rearrangement of the initial one. -/
theorem tgtL_perm (hK : 0 < K) (hℓ : 1 ≤ ℓ) : (tgtL D ℓ K long).Perm (initL D ℓ long) := by
  have h1 : (lstL D ℓ K long).Perm ((initL D ℓ long).filter fun x => decide (x ∈ lstL D ℓ K long)) := by
    rw [List.perm_ext_iff_of_nodup (lstL_nodup D ℓ K long) ((initL_nodup D ℓ long).filter _)]
    intro x; simp only [List.mem_filter, decide_eq_true_eq]
    exact ⟨fun h => ⟨lstL_sub_initL D ℓ K long hℓ x h, h⟩, fun h => h.2⟩
  have h2 : (chunkMajor D ℓ K).flatten.Perm
      ((initL D ℓ long).filter fun x => !decide (x ∈ lstL D ℓ K long)) := by
    have : ((initL D ℓ long).filter fun x => !decide (x ∈ lstL D ℓ K long)) =
        (initL D ℓ long).filter (· ∉ lstL D ℓ K long) := by
      apply List.filter_congr; intro x _; simp
    rw [this, initL_filter D ℓ K long hK hℓ]
    exact (chunk_perm D ℓ K hK).flatten
  exact (h1.append h2).trans (List.filter_append_perm _ _)

/-- The paper's layout is a slot arrangement. -/
theorem paperArr (hK : 0 < K) (hℓ : 1 ≤ ℓ) : SlotArr D (al D ℓ long) (tgtL D ℓ K long) where
  nodup := (tgtL_perm D ℓ K long hK hℓ).nodup_iff.2 (initL_nodup D ℓ long)
  mem_iff s := by rw [(tgtL_perm D ℓ K long hK hℓ).mem_iff, mem_initL]

/-- **The coordinate layout `L` of Lemma 16** (addresses of `|tgtL|` bits, slot order `tgtL`). -/
def paperLayout (hK : 0 < K) (hℓ : 1 ≤ ℓ) :
    CoordLayout (binSh (tgtL D ℓ K long).length) (al D ℓ long) :=
  slotLayout (paperArr D ℓ K long hK hℓ)

/-- The selected bits of round `h` in the paper's layout (axes with `a_i > h`). -/
def paperSel (hK : 0 < K) (hℓ : 1 ≤ ℓ) (h : ℕ) :
    SelBits (binSh (tgtL D ℓ K long).length) {i // i ∈ partAxes (al D ℓ long) h} :=
  slotSel (paperArr D ℓ K long hK hℓ) h (partAxes (al D ℓ long) h) partAxes_lt

variable {R : Type*} [CommRing R] [Algebra ℂ R] (y : Rˣ)

/-! ### The chunk rounds are Proposition 14's layer -/

end Paper

end GlueLayout

end

section
-- Source module: Solutions.P8LiftIface
/-!
# A restricted Proposition 8 interface (only the paper's cheap local operations)

`Butterfly.Prop8Iface` allows an arbitrary matrix as a non-recursive step, which makes it
vacuous. Here the non-recursive steps are restricted to the operations of §3.5/§5.1, each
of which touches every state entry a bounded number of times:

* `gate G` — a pointwise scalar gate across roles at equal addresses (`G` on the roles,
  identity on the addresses);
* `rowAdd w i j` — a binary row addition `x(j,c) += x(i,c)` (all columns `c`) on the address
  bits of role `w`, i.e. an address permutation;
* `phase w k₀ k₁ a` — the diagonal phase `i^{k₀ + k₁|κ| + Σ_{(r,c)} a_r x(r,c)}` on role `w`
  (covers `Z_u` with `a = 2u`, `S` on selected bits, scalars `i^{27|κ|}`);
* `perm σ` — a permutation of the roles;
* `child w h fwd` — a recursive call on slot `h` of role `w` (forward or the (17) inverse).

The step list is independent of the column type `κ` (the same program for every `κ`).

* `Prop8IfaceR.toIface`: the forgetful map to `Butterfly.Prop8Iface` (so `recOp_eq`,
  `recCost_le`, `prop14_output` apply).
* `den_row_card_le`: every local step's matrix has at most `|R|` nonzero entries per row
  (`≤ 1` except for gates), so one local step costs `≤ W · V` entry operations.
* `recCostR_le`: the §5.3 bound for the recursion whose local work is the actual work
  `≤ nloc · W · V` of the local steps (`nloc` = number of local steps, independent of `κ`).
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace P8Lift

open Matrix Finset Complex Butterfly

/-- Group addresses: slot `r : Fin m`, column `c : κ`. -/
abbrev Addr (m : ℕ) (κ : Type) := Fin m × κ → Fin 2

/-- Operators on roles `R` of functions of the group address. -/
abbrev ROp (R : Type) (m : ℕ) (κ : Type) := Matrix (R × Addr m κ) (R × Addr m κ) ℂ

/-- The binary row addition `x(j,c) += x(i,c)` on every column. -/
def rowAddr {m : ℕ} {κ : Type} (i j : Fin m) (x : Addr m κ) : Addr m κ :=
  fun p => if p.1 = j then x p + x (i, p.2) else x p

lemma rowAddr_invol {m : ℕ} {κ : Type} {i j : Fin m} (hij : i ≠ j) (x : Addr m κ) :
    rowAddr i j (rowAddr i j x) = x := by
  funext p
  simp only [rowAddr]
  by_cases hp : p.1 = j
  · simp only [hp, if_true, if_neg hij]
    generalize x p = a; generalize x (i, p.2) = b
    revert a b; decide
  · simp [hp]

/-- The restricted network steps (independent of the column type). -/
inductive RStep (R : Type) (m : ℕ) where
  | gate (G : Matrix R R ℂ)
  | rowAdd (w : R) (i j : Fin m) (hij : i ≠ j)
  | phase (w : R) (k₀ k₁ : ℕ) (a : Fin m → ℕ)
  | perm (σ : Equiv.Perm R)
  | child (w : R) (h : Fin m) (fwd : Bool)

namespace RStep

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

/-- Whether a step is a child call. -/
def isChild : RStep R m → Bool
  | child .. => true
  | _ => false

/-- The phase exponent `k₀ + k₁|κ| + Σ_{(r,c)} a_r x(r,c)`. -/
def phExp (k₀ k₁ : ℕ) (a : Fin m → ℕ) (x : Addr m κ) : ℕ :=
  k₀ + k₁ * Fintype.card κ + ∑ p : Fin m × κ, a p.1 * (x p).val

/-- Child call `T` on slot `h` of role `w` (as `Butterfly.childMat`, for any role type). -/
def childR (w : R) (h : Fin m) (T : Op κ) : ROp R m κ := fun p q =>
  if p.1 = q.1 then
    (if p.1 = w then
      (if ∀ h', h' ≠ h → ∀ j, p.2 (h', j) = q.2 (h', j) then T (slotBits p.2 h) (slotBits q.2 h)
        else 0)
     else if p.2 = q.2 then 1 else 0)
  else 0

/-- Matrices of the steps (`T` = operator computed by a forward child call). -/
noncomputable def den (T : Op κ) : RStep R m → ROp R m κ
  | gate G => fun p q => G p.1 q.1 * (if p.2 = q.2 then 1 else 0)
  | rowAdd w i j _ => fun p q => if p.1 = q.1 then
      (if p.1 = w then (if q.2 = rowAddr i j p.2 then 1 else 0) else if p.2 = q.2 then 1 else 0)
      else 0
  | phase w k₀ k₁ a => diagonal fun p => if p.1 = w then I ^ phExp k₀ k₁ a p.2 else 1
  | perm σ => fun p q => if p.1 = σ q.1 ∧ p.2 = q.2 then 1 else 0
  | child w h true => childR w h T
  | child w h false => childR w h ((-I) ^ Fintype.card κ • (ZD * T * ZD))

end RStep

open RStep

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

/-- Denotation of a list (first step applied first). -/
noncomputable def denListR (T : Op κ) (l : List (RStep R m)) : ROp R m κ :=
  l.foldl (fun acc st => st.den T * acc) 1

lemma foldl_den_mul (T : Op κ) (l : List (RStep R m)) (A : ROp R m κ) :
    l.foldl (fun acc st => st.den T * acc) A = denListR T l * A := by
  induction l generalizing A with
  | nil => simp [denListR]
  | cons s l ih =>
    simp only [denListR, List.foldl_cons]
    rw [ih, ih (s.den T * 1), mul_one, mul_assoc]

@[simp] lemma denListR_nil (T : Op κ) : denListR T ([] : List (RStep R m)) = 1 := rfl

lemma denListR_append (T : Op κ) (l₁ l₂ : List (RStep R m)) :
    denListR T (l₁ ++ l₂) = denListR T l₂ * denListR T l₁ := by
  simp only [denListR, List.foldl_append]
  exact foldl_den_mul _ _ _

/-- Forget the restriction: local steps become `NStep.loc` of their matrix. -/
noncomputable def RStep.toN {W : ℕ} : RStep (Fin W) m → NStep W m κ
  | child w h fwd => NStep.child w h fwd
  | st => NStep.loc (st.den (Butterfly.tens Butterfly.Cm))

lemma RStep.den_toN {W : ℕ} (st : RStep (Fin W) m) :
    (st.toN (κ := κ)).den (Butterfly.tens Butterfly.Cm) = st.den (Butterfly.tens Butterfly.Cm) := by
  cases st with
  | child w h fwd => cases fwd <;> rfl
  | _ => rfl

lemma RStep.isChild_toN {W : ℕ} (st : RStep (Fin W) m) :
    (st.toN (κ := κ)).isChild = st.isChild := by
  cases st <;> rfl

lemma denList_map_toN {W : ℕ} (l : List (RStep (Fin W) m)) :
    denList (Butterfly.tens Butterfly.Cm) (l.map (RStep.toN (κ := κ))) = denListR (Butterfly.tens Butterfly.Cm) l := by
  induction l using List.reverseRecOn with
  | nil => rfl
  | append_singleton l s ih =>
    rw [List.map_append, List.map_singleton, denListR_append]
    simp only [denList, List.foldl_append, List.foldl_cons, List.foldl_nil] at ih ⊢
    rw [ih, RStep.den_toN]
    simp [denListR]

/-- **The restricted Proposition 8 interface**: one `κ`-independent program of cheap steps
with exactly `s` child calls computing `C^{⊗(Fin m × κ)}` on every role for every `κ`. -/
structure Prop8IfaceR (W m s : ℕ) where
  net : List (RStep (Fin W) m)
  count : net.countP RStep.isChild = s
  correct : ∀ (κ : Type) [Fintype κ] [DecidableEq κ],
    denListR (Butterfly.tens Butterfly.Cm) net = (roleDiag (Butterfly.tens Butterfly.Cm) : NOp W (Fin m × κ))

namespace Prop8IfaceR

variable {W s : ℕ}

/-- **Forgetful map** to the (unrestricted) interface of `ButterflyProp14`. -/
noncomputable def toIface (J : Prop8IfaceR W m s) : Prop8Iface W m s where
  net κ _ _ := J.net.map RStep.toN
  count κ _ _ := by
    rw [List.countP_map]
    have : (NStep.isChild ∘ (RStep.toN (κ := κ) (W := W) (m := m))) = RStep.isChild := by
      funext st; exact RStep.isChild_toN st
    rw [this]; exact J.count
  correct κ _ _ := by rw [denList_map_toN]; exact J.correct κ

/-- Number of local steps (independent of `κ`). -/
def nloc (J : Prop8IfaceR W m s) : ℕ := J.net.countP fun st => !st.isChild

variable [NeZero W]

end Prop8IfaceR

end P8Lift

end