κ-improvement development, part 15: item-language primitives

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

import Definitions.Def_IntMul_Kappa_K_Defs14

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFGPair
/-!
# KFG: a pointwise `2 × 2` scaled operation on two arbitrary role streams

On the data stack `pre ++ u ++ gap ++ v ++ rest` (Gaussian streams, `|u| = |v| = L`),
`pairP M` replaces `(u, v)` by `(M₀₀ u + M₀₁ v, M₁₀ u + M₁₁ v)` entrywise, with scaled
coefficients `c · 2^d`.  The prefix and the gap are parked on `qP`, `qQ`; the two streams are
distributed to the four register stacks and zipped (`zipRS`); the outputs are moved back.  Every
entry is touched `O(1)` times; the loop counters are unary (`fcode`).

* `pairP_fsp`: the program with its counts in registers (`rB = L`, `rB2 = 2L`, `rP = 2|pre|`,
  `rQ = 2|gap|`).
* `pairW M L i j`: the self-contained wrapper on roles `i < j` of a role-major array with role
  length `L` (`pairW_fsp`, list semantics `pairL`, `getD_pairL`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

abbrev qQ : Fin HS := 41
abbrev rQ : Fin HS := 54

/-- The pair pass (counts in registers). -/
def pairP (M : M2S) : RS :=
  .seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.reg NLogN.KFB.rP)) (.seq (NLogN.mvK 0 qT NLogN.KFB.c1 (.reg rB2)) (.seq (distB rUr rUi)
  (.seq (NLogN.mvK 0 qQ NLogN.KFB.c1 (.reg rQ)) (.seq (NLogN.mvK 0 qT NLogN.KFB.c1 (.reg rB2)) (.seq (distB rVr rVi)
  (.seq (zipRS M) (.seq (NLogN.mvK qO 0 NLogN.KFB.c1 (.reg rB2)) (.seq (NLogN.mvK qQ 0 NLogN.KFB.c1 (.reg rQ))
  (.seq (NLogN.mvK qA 0 NLogN.KFB.c1 (.reg rB2)) (NLogN.mvK qP 0 NLogN.KFB.c1 (.reg NLogN.KFB.rP)))))))))))

/-- Cost of the pair pass. -/
def cPair (w D L P Q : ℕ) : ℕ :=
  4 * (P + Q) * (250 * w + 2600) +
    L * (cZS w D + 2 * cD w + 8 * (250 * (w + D + 2) + 2600) + 10000) + 30000

/-- The registers of the pair pass hold the counts. -/
structure PCnt (s : IState HS) (f : Fin HS → List ℤ) (L P Q : ℕ) : Prop where
  b : ∃ R, (encK s f).G NLogN.KFB.rB = .int (L : ℤ) :: R
  b2 : ∃ R, (encK s f).G rB2 = .int ((2 * L : ℕ) : ℤ) :: R
  p : ∃ R, (encK s f).G NLogN.KFB.rP = .int ((2 * P : ℕ) : ℤ) :: R
  q : ∃ R, (encK s f).G rQ = .int ((2 * Q : ℕ) : ℤ) :: R

end NLogN.KFG

end

section
-- Source module: Solutions.KFGPairW
/-!
# KFG: the self-contained pair pass on roles `i < j`, and the scaled role encoding

* `pairW M L i j`: pushes the counts (`L`, `2L`, `2iL`, `2(j-i-1)L`), runs `pairP M`, pops.
  `pairW_fsp`: on a clean managed state, stack `0` holding a role-major array `zs` (role length
  `L`), the result is `pairL M L i j zs`; cost `cPairW`.  `fok_pairW`: it compiles finely.
* `getD_pairL`: index semantics.
* `Enc W L e zs y`: `zs` encodes the role streams `y w p` (`w < W`, `p < L`) with per-role scale
  `2^{-e w}`: `toC zs[w L + p] = 2^{e w} y w p`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-! ## Compilability -/

/-! ## The wrapper -/

/-- The pair pass on roles `i < j` of a role-major array with role length `L`. -/
def pairW (M : M2S) (L i j : ℕ) : RS :=
  .seq (.push NLogN.KFB.rB (.lit (L : ℤ))) (.seq (.push rB2 (.lit ((2 * L : ℕ) : ℤ)))
  (.seq (.push NLogN.KFB.rP (.lit ((2 * (i * L) : ℕ) : ℤ))) (.seq (.push rQ (.lit ((2 * ((j - i - 1) * L) : ℕ) : ℤ)))
  (.seq (pairP M) (.seq (.pop NLogN.KFB.rB) (.seq (.pop rB2) (.seq (.pop NLogN.KFB.rP) (.pop rQ))))))))

/-- Clean scratch of the pair pass. -/
structure ClP (f : Fin HS → List ℤ) : Prop where
  scr : Scr f
  a : f qA = []
  p : f qP = []
  q : f qQ = []
  b : f NLogN.KFB.rB = []
  b2 : f rB2 = []
  rp : f NLogN.KFB.rP = []
  rq : f rQ = []

/-- List semantics of the pair pass on roles `i < j`. -/
def pairL (M : M2S) (L i j : ℕ) (zs : List GI) : List GI :=
  zs.take (i * L) ++ (List.zipWith (lin2S M.a M.b) ((zs.drop (i * L)).take L) ((zs.drop (j * L)).take L) ++
    ((zs.drop ((i + 1) * L)).take ((j - i - 1) * L) ++
      (List.zipWith (lin2S M.c M.d) ((zs.drop (i * L)).take L) ((zs.drop (j * L)).take L) ++
        zs.drop ((j + 1) * L))))

/-- Cost of the wrapper. -/
def cPairW (w D L i j Z : ℕ) : ℕ :=
  cPair w D L (i * L) ((j - i - 1) * L) + 4 * (4000 * (Z + 2) + 2 * Z + 2)

end NLogN.KFG

end

section
-- Source module: Solutions.KFGEnc
/-!
# KFG: index semantics of the pair pass, and the scaled role encoding

* `getD_pairL`: entry `p` of `pairL M L i j zs`.
* `Enc W L e zs y`: the Gaussian array `zs` (role-major, role length `L`) encodes the complex
  role streams `y w p` with per-role scales: `toC zs[w L + p] = 2^{e w} · y w p`.
* `enc_pairL`: the pair pass on the encoding, for any new scales `e'` and streams `y'` that
  satisfy the two pointwise identities.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-! ## The scaled encoding -/

/-- `zs` encodes the role streams `y` with per-role scales `e`. -/
def Enc (W L : ℕ) (e : ℕ → ℕ) (zs : List GI) (y : ℕ → ℕ → ℂ) : Prop :=
  zs.length = W * L ∧ ∀ w < W, ∀ p < L, NLogN.KFB.toC (zs.getD (w * L + p) 0) = 2 ^ (e w) * y w p

end NLogN.KFG

end

section
-- Source module: Solutions.KFGTab
/-!
# KFG: the phase table by doubling (task B, part 1)

The phase `i^{u + Σ_t aE t · bit_t q}` of entry `q < 2^K` is produced as a table of exponents
`mod 4` on stack `xA`, built by `K` doubling steps: `T ↦ T ++ map (· + aE t mod 4) T`.  The
doubling step on a table of length `M` costs `O(M) + O(Z)`, so the whole table costs
`O(2^K) + O(K Z)` (geometric sum), i.e. `O(1)` per table entry.

* `tbl u aE t`, `getD_tbl`: the table and its entries.
* `dblP ε`, `dbl_fsp`: one doubling step (fine, unary loop counter `c1`).
* `tblP aE K`, `tblP_fsp`: all `K` steps, cost `cTbl K Z = 120000 · 2^K + K · O(Z)`.
* `rot e z = z · i^e` by three selections (`reE`, `imE`), `mulB`, `mulLoop_fsp`: every entry
  rotated by its table entry (`O(w)` bit operations per entry).
* `phT K off u aE`, `phT_fsp`: the whole phase on the `2^K` entries at offset `off`
  (park prefix, table, multiply, drop table, unpark); cost `cPhT w K off Z =
  O(off · w + 2^K · w) + cTbl K Z + O(Z)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

abbrev xA : Fin HS := 42
abbrev xB : Fin HS := 43
abbrev xC : Fin HS := 44
abbrev xD : Fin HS := 45
abbrev xN : Fin HS := 55

/-! ## The table on the list level -/

/-- One doubling of the exponent table. -/
def dblL (ε : ℤ) (T : List ℤ) : List ℤ := T ++ T.map (fun x => (x + ε) % 4)

/-- The exponent table after `t` doublings. -/
def tbl (u : ℕ) (aE : ℕ → ℕ) : ℕ → List ℤ
  | 0 => [((u % 4 : ℕ) : ℤ)]
  | t + 1 => dblL ((aE t % 4 : ℕ) : ℤ) (tbl u aE t)

/-! ## The doubling step -/

/-- New entry `(b + ε) mod 4` from the top of `xB`. -/
def dE (ε : ℤ) : Ex HS := .rem (.add (.reg xB) (.lit ε)) (.lit 4)

/-- Body of the doubling loop. -/
def dblB (ε : ℤ) : RS := .seq (.push xA (NLogN.KFG.dE ε)) (.seq (.push xC (.reg xB)) (.pop xB))

/-- One doubling of the table on `xA` (length in `xN`; doubles `xN`). -/
def dblP (ε : ℤ) : RS :=
  .seq (NLogN.mvK xA xB NLogN.KFB.c1 (.reg xN))
  (.seq (.rep NLogN.KFB.c1 (.reg xN) (dblB ε))
  (.seq (NLogN.mvK xC xD NLogN.KFB.c1 (.reg xN)) (.seq (NLogN.mvK xD xA NLogN.KFB.c1 (.reg xN))
  (.asg xN (.add (.reg xN) (.reg xN))))))

/-- Entries in `[0, 4)`. -/
def Ex4 (T : List ℤ) : Prop := ∀ a ∈ T, 0 ≤ a ∧ a < 4

/-- Cost of one doubling of a table of length `M`. -/
def cDbl (M Z : ℕ) : ℕ := 120000 * M + 10000 * (Z + 2) + 20000

/-! ## All doublings -/

/-- `K` doublings, with the bit exponents `aE t`. -/
def tblP (aE : ℕ → ℕ) : ℕ → RS
  | 0 => .skip
  | t + 1 => .seq (tblP aE t) (dblP ((aE t % 4 : ℕ) : ℤ))

/-- Cost of the table. -/
def cTbl (t Z : ℕ) : ℕ := 120000 * 2 ^ t + t * (10000 * (Z + 2) + 20001)

/-! ## Rotation of an element by `i^e` -/

/-- `z · i^e` for `e ∈ [0, 4)`, as the selections compute it. -/
def rot (e : ℤ) (z : GI) : GI :=
  (if e < 1 then z.1 else if e < 2 then -z.2 else if e < 3 then -z.1 else z.2,
   if e < 1 then z.2 else if e < 2 then z.1 else if e < 3 then -z.2 else -z.1)

/-- `e - k` for the exponent `e` on top of `xA`. -/
def ckE (k : ℤ) : Ex HS := .add (.reg xA) (.lit (-k))

/-- Real part of `(x + i y) i^e`. -/
def reE : Ex HS := .sel (ckE 1) (.reg rUr) (.sel (ckE 2) (.neg (.reg rUi))
  (.sel (ckE 3) (.neg (.reg rUr)) (.reg rUi)))

/-- Imaginary part of `(x + i y) i^e`. -/
def imE : Ex HS := .sel (ckE 1) (.reg rUi) (.sel (ckE 2) (.reg rUr)
  (.sel (ckE 3) (.neg (.reg rUi)) (.neg (.reg rUr))))

section
variable {s : IState HS} {f : Fin HS → List ℤ} {e x y : ℤ} {A X Y : List ℤ}

end

/-- The multiply body: one element of stack `0` times `i^e`, `e` moved from `xA` to `xB`,
results pushed on `qO`. -/
def mulB : RS :=
  .seq (.push rUr (.reg 0)) (.seq (.pop 0) (.seq (.push rUi (.reg 0)) (.seq (.pop 0)
  (.seq (.push qO NLogN.KFG.reE) (.seq (.push qO NLogN.KFG.imE) (.seq (.pop rUr) (.seq (.pop rUi)
  (.seq (.push xB (.reg xA)) (.pop xA)))))))))

/-- Cost of one element. -/
def cMulB (w : ℕ) : ℕ := 200000 * (w + 5)

/-! ## The multiply loop -/

/-- Entrywise rotation of a stream by a table. -/
def mulL (T : List ℤ) (zs : List GI) : List GI := List.zipWith rot T zs

/-! ## The phase program on one segment -/

/-- Phase `i^{u + Σ_t aE t · bit_t q}` on the `2^K` entries at offset `off` of the data stack. -/
def phT (K off u : ℕ) (aE : ℕ → ℕ) : RS :=
  .seq (.push NLogN.KFB.rP (.lit ((2 * off : ℕ) : ℤ)))
  (.seq (NLogN.mvK 0 qP NLogN.KFB.c1 (.reg NLogN.KFB.rP))
  (.seq (.push xA (.lit ((u % 4 : ℕ) : ℤ)))
  (.seq (.push xN (.lit 1))
  (.seq (tblP aE K)
  (.seq (.rep NLogN.KFB.c1 (.reg xN) mulB)
  (.seq (.rep NLogN.KFB.c1 (.reg xN) (.pop xB))
  (.seq (.asg xN (.add (.reg xN) (.reg xN)))
  (.seq (NLogN.mvK qO 0 NLogN.KFB.c1 (.reg xN))
  (.seq (.pop xN)
  (.seq (NLogN.mvK qP 0 NLogN.KFB.c1 (.reg NLogN.KFB.rP)) (.pop NLogN.KFB.rP)))))))))))

/-- Clean scratch of the phase program. -/
structure ClT (f : Fin HS → List ℤ) : Prop where
  a : f xA = []
  b : f xB = []
  c : f xC = []
  d : f xD = []
  n : f xN = []
  c1 : f NLogN.KFB.c1 = []
  ur : f rUr = []
  ui : f rUi = []
  o : f qO = []
  p : f qP = []
  rp : f NLogN.KFB.rP = []

/-- Cost of the phase program: `O(off + 2^K)` element moves of `O(w)` bits, plus `O(K Z)`. -/
def cPhT (w K off Z : ℕ) : ℕ :=
  2000 * off * (w + 11) + 2 ^ K * (cMulB w + 400000 + 2000 * (w + 11)) + cTbl K Z +
    100000 * (Z + 2)

end NLogN.KFG

end

section
-- Source module: Solutions.KFGPhase
/-!
# KFG: the `phase` step in `O(V)` (task B)

`phaseT m n w k₀ k₁ a := phT (mn) (w 2^{mn}) (k₀ + k₁ n) (phE a)` builds the exponent table of
the `2^{mn}` addresses by doubling (`O(2^{mn})` total) and rotates every entry of role `w` by its
table entry (`O(1)` selections per entry).  Compared with `KFB.phaseProg` (one pass per address
bit, `O(mn · 2^{mn})`), the per-element `mn = log V` factor is gone.

* `getD_phOut`: index semantics (same as `KFB.phL_spec`).
* `phaseT_iruns`: the fine IL run, its output encodes `RStep.den T (phase w k₀ k₁ a) *ᵥ x`,
  cost `cPhT wb (mn) (w 2^{mn}) Z = O((w + 1) 2^{mn} · (wb + 1)) + O(mn · Z)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-- The output list of the phase program on the segment `[off, off + 2^K)`. -/
def phOut (K off u : ℕ) (aE : ℕ → ℕ) (zs : List GI) : List GI :=
  zs.take off ++ (mulL (tbl u aE K) ((zs.drop off).take (2 ^ K)) ++ zs.drop (off + 2 ^ K))

/-- The `phase` program in `O(V)`. -/
def phaseT (m n : ℕ) (w k₀ k₁ : ℕ) (a : Fin m → ℕ) : RS :=
  phT (m * n) (w * 2 ^ (m * n)) (k₀ + k₁ * n) (phE (n := n) a)

end NLogN.KFG

end

section
-- Source module: Solutions.KFGChild
/-!
# KFG: the IL wrapper of a `child` call (task D)

`childI off A q M PT` on the data stack `pre ++ seg ++ rest` (Gaussian, `|pre| = off`,
`|seg| = A·M·q`):

1. parks `pre` on `qP` (`parkP`, fine, `O(off)` moves);
2. pushes `[A, q, M, 2]` on stacks `1..4` and runs KFC's `fwdF` (the Gaussian field move
   `[A, q, M] → [A, M, q]`, `KFB.mvL M q`), pops;
3. runs the called program `PT` on the data stack `mvL M q seg ++ rest`;
4. pushes `[A, q, M, 2]`, runs `bwdF` (`mvL q M`), pops;
5. unparks `pre` (`unparkP`).

`childI_runs`: the IL run, given a run specification of `PT` with list semantics `F`; the
result is `KFB.chL N off q M F zs`, so `KFB.rep_child` gives the denotation
`RStep.den T (child w h fwd)` (`child_iruns`).  Cost: `O(off · w)` for the park, `cFwdF + cBwdF`
(`O(q · kap w · |seg|)`, `q = 2^n` constant), `O((Bb+2)^2)` for the parameter pushes, plus the
cost of `PT`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

section
variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- The four parameters of the field moves. -/
def l4 (A q M : ℕ) : List (Fin HS × Ex HS) :=
  [(1, .lit (A : ℤ)), (2, .lit (q : ℤ)), (3, .lit (M : ℤ)), (4, .lit 2)]

/-! ## Parking the prefix -/

/-- Park the first `off` Gaussian entries of the data stack on `qP` (count in `rP`). -/
def parkP (off : ℕ) : RS := .seq (.push NLogN.KFB.rP (.lit ((2 * off : ℕ) : ℤ))) (NLogN.mvK 0 qP NLogN.KFB.c1 (.reg NLogN.KFB.rP))

/-- Move them back and drop the count. -/
def unparkP : RS := .seq (NLogN.mvK qP 0 NLogN.KFB.c1 (.reg NLogN.KFB.rP)) (.pop NLogN.KFB.rP)

/-! ## The wrapper -/

section
variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- The `child` wrapper around the called program `PT`. -/
def childI (off A q M : ℕ) (PT : IStmt HS Pp) : IStmt HS Pp :=
  .seq (fcode (parkP off)) (.seq (withPushI (l4 A q M) KFC.fwdF) (.seq PT
    (.seq (withPushI (l4 A q M) KFC.bwdF) (fcode unparkP))))

/-- Cost of the wrapper (without the called program). -/
def cChild (w w' off A q M Bb Z : ℕ) : ℕ :=
  2 * (2 * 2000 * (Z + 2) + (2 * off * (250 * (w + w') + 2600) + 2000) + (2 * Z + 2)) +
    2 * (pushCost Bb (l4 A q M) + (4 * (2 * Bb + 2) + 1)) + KFC.cFwdF A q M 2 w +
    KFC.cBwdF A q M 2 w'

end

end
end NLogN.KFG

end

section
-- Source module: Solutions.KFGGate
/-!
# KFG: dyadic gate atoms, dyadic gates, unit scalings, role swaps and role permutations

All steps are pair passes (`pairW`) on the scaled encoding `Enc`.

* `Atom` `⟨t, s, c, h⟩`: `y_t += (c / 2^h) y_s` (`c ∈ {-1,0,1}²`).  The target's scale becomes
  `E = max (e t) (e s + h)`; the pass computes `z_t ← 2^{E - e t} z_t + c · 2^{E - e s - h} z_s`
  (exact).  `enc_atom`.
* `gateP as e L`: a list of atoms whose targets are never sources (`NoTS`), applied in order;
  `enc_gate`: the result encodes `y_w + Σ_{a.t = w} (c_a / 2^{h_a}) y_{a.s}`, i.e. the gate matrix
  `1 + Σ_a (c_a/2^{h_a}) e_{a.t, a.s}` applied pointwise across roles.
* `UAtom` `⟨t, u, t'⟩`: `y_t ← u y_t` (a pass on `t` and a partner role `t'`); `enc_uatom`.
* `swapW`: exchange of two roles (and their scales); `perm_swaps`: every permutation of `Fin W`
  is a product of at most `|supp σ| ≤ W` swaps; `enc_permP`: the role permutation.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

def scOne : SC := ⟨(1, 0), 0⟩
def scZero : SC := ⟨(0, 0), 0⟩

@[simp] theorem scOne_d : scOne.d = 0 := rfl

@[simp] theorem scZero_d : scZero.d = 0 := rfl

/-! ## Gate atoms -/

/-- `y_t += (c / 2^h) · y_s`. -/
structure Atom where
  t : ℕ
  s : ℕ
  c : GI
  h : ℕ

namespace Atom

variable (a : Atom) (e : ℕ → ℕ)

def Ok (W : ℕ) : Prop := a.t < W ∧ a.s < W ∧ a.t ≠ a.s ∧ Sm a.c

/-- The new scale of the target. -/
def E : ℕ := max (e a.t) (e a.s + a.h)
def al : ℕ := a.E e - e a.t
def be : ℕ := a.E e - e a.s - a.h

/-- Scales after the atom. -/
def sc : ℕ → ℕ := Function.update e a.t (a.E e)

/-- The matrix of the pass. -/
def M : M2S :=
  if a.t < a.s then ⟨⟨(1, 0), a.al e⟩, ⟨a.c, a.be e⟩, scZero, scOne⟩
  else ⟨scOne, scZero, ⟨a.c, a.be e⟩, ⟨(1, 0), a.al e⟩⟩

/-- The pass. -/
def P (L : ℕ) : RS := if a.t < a.s then pairW (a.M e) L a.t a.s else pairW (a.M e) L a.s a.t

/-- The new streams. -/
noncomputable def y (y : ℕ → ℕ → ℂ) : ℕ → ℕ → ℂ :=
  Function.update y a.t (fun p => y a.t p + NLogN.KFB.toC a.c / 2 ^ a.h * y a.s p)

/-- Shift bound of the pass. -/
def D : ℕ := max (a.al e) (a.be e)

end Atom

/-- The list after one atom. -/
def Atom.L (a : Atom) (e : ℕ → ℕ) (L : ℕ) (zs : List GI) : List GI :=
  if a.t < a.s then pairL (a.M e) L a.t a.s zs else pairL (a.M e) L a.s a.t zs

/-- Cost of one atom: a pair pass on the first `W` roles. -/
def cAtom (w D L W Z : ℕ) : ℕ := cPairW w D L 0 W Z

/-! ## Gates: lists of atoms -/

/-- Scales after a list of atoms. -/
def gateE : List Atom → (ℕ → ℕ) → (ℕ → ℕ)
  | [], e => e
  | a :: as, e => gateE as (a.sc e)

/-- The program of a list of atoms. -/
def gateP : List Atom → (ℕ → ℕ) → ℕ → RS
  | [], _, _ => .skip
  | a :: as, e, L => .seq (a.P e L) (gateP as (a.sc e) L)

/-- The list after a list of atoms. -/
def gateL : List Atom → (ℕ → ℕ) → ℕ → List GI → List GI
  | [], _, _, zs => zs
  | a :: as, e, L, zs => gateL as (a.sc e) L (a.L e L zs)

/-- Bit bound after a list of atoms. -/
def gateW : List Atom → (ℕ → ℕ) → ℕ → ℕ
  | [], _, w => w
  | a :: as, e, w => gateW as (a.sc e) (w + a.D e + 2)

/-- Cost of a list of atoms. -/
def cGate : List Atom → (ℕ → ℕ) → ℕ → ℕ → ℕ → ℕ → ℕ
  | [], _, _, _, _, _ => 1
  | a :: as, e, w, L, W, Z => cAtom w (a.D e) L W Z + cGate as (a.sc e) (w + a.D e + 2) L W Z

/-- Targets are never sources. -/
def NoTS (as : List Atom) : Prop := ∀ a ∈ as, ∀ b ∈ as, a.t ≠ b.s

/-- The streams after a gate (simultaneous update). -/
noncomputable def gateY (as : List Atom) (y : ℕ → ℕ → ℂ) : ℕ → ℕ → ℂ := fun w p =>
  y w p + (as.map fun a => if a.t = w then NLogN.KFB.toC a.c / 2 ^ a.h * y a.s p else 0).sum

/-! ## The gate matrix -/

/-- The role matrix of a gate: `1 + Σ_a (c_a / 2^{h_a}) e_{t_a, s_a}`. -/
noncomputable def gateMat (W : ℕ) (as : List Atom) : Matrix (Fin W) (Fin W) ℂ := fun p q =>
  (if p = q then 1 else 0) +
    (as.map fun a => if a.t = p.val ∧ a.s = q.val then NLogN.KFB.toC a.c / 2 ^ a.h else 0).sum

/-! ## Unit scalings -/

/-! ## Role swaps and permutations -/

/-- The swap matrix. -/
def swS : M2S := ⟨scZero, scOne, scOne, scZero⟩

/-- `ℕ`-level swap of two indices. -/
def nsw (i j w : ℕ) : ℕ := if w = i then j else if w = j then i else w

/-- A list of swaps `(i, j)` with `i < j`, applied in order. -/
def swapsP (L : ℕ) : List (ℕ × ℕ) → RS
  | [] => .skip
  | (i, j) :: l => .seq (pairW swS L i j) (swapsP L l)

def swapsL (L : ℕ) : List (ℕ × ℕ) → List GI → List GI
  | [], zs => zs
  | (i, j) :: l, zs => swapsL L l (pairL swS L i j zs)

/-- The composite index map `τ₁ ∘ τ₂ ∘ ⋯`. -/
def swapsF : List (ℕ × ℕ) → ℕ → ℕ
  | [] => id
  | (i, j) :: l => nsw i j ∘ swapsF l

/-- Every permutation of `Fin W` is a product of at most `|supp σ|` swaps `(i, j)`, `i < j`. -/
theorem perm_swaps {W : ℕ} : ∀ (n : ℕ) (σ : Equiv.Perm (Fin W)), σ.support.card ≤ n →
    ∃ l : List (Fin W × Fin W), l.length ≤ σ.support.card ∧ (∀ p ∈ l, p.1 < p.2) ∧
      (l.map fun p => Equiv.swap p.1 p.2).prod = σ
  | 0, σ, hn => by
    have : σ = 1 := by
      rw [← Equiv.Perm.support_eq_empty_iff]; exact Finset.card_eq_zero.1 (by omega)
    exact ⟨[], by simp, by simp, by simp [this]⟩
  | n + 1, σ, hn => by
    by_cases h1 : σ = 1
    · exact ⟨[], by simp, by simp, by simp [h1]⟩
    obtain ⟨x, hx⟩ : ∃ x, σ x ≠ x := by
      by_contra hc; push Not at hc; exact h1 (Equiv.ext hc)
    have hlt := Equiv.Perm.card_support_swap_mul hx
    obtain ⟨l, hl1, hl2, hl3⟩ := perm_swaps n (Equiv.swap x (σ x) * σ) (by omega)
    have hsw : σ = Equiv.swap x (σ x) * (Equiv.swap x (σ x) * σ) := by
      rw [← mul_assoc, Equiv.swap_mul_self, one_mul]
    rcases lt_or_gt_of_ne (Ne.symm hx) with h | h
    · refine ⟨(x, σ x) :: l, by simp; omega, ?_, ?_⟩
      · intro p hp
        rcases List.mem_cons.1 hp with rfl | hp
        · exact h
        · exact hl2 p hp
      · simp only [List.map_cons, List.prod_cons, hl3]
        exact hsw.symm
    · refine ⟨(σ x, x) :: l, by simp; omega, ?_, ?_⟩
      · intro p hp
        rcases List.mem_cons.1 hp with rfl | hp
        · exact h
        · exact hl2 p hp
      · simp only [List.map_cons, List.prod_cons, hl3]
        rw [Equiv.swap_comm]; exact hsw.symm

/-- `ℕ`-level action of a permutation of `Fin W`. -/
def pf {W : ℕ} (σ : Equiv.Perm (Fin W)) (w : ℕ) : ℕ := if h : w < W then (σ ⟨w, h⟩).val else w

/-- The swap list of a permutation (the inverse is decomposed: the passes realise `y ∘ σ⁻¹`). -/
noncomputable def permSw {W : ℕ} (σ : Equiv.Perm (Fin W)) : List (ℕ × ℕ) :=
  ((perm_swaps W σ⁻¹ (Finset.card_le_univ _ |>.trans (by simp))).choose).map
    fun p => (p.1.val, p.2.val)

end NLogN.KFG

end

section
-- Source module: Solutions.KFGLink
/-!
# KFG: the steps `gate` and `perm` as `RStep.den` on the (scaled) state encoding

`ScRep e zs x`: role-major Gaussian array `zs` (role length `2^{mn}`, address index `aidx`)
encodes the state `x` with per-role scale `2^{-e w}`: `toC zs[w 2^{mn} + aidx b] = 2^{e w} x(w,b)`
(`ScRep 0 = KFB.Rep`, `scRep_zero`).

* `gate_iruns`: for a list of dyadic atoms `as` (targets never sources), the fine IL run of
  `gateP as e 2^{mn}` leaves an encoding (new scales `gateE as e`) of
  `RStep.den T (gate (gateMat W as)) *ᵥ x`, cost `cGate`, `O(|as|)` pair passes, each
  `O(1)` per element per touched role (`O(w + D)` bit operations per element).
* `perm_iruns`: the role permutation, `≤ W` swaps, encodes `RStep.den T (perm σ) *ᵥ x`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFG

open IntMul RS AModel KIL KFD KFB Complex

/-- The scaled state encoding. -/
def ScRep {W m n : ℕ} (e : ℕ → ℕ) (zs : List GI) (x : Fin W × (Fin m × Fin n → Fin 2) → ℂ) :
    Prop :=
  zs.length = W * 2 ^ (m * n) ∧
    ∀ w b, NLogN.KFB.toC (zs.getD (w.val * 2 ^ (m * n) + aidx b) 0) = 2 ^ (e w.val) * x (w, b)

/-- The streams read off an array. -/
noncomputable def yOf (L : ℕ) (e : ℕ → ℕ) (zs : List GI) : ℕ → ℕ → ℂ := fun w p =>
  NLogN.KFB.toC (zs.getD (w * L + p) 0) / 2 ^ (e w)

end NLogN.KFG

end

section
-- Source module: Solutions.KFHRec
/-!
# KFH: fixed-point record arrays and the `Q_p` truncation pass

**Record format.**  The array of Proposition 14 is a list `xs` of records, each record `r`
complex coefficients (`Fin r → ℂ`).  On the item stacks it is a Gaussian list `zs` (two integer
items per coefficient, `gflat`), records consecutive, coefficients in order, with one common
scale: `RecRep r F zs xs` says coefficient `c` of record `j` is `toC zs[j r + c] / 2^F`.
The scale `F` (number of fractional bits) is bookkeeping, not data: every halving of §5.4 is
`F ↦ F + 1` and costs nothing; the numerators stay exact (this is the guarded format of §5.4
with the reserved fractional zeros left implicit).  Bit widths are tracked by `GB w zs`.

**Truncation.** `truncP Δ N` replaces each of the `N` integer items `a` on top of stack `0` by
`a.tdiv 2^Δ` (truncation toward zero, `Ex.quo`, linear cost).  On numerators at scale
`p + Δ` this is exactly `Q_p` (`truncNum_tdiv`): `truncP_rep` gives
`RecRep r p zs' (QpArr p xs)` from `RecRep r (p + Δ) zs xs`.  Cost `O(N (w + Δ)) + O(Z)`
(`truncP_fsp`), no `Bb` factor.  The same pass with an exact divisor is the renormalisation step.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFH

open IntMul RS AModel KIL KFD KFB Complex

/-! ## The record format -/

/-- `zs` stores the records `xs` (each `r` coefficients) at scale `2^{-F}`. -/
def RecRep (r F : ℕ) (zs : List GI) (xs : List (Fin r → ℂ)) : Prop :=
  zs.length = xs.length * r ∧
    ∀ j (c : Fin r), j < xs.length → NLogN.KFB.toC (zs.getD (j * r + c) 0) = 2 ^ F * xs.getD j 0 c

/-! ## Truncation toward zero of a dyadic numerator -/

/-- Componentwise truncation of a Gaussian numerator. -/
def gtr (Δ : ℕ) (z : GI) : GI := (z.1.tdiv (2 ^ Δ), z.2.tdiv (2 ^ Δ))

/-! ## The truncation pass -/

/-- `a ↦ a.tdiv 2^Δ` on the top item of stack `0`. -/
def quoE (Δ : ℕ) : Ex HS := .quo (.reg 0) (.lit ((2 : ℤ) ^ Δ))

/-- The loop: `N` times, push the truncated top of `0` on `qO` and pop it. -/
def truncLoop (Δ : ℕ) : RS := .rep NLogN.KFB.c1 (.reg rN) (.seq (.push qO (quoE Δ)) (.pop 0))

/-- **The truncation pass** on the top `N` items of stack `0`. -/
def truncP (Δ N : ℕ) : RS :=
  .seq (.push rN (.lit (N : ℤ))) (.seq (truncLoop Δ) (.seq (NLogN.mvK qO 0 NLogN.KFB.c1 (.reg rN)) (.pop rN)))

/-- Cost of the truncation pass on `N` items of at most `w` bits. -/
def cTrunc (w Δ N Z : ℕ) : ℕ :=
  N * (8100 * (w + Δ + 3) + 2100) + 2000 + (N * (250 * w + 2600) + 2000) + 4000 * (Z + 2) + 2 * Z + 2

end NLogN.KFH

end

section
-- Source module: Solutions.KFHKer
/-!
# KFH: an individual selected-bit kernel as one `qop` pass on a record array

The layer layout of Proposition 14 is `[P] × [2^K]^D × [S]` (`Butterfly.layerShape`), records in
lexicographic order.  Bit `ρ` of chunk `i` has record stride `stT = 2^ρ · vol(suffix after
chunk i)`: the address with that bit flipped is `j ± stT` (`layer_partner`).  On the flat
Gaussian list (records of `r` coefficients) the paired entries are at distance `B1 = stT · r`,
in blocks of `2 B1`; this is exactly the two-stream pass `KFB.qop [M] B1 G` with
`G = P · 2^{K i} · 2^{K-ρ-1}` groups.

* `encode_set`, `getD_digit`: mixed-radix arithmetic of one changed digit.
* `layer_partner`: selected bit `i` of record `j` is `j / stT % 2`; flipping it moves to
  `j + (b - bit) stT`.
* `ker_rep`: if `toC M = 2^e A` entrywise (`KerOk`), `qopL [M] (stT r) G` represents
  `selOp (tensF (kerAt i A))` (the kernel (25) of `A` on bit `i`) at scale `F + e`.
* `kerP_fsp`: the self-contained pass `kerP = qopW [M] (stT r) G 0`, cost `cQopW`, width `+2`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFH

open IntMul RS AModel KIL KFD KFB Complex

/-! ## Mixed-radix arithmetic -/

/-! ## The layer layout -/

variable {P K D S ρ : ℕ}

/-- Volume of the fields after chunk `i`. -/
def stS (K D S : ℕ) (i : ℕ) : ℕ := vol ((Butterfly.layerShape 1 K D S).drop (i + 2))

/-- Record stride of selected bit `ρ` of chunk `i`. -/
def stT (K D S ρ : ℕ) (i : ℕ) : ℕ := 2 ^ ρ * stS K D S i

/-- Number of `2 stT` blocks. -/
def grpG (P K ρ : ℕ) (i : ℕ) : ℕ := P * (2 ^ K) ^ i * 2 ^ (K - ρ - 1)

/-! ## The kernel pass on records -/

/-- Entries of a Gaussian `2 × 2` matrix. -/
def Ment (M : M2G) (x y : Fin 2) : GI :=
  if x = 0 then (if y = 0 then M.a else M.b) else (if y = 0 then M.c else M.d)

/-- `M = 2^e A` entrywise. -/
def KerOk (M : M2G) (A : Butterfly.M2) (e : ℕ) : Prop := ∀ x y, NLogN.KFB.toC (Ment M x y) = 2 ^ e * A x y

/-! ## The self-contained pass -/

/-- The kernel pass on bit `ρ` of chunk `i` (records of `r` coefficients). -/
def kerP (K D S ρ P r : ℕ) (i : ℕ) (M : M2G) : RS :=
  qopW [M] (stT K D S ρ i * r) (grpG P K ρ i) 0

/-- Its descriptor size bound. -/
def kerZ (K D S ρ P r : ℕ) (i : ℕ) : ℕ :=
  4 * (stT K D S ρ i * r + grpG P K ρ i + grpG P K ρ i * (2 * (stT K D S ρ i * r) * 1) + 0)

/-- Its cost. -/
def cKer (K D S ρ P r : ℕ) (i : ℕ) (w Z : ℕ) : ℕ :=
  cQopW w (stT K D S ρ i * r) 1 (grpG P K ρ i) 0 Z

end NLogN.KFH

end

section
-- Source module: Solutions.KFHLeaf
/-!
# KFH: Proposition 14 by individual kernels, as one fine IL program

`LImpl T code e dw cost`: on any clean managed state holding a record array `zs` (scale `F`,
width `w`) that represents `xs` on the layer layout, `code` leaves an array representing
`selOp T xs` at scale `F + e` and width `w + dw`, at cost `cost w`.  Implementations compose
(`LImpl.comp`: matrix product, scales and widths add), and each kernel pass is one
(`limpl_ker`), so `A^{⊗D}` is `D` passes (`limpl_tens`, via `Butterfly.tens_eq_prod_kerAt`).

The program `p14Leaf` is the paper's §5.1/§5.3 procedure with every `C^{⊗}` evaluated by
individual kernels (25) (the `D ≤ q₀` branch and the leaves of the recursion):

```
S_D (D passes, diag(1,i))  ;  C^{⊗D} (D passes, 2C = [[1+i,1-i],[1-i,1+i]], F += 1 each)
S_D (D passes)  ;  b^D (one pass, (-i)^t (1-i)^{D mod 2}, F += t + D mod 2, t = ⌊D/2⌋)
Q_p (truncation pass: numerators tdiv 2^{F-p})
```

`p14Leaf_run`: output `= Q_p(H0^{⊗D} z) = prop14Proc … (tens Cm) z` (by (18) `tens_H0` and
`prop14_output`), disk outputs, error `< √2·2^{-p}`, output numerators of `≤ p + 1` bits, cost
`≤ 10^6 (D+1)(w + 8D + 8)·N + 10^5 (D+1)(Z + 2)` with `N = vol·r` Gaussian entries and
`Z = zsz(12 N)`: linear in the data per pass, `O(D)` passes (the `C₀ V u` leaf cost of
`recCost_le`), no `Bb` factor.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFH

open IntMul RS AModel KIL KFD KFB Complex

variable {P K D S ρ : ℕ}

/-! ## Layer implementations -/

/-- `code` implements the selected-bit operator `T` on record arrays of the layer layout. -/
def LImpl (hρ : ρ < K) (r : ℕ) (T : Butterfly.Op (Fin D)) (code : RS) (e dw : ℕ)
    (cost : ℕ → ℕ) : Prop :=
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ) (xs : List (Fin r → ℂ))
    (F w : ℕ), Cl f → f 0 = gflat zs ++ l → RecRep r F zs xs →
    xs.length = vol (Butterfly.layerShape P K D S) → NLogN.KFB.GB w zs →
    ∃ zs', RecRep r (F + e) zs' (Butterfly.selOp (Butterfly.layerSel P K D S ρ hρ) T xs) ∧
      NLogN.KFB.GB (w + dw) zs' ∧
      FSp code (encK s f) (encK s (Function.update f 0 (gflat zs' ++ l))) (cost w)

/-- Cost bound of one kernel pass on `N` Gaussian entries. -/
def ck (N Z w : ℕ) : ℕ := 100000 * (w + 3) * N + 30000 * (Z + 2)

/-- The kernel passes of a list of bits. -/
def kersP (K D S ρ P r : ℕ) (M : M2G) : List ℕ → RS
  | [] => .skip
  | i :: l => .seq (kersP K D S ρ P r M l) (kerP K D S ρ P r i M)

/-! ## The concrete kernels -/

/-- `diag(1, i)` (the phase `S`), exact. -/
def MS : M2G := ⟨(1, 0), 0, 0, (0, 1)⟩

/-- `2C = [[1+i, 1-i], [1-i, 1+i]]`. -/
def MC : M2G := ⟨(1, 1), (1, -1), (1, -1), (1, 1)⟩

/-- The Gaussian numerator of `b^D`: `(-i)^t (1-i)^{D mod 2}`, `t = ⌊D/2⌋`. -/
def gD (D : ℕ) : GI := if D % 2 = 0 then ipw (3 * (D / 2)) else gmul (ipw (3 * (D / 2))) (1, -1)

/-- The scalar pass `dg g g` on chunk `0`. -/
def MD (D : ℕ) : M2G := dg (NLogN.KFH.gD D) (NLogN.KFH.gD D)

/-! ## The program -/

/-- Fractional bits added inside one layer: `D` kernels `C` and the scalar `2^{-(t + D mod 2)}`. -/
def dScale (D : ℕ) : ℕ := D + (D / 2 + D % 2)

/-- **Proposition 14 by individual kernels** on `[P] × [2^K]^D × [S]`, records of `r`
coefficients, as one register program (compiled with `fcode`). -/
def p14Leaf (P K D S ρ r : ℕ) : RS :=
  .seq (.seq (.seq (.seq (kersP K D S ρ P r MS ((List.finRange D).map Fin.val))
    (kersP K D S ρ P r MC ((List.finRange D).map Fin.val)))
    (kersP K D S ρ P r MS ((List.finRange D).map Fin.val)))
    (kerP K D S ρ P r 0 (MD D)))
    (truncP (dScale D) (2 * (vol (Butterfly.layerShape P K D S) * r)))

/-- The exact cost expression. -/
def cP14 (D N Z w : ℕ) : ℕ :=
  (D * ck N Z (w + 2 * D) + 1) + (D * ck N Z (w + 2 * D + 2 * D) + 1) +
  (D * ck N Z (w + (2 * D + 2 * D) + 2 * D) + 1) + ck N Z (w + (2 * D + 2 * D + 2 * D)) +
  cTrunc (w + 6 * D + 2) (dScale D) (2 * N) Z

end NLogN.KFH

end

section
-- Source module: Solutions.KFHComp
/-!
# KFH: the Proposition 14 wrapper around an arbitrary `C^{⊗D}` implementation

The recursion of §5.3 replaces the `D` individual `C` kernels by the batched network.  The
interface between the two is `LImpl … (tens Cm) Cprog D dC kC`: any fine IL code that maps a
record array at scale `F` to the `C`-layer at scale `F + D` (exact, the `D` halvings of the
kernels (25) kept as scale), widths `+dC`, cost `kC w`.

* `p14With Cprog = S_D ; Cprog ; S_D ; b^D ; Q_p` and **`p14With_iruns`**: for every such
  `Cprog`, the output represents `prop14Proc … (tens Cm) z = Q_p(H0^{⊗D} z)`, disk, error
  `< √2·2^{-p}`, `≤ p+1`-bit numerators; cost `kC (w + 2D) + O(D (w + D + dC) N) + O(D Z)`.
* `limpl_cinv`: from any `C`-layer implementation, the inverse layer `C^{-1 ⊗D}` by (17):
  `Z_D ; Cprog ; Z_D ; (-i)^D` (two sign passes per chunk and one unit pass, no extra call).
  This is the inverse-child wrapper of the network steps `NStep.child w h false`.
* `limpl_leafC`: the individual-kernel `C`-layer (`kersP MC`), the leaf of the recursion.
* `renorm_rep`: exact renormalisation (`truncP` on numerators divisible by `2^Δ`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFH

open IntMul RS AModel KIL KFD KFB Complex

variable {P K D S ρ : ℕ}

/-! ## Exact renormalisation -/

/-! ## The wrapper -/

/-- Proposition 14 around a `C`-layer implementation `Cprog`. -/
def p14With (P K D S ρ r : ℕ) (Cprog : RS) : RS :=
  .seq (.seq (.seq (.seq (kersP K D S ρ P r MS ((List.finRange D).map Fin.val)) Cprog)
    (kersP K D S ρ P r MS ((List.finRange D).map Fin.val)))
    (kerP K D S ρ P r 0 (MD D)))
    (truncP (dScale D) (2 * (vol (Butterfly.layerShape P K D S) * r)))

/-- Cost of the wrapper (without the `C`-layer). -/
def cWrap (D N Z w dC : ℕ) : ℕ :=
  (D * ck N Z (w + 2 * D) + 1) + (D * ck N Z (w + (2 * D + dC) + 2 * D) + 1) +
    ck N Z (w + (2 * D + dC + 2 * D)) + cTrunc (w + 4 * D + dC + 2) (dScale D) (2 * N) Z

/-! ## The leaf `C`-layer and the inverse layer -/

/-- The sign `Z = diag(1, -1)`. -/
def MZ : M2G := ⟨(1, 0), 0, 0, (-1, 0)⟩

/-- The unit `(-i)^D`. -/
def MU (D : ℕ) : M2G := dg (ipw (3 * D)) (ipw (3 * D))

end NLogN.KFH

end

section
-- Source module: Solutions.KFLConv
/-!
# KFL §16: integers ↔ bits in IL, linear in the bit width

`encW w z = sgn :: [bit 0, …, bit (w-1)]` of `|z|` (`zsz z ≤ w`), `bitsL w xs = xs.flatMap (encW w)`.

* `toBitsI N w`: `ints xs ++ R ↦ ints (bitsL w xs) ++ R` on stack `0` (`|xs| = N`): per integer a
  sign test and `w` splits by the ruler `2` (`split` costs `O(m)` for the ruler `2^m`), so
  `O(w + 1)` per integer — no per-element logarithm, the counters are unary.
* `fromBitsI N w`: the inverse: per integer `w` `cat`s by the ruler `2` from the sentinel `1`,
  one `split` by the ruler `2^w` (built once by `w` shifts), and a sign test.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unnecessarySeqFocus false
set_option maxRecDepth 20000

namespace NLogN.KFL

open IntMul KFC

variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

abbrev cO : Fin HS := 40
abbrev cW : Fin HS := 41
abbrev cR2 : Fin HS := 42
abbrev cRW : Fin HS := 43
abbrev cT : Fin HS := 44
abbrev cSg : Fin HS := 45
abbrev cZ : Fin HS := 46
abbrev cPN : Fin HS := 47
abbrev cPW : Fin HS := 48
abbrev cCN : Fin HS := 49
abbrev cDN : Fin HS := 50
abbrev cCw : Fin HS := 51
abbrev cDw : Fin HS := 52

/-! ## The encoding -/

def bitN (n i : ℕ) : ℤ := ((n / 2 ^ i % 2 : ℕ) : ℤ)
def sgnB (z : ℤ) : ℤ := if z < 0 then 1 else 0

def encW (w : ℕ) (z : ℤ) : List ℤ := sgnB z :: (List.range w).map (bitN z.natAbs)
def bitsL (w : ℕ) (xs : List ℤ) : List ℤ := xs.flatMap (encW w)

/-! ## Peeling `w` bits -/

def isNg (c : Cls) : Bool := c == .neg

def peel : IStmt HS Pp := uloop cCw cDw (.op (.split cW cR2 cO))

def cPeel (w : ℕ) : ℕ := w * (splitCost 1 + 17 + 64) + 60

/-! ## One integer to bits -/

def sgnStep : IStmt HS Pp :=
  .ite cW isNg (.seq (.op (.neg cW)) (.op (.const cO 1))) (.op (.const cO 0))

def body1 : IStmt HS Pp :=
  .seq (.op (.mv 0 cW)) (.seq sgnStep (.seq peel (.op (.drop cW))))

def cBody1 (w : ℕ) : ℕ := 76 * w + 101 + 48 + cPeel w + 2

/-! ## All integers -/

def fwdLoop : IStmt HS Pp := uloop cCN cDN body1

/-- Move `n (w + 1)` integers from `cO` to `0` (counters `cCN` = `n`, `cCw` = `w`). -/
def mvBack : IStmt HS Pp := uloop cCN cDN (.seq (.op (.mv cO 0)) (mvB cO 0 cCw cDw))

def cMB (w : ℕ) : ℕ := kap 1 + w * kap 1 + 60

/-! ## The forward conversion -/

def toCore : IStmt HS Pp := .seq fwdLoop mvBack

/-- **Integers to bits.** -/
def toBitsI (N w : ℕ) : IStmt HS Pp :=
  .seq (.op (.const cPN N)) (.seq (.op (.const cPW w)) (.seq (.op (.const cR2 2))
    (.seq (NLogN.KFC.wrap cPN cCN cDN (NLogN.KFC.wrap cPW cCw cDw toCore))
    (.seq (.op (.drop cR2)) (.seq (.op (.drop cPW)) (.op (.drop cPN)))))))

def cTo (N w : ℕ) : ℕ :=
  (zsz (N : ℤ) + 2) + (zsz (w : ℤ) + 2) + 4 +
    ((N * (cBody1 w + 64) + 60 + (N * (cMB w + 64) + 60) + cUL w + w * 31 + 38) + cUL N + N * 31
      + 38) + 6 + (2 * zsz (w : ℤ) + 2) + (2 * zsz (N : ℤ) + 2)

/-! ## Bits to one integer -/

def mkRW : IStmt HS Pp := .seq (.op (.const cRW 1)) (uloop cCw cDw (.op (.shl cRW cR2)))

def cMkRW (w : ℕ) : ℕ := 3 + (w * (shlCost 1 + 64) + 60)

def catLoop : IStmt HS Pp := uloop cCw cDw (.op (.cat cW cR2 cT))

def isPs (c : Cls) : Bool := c == .pos

def sgnFix : IStmt HS Pp := .ite cSg isPs (.op (.neg NLogN.KFL.cZ)) .skip

def body2 : IStmt HS Pp :=
  .seq (.op (.mv 0 cSg)) (.seq (mvB 0 cT cCw cDw) (.seq (.op (.const cW 1)) (.seq catLoop
    (.seq (.op (.split cW cRW NLogN.KFL.cZ)) (.seq (.op (.drop cW)) (.seq sgnFix (.seq (.op (.drop cSg))
      (.op (.mv NLogN.KFL.cZ cO)))))))))

def cBody2 (w : ℕ) : ℕ := 177 + (w * kap 1 + 60) + 3 + (w * (catCost 1 + 64) + 60) +
  (splitCost w + 17) + 4 + 45 + 4 + (76 * w + 101)

def fromLoop : IStmt HS Pp := uloop cCN cDN body2

def fromCore : IStmt HS Pp := .seq mkRW (.seq fromLoop (.seq (mvB cO 0 cCN cDN) (.op (.drop cRW))))

/-- **Bits to integers.** -/
def fromBitsI (N w : ℕ) : IStmt HS Pp :=
  .seq (.op (.const cPN N)) (.seq (.op (.const cPW w)) (.seq (.op (.const cR2 2))
    (.seq (NLogN.KFC.wrap cPN cCN cDN (NLogN.KFC.wrap cPW cCw cDw fromCore))
    (.seq (.op (.drop cR2)) (.seq (.op (.drop cPW)) (.op (.drop cPN)))))))

def cFrom (N w : ℕ) : ℕ :=
  (zsz (N : ℤ) + 2) + (zsz (w : ℤ) + 2) + 4 +
    ((cMkRW w + (N * (cBody2 w + 64) + 60) + (N * kap w + 60) + (2 * (w + 1) + 2) + cUL w +
      w * 31 + 38) + cUL N + N * 31 + 38) + 6 + (2 * zsz (w : ℤ) + 2) + (2 * zsz (N : ℤ) + 2)

/-! ## Cost bounds: `O((N + 1)(w + 1))` -/

end NLogN.KFL

end

section
-- Source module: Solutions.KFLIdx
/-!
# KFL §1: index arithmetic of chunked arrays

The Lemma 13 rectangle (records of `B` items) is the flat array of shape
`sh8 = [N0, Q, N1, Q, N2, Q, N3, B]` with three chunk fields of size `Q = 2^L`
(slots `0, 1, 2` at list positions `1, 3, 5`).  Chunk `t` has item stride `S_t`
(`S2 = N3 B`, `S1 = Q S2 N2`, `S0 = Q S1 N1`, `V = Q S0 N0`).

* `gc Q S n = n / S % Q` (the chunk value) and `sc Q S v n` (the index with that chunk set to `v`),
  with the integer formula `sc_int` and the independence lemmas `gc_sc_self/hi/lo`.
* `rix V φ d xs` reindexes an array (`out[j] = xs[φ j]`); `rix_rix` composes.
* A triple map `F : Tr → Tr` (on the three chunk values) induces the index map `phi g F`;
  `phi_phi : phi F ∘ phi G = phi (F ∘ G)`, so `rixT G (rixT F xs) = rixT (F ∘ G) xs`.
* The array operations used by Lemma 13 as triple maps: the Lemma 12 swaps
  (`pull sh8 sh8 (swapPos 1 3 | 3 5 | 1 5)`, `pull_swap*`), the block rotation of `rotF` on the
  view `[P, Q, S2]` (`pull_rot`), and the rotation of one bit field (`[P', 2, D S1]`, `pull_bit`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open AModel

/-! ## Chunk get / set -/

/-- The chunk value at stride `S` (size `Q`). -/
def gc (Q S n : ℕ) : ℕ := n / S % Q

/-- `n` with the chunk at stride `S` (size `Q`) set to `v`. -/
def sc (Q S v n : ℕ) : ℕ := n / (Q * S) * (Q * S) + v * S + n % S

/-! ## Reindexing -/

/-- `out[j] = xs[φ j]` for `j < V`. -/
def rix {α : Type*} (V : ℕ) (φ : ℕ → ℕ) (d : α) (xs : List α) : List α :=
  (List.range V).map fun j => xs.getD (φ j) d

/-! ## Geometry and triple maps -/

/-- The rectangle `[N0] × [Q] × [N1] × [Q] × [N2] × [Q] × [N3]` with `B` items per record. -/
structure Geo where
  Q : ℕ
  N0 : ℕ
  N1 : ℕ
  N2 : ℕ
  N3 : ℕ
  B : ℕ

namespace Geo

variable (g : Geo)

def S2 : ℕ := g.N3 * g.B
def S1 : ℕ := g.Q * g.S2 * g.N2
def S0 : ℕ := g.Q * g.S1 * g.N1
def V : ℕ := g.Q * g.S0 * g.N0

/-- The flat shape. -/
def sh8 : Shape := [g.N0, g.Q, g.N1, g.Q, g.N2, g.Q, g.N3, g.B]

/-- All lengths positive. -/
structure Pos : Prop where
  q : 0 < g.Q
  n0 : 0 < g.N0
  n1 : 0 < g.N1
  n2 : 0 < g.N2
  n3 : 0 < g.N3
  b : 0 < g.B

variable {g}

end Geo

/-- Chunk-value triples `(slot 0, slot 1, slot 2)`. -/
abbrev Tr := ℕ × ℕ × ℕ

/-- All three values below `Q`. -/
def TV (Q : ℕ) (c : Tr) : Prop := c.1 < Q ∧ c.2.1 < Q ∧ c.2.2 < Q

/-- The chunk triple of index `n`. -/
def trOf (g : Geo) (n : ℕ) : Tr := (gc g.Q g.S0 n, gc g.Q g.S1 n, gc g.Q g.S2 n)

/-- The index map induced by a triple map. -/
def phi (g : Geo) (F : Tr → Tr) (n : ℕ) : ℕ :=
  NLogN.KFL.sc g.Q g.S0 (F (trOf g n)).1 (NLogN.KFL.sc g.Q g.S1 (F (trOf g n)).2.1 (NLogN.KFL.sc g.Q g.S2 (F (trOf g n)).2.2 n))

section Phi

variable {g : Geo} (hg : g.Pos)
include hg

end Phi

/-- The array after the pull by the triple map `F`. -/
def rixT (g : Geo) (F : Tr → Tr) (xs : List ℤ) : List ℤ := rix g.V (phi g F) 0 xs

/-! ## Lemma 12 swaps as triple maps -/

/-- The slot swaps. -/
def sw01 (c : Tr) : Tr := (c.2.1, c.1, c.2.2)
def sw12 (c : Tr) : Tr := (c.1, c.2.2, c.2.1)
def sw02 (c : Tr) : Tr := (c.2.2, c.2.1, c.1)

/-! ## Block rotations (`rotF` views) -/

/-- The rotation of slot 2 by an amount depending on slots 0, 1 (pull map). -/
def rot2 (Q : ℕ) (a : ℕ → ℕ → ℕ) (c : Tr) : Tr := (c.1, c.2.1, (c.2.2 + Q - a c.1 c.2.1) % Q)

/-! ## Rotation of one bit field of slot 1 -/

/-- Flip of the bit at `D` of slot 1 controlled by the table value `b c0`. -/
def flip1 (D : ℕ) (b : ℕ → ℕ) (c : Tr) : Tr :=
  (c.1, NLogN.KFL.sc 2 D ((gc 2 D c.2.1 + 2 - b c.1) % 2) c.2.1, c.2.2)

end NLogN.KFL

end

section
-- Source module: Solutions.KFLMath
/-!
# KFL §2: Lemma 13 on chunk triples (canonical layout `x, y, z` = slots `0, 1, 2`)

* `selM K ρ n = Σ_{i<n} 2^{ρ+iK}` (the selected-position mask), `testBit_selM`.
* `T0y_eq`: the ideal map `T0` of `ButterflyPacked` is `y ↦ y ⊕ (x ∧ selM (f-1))`;
  `guardN_xor`: it preserves every guard.
* The pull maps of the eight lines as triple maps `lineP` (realised by `rot2`, after a slot
  swap for the `y`-lines), their composite `FSc = lineP l₁ ∘ … ∘ lineP l₈` and the forward
  `S0c`; `S0c_FSc`, `FSc_S0c` (inverse bijections), `FSc_good` (off the guard set `FSc = T0c`).
* The correction: `keyc = T0c ∘ FSc`; `bad_keyc` (bad to bad), `keyc_S0c` (onto the bad set),
  `T0c_keyc` (the corrected pull is `T0c`).
* `T0c_XLc`: with the last selected bit (`flip1`) the composite is the row addition `RAc`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open Butterfly

/-! ## Bits -/

/-- The selected mask `Σ_{i<n} 2^{ρ + iK}`. -/
def selM (K ρ n : ℕ) : ℕ := ((List.range n).map fun i => 2 ^ (ρ + i * K)).sum

/-! ## `T0` is an XOR -/

/-- `T0` with `n` segments. -/
def T0n (K ρ n : ℕ) (x y : ℤ) : ℤ :=
  y + ∑ i ∈ Finset.range n, (2 : ℤ) ^ (ρ + i * K) *
    ((if ctl K ρ x i then 1 else 0) * (1 - 2 * bitAt y (ρ + i * K)))

/-! ## Canonical triple maps -/

section Canon

variable (K ρ f : ℕ)

/-- The rotation amount of line `l` (target value moves by `-(sgn · off)` in the pull map). -/
def amt (l : Line) (x t : ℕ) : ℕ := ((l.sgn * Butterfly.offset K ρ f x t l.b) % 2 ^ (f * K)).toNat

/-- The pull map of line `l` on a canonical triple `(x, y, z)`. -/
def lineP (l : Line) (c : Tr) : Tr :=
  if l.tgtY then (c.1, (c.2.1 + 2 ^ (f * K) - amt K ρ f l c.1 c.2.2) % 2 ^ (f * K), c.2.2)
  else (c.1, c.2.1, (c.2.2 + 2 ^ (f * K) - amt K ρ f l c.1 c.2.1) % 2 ^ (f * K))

/-- The composite pull map of the eight lines (line `l₁` executed first). -/
def FSc (c : Tr) : Tr := lines.foldr (fun l c => lineP K ρ f l c) c

/-- The forward map `S0` on canonical triples. -/
def S0c (c : Tr) : Tr :=
  (c.1, (S0 K ρ f c.1 (c.2.1, c.2.2)).1.toNat, (S0 K ρ f c.1 (c.2.1, c.2.2)).2.toNat)

/-- The ideal map `T0` on canonical triples. -/
def T0c (c : Tr) : Tr := (c.1, c.2.1 ^^^ (c.1 &&& selM K ρ (f - 1)), c.2.2)

/-- Outside the guard set `B`. -/
def good (c : Tr) : Prop := GuardOK K ρ f (c.2.1 : ℤ) ∧ GuardOK K ρ f (c.2.2 : ℤ)

/-- The destination of a bad record in the correction. -/
def keyc (c : Tr) : Tr := T0c K ρ f (FSc K ρ f c)

/-- The full row addition `y ← y ⊕ (x ∧ selM f)`. -/
def RAc (c : Tr) : Tr := (c.1, c.2.1 ^^^ (c.1 &&& selM K ρ f), c.2.2)

end Canon

section Facts

variable {K ρ f : ℕ}

/-- Negated line. -/
def negL (l : Line) : Line := ⟨l.tgtY, -l.sgn, l.b⟩

/-- Valid pairs. -/
def VP (Q : ℤ) (p : ℤ × ℤ) : Prop := 0 ≤ p.1 ∧ p.1 < Q ∧ 0 ≤ p.2 ∧ p.2 < Q

/-- The composite of the negated lines (line `l₁` outermost). -/
def FSb (K ρ f : ℕ) (x : ℤ) (p : ℤ × ℤ) : ℤ × ℤ := lines.foldr (fun l p => pline K ρ f x (negL l) p) p

/-- Triple from a pair. -/
def ofP (x : ℕ) (p : ℤ × ℤ) : Tr := (x, p.1.toNat, p.2.toNat)
/-- Pair of a triple. -/
def nmP (c : Tr) : ℤ × ℤ := ((c.2.1 : ℤ), (c.2.2 : ℤ))

/-! ### `T0` and the guard set -/

/-! ### The last selected bit -/

/-- The pull map of the two-bit XOR on the last selected position. -/
def XLc (K ρ f : ℕ) : Tr → Tr :=
  flip1 (2 ^ (ρ + (f - 1) * K)) (fun x => gc 2 (2 ^ (ρ + (f - 1) * K)) x)

end Facts

end NLogN.KFL

end

section
-- Source module: Solutions.KFLLoop
/-!
# KFL §3: the segment loop (register programs)

All per-record / per-block arithmetic of Lemma 13 is a loop over the `n` segments
`j_i = ρ + iK`: with `u, v` the two chunk values shifted by `2^ρ`, `w = 2^ρ`,

  `segBody K T : acc += w · T(u, v); u /= 2^K; v /= 2^K; w *= 2^K`

repeated `n` times (`segLoop`).  `segLoop_fsp` gives the closed form (for any expression `T` of
the two registers `rU, rV` with semantics `tf`) and the fine cost `O(n (|T| + 1) (Bb+2)^2)`.
Instances: the offset of a line (`offE`), the `T0` correction (`t0E`), the guard test (`grdE`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD

/-! ## Registers -/

abbrev rA : Fin HS := 30
abbrev rC : Fin HS := 31
abbrev rU : Fin HS := 50
abbrev rV : Fin HS := 51
abbrev rW : Fin HS := 52

/-! ## The loop -/

def segBody (K : ℕ) (T : Ex HS) : RS :=
  .seq (.asg rA (.add (.reg rA) (.mul (.reg NLogN.KFL.rW) T)))
    (.seq (.asg NLogN.KFL.rU (.quo (.reg NLogN.KFL.rU) (.lit (2 ^ K))))
    (.seq (.asg NLogN.KFL.rV (.quo (.reg NLogN.KFL.rV) (.lit (2 ^ K))))
      (.asg NLogN.KFL.rW (.mul (.reg NLogN.KFL.rW) (.lit (2 ^ K))))))

def segLoop (K n : ℕ) (T : Ex HS) : RS := .rep rC (.lit (n : ℤ)) (segBody K T)

/-- Partial sums of the loop. -/
def segAcc (K : ℕ) (tf : ℤ → ℤ → ℤ) (u v w : ℤ) (k : ℕ) : ℤ :=
  ∑ i ∈ Finset.range k, w * 2 ^ (i * K) * tf (u / 2 ^ (i * K)) (v / 2 ^ (i * K))

/-- The managed state after `k` iterations. -/
def segF (F : Fin HS → List ℤ) (K : ℕ) (tf : ℤ → ℤ → ℤ) (a u v w : ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update F
    rA ((a + segAcc K tf u v w k) :: (F rA).tail))
    NLogN.KFL.rU ((u / 2 ^ (k * K)) :: (F NLogN.KFL.rU).tail))
    NLogN.KFL.rV ((v / 2 ^ (k * K)) :: (F NLogN.KFL.rV).tail))
    NLogN.KFL.rW ((w * 2 ^ (k * K)) :: (F NLogN.KFL.rW).tail)

/-- Hypotheses on the expression `T`: it reads `rU, rV` only, has semantics `tf` (|tf| ≤ 2),
and its evaluation is executable and bounded by `Bb` bits when `u, v` are. -/
structure SegT (T : Ex HS) (tf : ℤ → ℤ → ℤ) (Bb : ℕ) : Prop where
  regs : ∀ i ∈ T.regs, i = NLogN.KFL.rU ∨ i = NLogN.KFL.rV
  ev : ∀ (t : IState HS) (a b : ℤ), t.top NLogN.KFL.rU = a → t.top NLogN.KFL.rV = b → T.eval t = tf a b
  ok : ∀ (t : IState HS) (a b : ℤ), t.top NLogN.KFL.rU = a → t.top NLogN.KFL.rV = b → t.IsInt NLogN.KFL.rU → t.IsInt NLogN.KFL.rV →
    0 ≤ a → 0 ≤ b → zsz a ≤ Bb → zsz b ≤ Bb → T.ok t ∧ T.Bnd t Bb
  abs : ∀ a b, |tf a b| ≤ 2

/-- Cost of one iteration. -/
def cSegIt (T : Ex HS) (Bb : ℕ) : ℕ := (T.size + 8) * 4 * 10 ^ 7 * (Bb + 2) ^ 2
end NLogN.KFL
end

section
-- Source module: Solutions.KFLTab
/-!
# KFL §5: the offset tables of the rotations (22)

* `offE b`, `tfB b`: the per-segment expression of `offset` (reads `rU, rV`), `segT_offE`.
* `offset_seg`: `offset K ρ f x t b = segAcc K (tfB b) (x / 2^ρ) (t / 2^ρ) (2^ρ) (f - 1)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly

/-- Segment test: `[u odd ∧ v ≡ b]` (for `b ∈ {0,1}`). -/
def tfB (b : ℤ) (u v : ℤ) : ℤ := (u % 2) * (if b = 1 then v % 2 else 1 - v % 2)

/-- The segment expression. -/
def offE (b : ℤ) : Ex HS :=
  if b = 1 then .mul (.nmod (.reg NLogN.KFL.rU) (.lit 2)) (.nmod (.reg NLogN.KFL.rV) (.lit 2))
  else .mul (.nmod (.reg NLogN.KFL.rU) (.lit 2)) (.add (.lit 1) (.neg (.nmod (.reg NLogN.KFL.rV) (.lit 2))))

/-! ## The table loop -/

/-- Outer counter of the table loop. -/
abbrev rq : Fin HS := 45

/-- Block counter read by the body (the loop counter itself may not be read). -/
abbrev rP : Fin HS := 46

/-- The chunk `(p / d) % Q` of block `p = N - 1`, shifted by `ρ`. -/
def uE (d Q ρ : ℕ) : Ex HS :=
  .ndiv (.nmod (.ndiv (.add (.reg NLogN.KFL.rP) (.lit (-1))) (.lit (d : ℤ))) (.lit (Q : ℤ))) (.lit (2 ^ ρ))

/-- `(sgn · rA) mod M` (computed as `(sgn · rA + M) mod M`). -/
def amtE (sgn : ℤ) (M : ℕ) : Ex HS :=
  .nmod (.add (.mul (.lit sgn) (.reg rA)) (.lit (M : ℤ))) (.lit (M : ℤ))

/-- One table entry: block `p = N - 1` (`N` on top of `rq`). -/
def tabBody (K ρ f d0 d1 Q : ℕ) (sgn b : ℤ) : RS :=
  .seq (.push NLogN.KFL.rU (uE d0 Q ρ)) (.seq (.push NLogN.KFL.rV (uE d1 Q ρ)) (.seq (.push NLogN.KFL.rW (.lit (2 ^ ρ)))
  (.seq (.push rA (.lit 0)) (.seq (segLoop K (f - 1) (offE b))
  (.seq (.push 14 (amtE sgn (2 ^ (f * K)))) (.seq (.pop rA) (.seq (.pop NLogN.KFL.rW)
  (.seq (.pop NLogN.KFL.rV) (.seq (.pop NLogN.KFL.rU) (.asg NLogN.KFL.rP (.add (.reg NLogN.KFL.rP) (.lit (-1)))))))))))))

/-- Cost of one table entry. -/
def cTab (K f : ℕ) (b : ℤ) (Bb : ℕ) : ℕ :=
  cLp (f - 1) (cSegIt (offE b) Bb) Bb + 50 * 10 ^ 7 * (Bb + 2) ^ 2

/-- **The table generator**: pushes `as (P-1), …, as 0` on stack 14 (block counter `rP`). -/
def tabGen (K ρ f d0 d1 Q P : ℕ) (sgn b : ℤ) : RS :=
  .seq (.push NLogN.KFL.rP (.lit (P : ℤ))) (.seq (.rep NLogN.KFL.rq (.lit (P : ℤ)) (tabBody K ρ f d0 d1 Q sgn b))
    (.pop NLogN.KFL.rP))

/-- The table entries of line `l`: block `p` with controls `(p / d0) % Q`, `(p / d1) % Q`. -/
def tabA (K ρ f d0 d1 Q : ℕ) (l : Line) (p : ℕ) : ℕ := amt K ρ f l (p / d0 % Q) (p / d1 % Q)

/-- The partial table after `k` iterations. -/
def tabL (as : ℕ → ℕ) (P k : ℕ) : List ℤ := (List.range k).map fun i => (as (P - k + i) : ℤ)

/-- The cost of a table. -/
def cTabGen (K f : ℕ) (b : ℤ) (P Bb : ℕ) : ℕ := cLp P (NLogN.KFL.cTab K f b Bb) Bb + 3 * 10 ^ 7 * (Bb + 2) ^ 2
end NLogN.KFL
end

section
-- Source module: Solutions.KFLSwap
/-!
# KFL §4: the Lemma 12 swap wrapper

`swI q vP vG vS u` pushes the four Lemma 12 parameters (`kP, kG, kS, aU`), runs the fine
Lemma 12 program `l12F`, and pops them.  `swI_spec`: on any state whose stack 0 is
`enc2 xs ++ R` (arbitrary tail `R`, no other hypothesis) it interchanges the two `[2^u]`
fields of `binShape Pre G Suf (2^u)` at cost `≤ K (V + 1) u^τ`, `τ = 1 - 2^-50`.

Then `swI_rixT`: on the 8-axis geometry `sh8 = [N0,Q,N1,Q,N2,Q,N3,B]`, `Q = 2^u`, the three
swaps (axes 1↔3, 3↔5, 1↔5) are `rixT g sw01 / sw12 / sw02`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD KFE KFB AModel AModelL12 NetProp7

/-- The program table of Lemma 12. -/
noncomputable abbrev l12prog : Fin 1 → IStmt HS 1 := progF ((bitQC NetProp7.ind₀).toCircuit 1) (edgeTplQ (bitQC NetProp7.ind₀))

/-- Lemma 12's code. -/
abbrev l12C (q : ℕ) : IStmt HS 1 := l12F q (100 ^ 3) (Fintype.card (NetScalar.Wr NetProp7.ind₀))

/-- The parameter pushes. -/
def swL (vP vG vS u : ℕ) : List (Fin HS × Ex HS) :=
  [(kP, .lit (vP : ℤ)), (kG, .lit (vG : ℤ)), (NLogN.KIL.kS, .lit (vS : ℤ)), (aU, .lit (u : ℤ))]

/-- **The swap wrapper.** -/
def swI (q vP vG vS u : ℕ) : IStmt HS 1 := withPushI (swL vP vG vS u) (l12C q)

/-- The managed lists inside the wrapper. -/
def Fsw (ys : List ℤ) (vP vG vS u : ℕ) : Fin HS → List ℤ := fun i =>
  if i = 0 then ys else if i = kP then [(vP : ℤ)] else if i = kG then [(vG : ℤ)] else
  if i = NLogN.KIL.kS then [(vS : ℤ)] else if i = aU then [(u : ℤ)] else []

/-! ## Bit lists and the three geometry swaps -/

/-- A list of bits `0/1`. -/
def Bits (zs : List ℤ) : Prop := ∀ z ∈ zs, z = 0 ∨ z = 1
end NLogN.KFL
end

section
-- Source module: Solutions.KFLRot
/-!
# KFL §6: one line (22) as an item program

`rotStep … l` pushes `P, Q, B` on stacks 1, 2, 4, generates the table of line `l`
(`tabGen`, stack 14) and runs the fine block rotation `KFC.rotF` (which consumes the table).
`rotStep_runs`: the result is the rotation pull `pull [P,Q,B] [P,Q,B] (rotAddr Q (tabA …))`.
`rotStep_geo`: on the 8-axis geometry with `d0 = N2·Q·N1`, `d1 = N2`, this is
`rixT g (rot2 Q (amt K ρ f l))`, i.e. the line acting on the third slot `z` with controls
`(x, y)` (canonical slots).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option maxRecDepth 20000

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly AModel

section
variable {Pp : ℕ} {I : List Sym} {iprog : Fin Pp → IStmt HS Pp}

/-- Shape pushes for `rotF`. -/
def rotL (P Q B : ℕ) : List (Fin HS × Ex HS) := [(1, .lit (P : ℤ)), (2, .lit (Q : ℤ)), (4, .lit (B : ℤ))]

/-- **One line**: table, then block rotation. -/
def rotStep (K ρ f d0 d1 Q P B : ℕ) (sgn b : ℤ) : IStmt HS Pp :=
  KFD.withPushI (rotL P Q B) (.seq (fcode (tabGen K ρ f d0 d1 Q P sgn b)) KFC.rotF)

/-- Its cost. -/
def cRot (K f : ℕ) (b : ℤ) (P Q B w Bb : ℕ) : ℕ :=
  cTabGen K f b P Bb + KFC.cRotF P Q B w + 10 ^ 8 * (Bb + 2) ^ 2

end
end NLogN.KFL
end

section
-- Source module: Solutions.KFLLines
/-!
# KFL §7: the eight packed lines (22) on the canonical layout

On the 8-axis geometry `sh8 = [N0,Q,N1,Q,N2,Q,N3,B]` (`Q = 2^{fK}`; slots `x, y, z` = axes
1, 3, 5) a `tgtZ` line is one `rotStep` (rotation of slot `z` by `amt l x y`), a `tgtY` line is
`swap(3,5); rotStep; swap(3,5)` (Lemma 12 swaps).  `linesI_spec`: running the lines `ls` in
order gives `rixT g (ls.foldr (lineP K ρ f))`; for `ls = lines` this is `rixT g (FSc K ρ f)`
(`KFLMath`: `FSc` is the pull map of `S0`).  Cost: per line `≤ 2 K_s (V+1) (fK)^τ + cRot`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option maxRecDepth 20000

namespace NLogN.KFL

open IntMul RS KIL KFD KFE KFB Butterfly AModel

/-- The rotation step of line `l` on the geometry. -/
def rotG (K ρ f : ℕ) (g : Geo) (l : Line) : IStmt HS 1 :=
  NLogN.KFL.rotStep K ρ f (g.N2 * g.Q * g.N1) g.N2 g.Q (g.N0 * g.Q * g.N1 * g.Q * g.N2) g.S2 l.sgn l.b

/-- The slot swap `y ↔ z` (axes 3 ↔ 5). -/
def swYZ (q K f : ℕ) (g : Geo) : IStmt HS 1 :=
  swI q (vol [g.N0, g.Q, g.N1]) g.N2 (vol [g.N3, g.B]) (f * K)

/-- **One line.** -/
def lineI (q K ρ f : ℕ) (g : Geo) (l : Line) : IStmt HS 1 :=
  if l.tgtY then .seq (swYZ q K f g) (.seq (NLogN.KFL.rotG K ρ f g l) (swYZ q K f g)) else NLogN.KFL.rotG K ρ f g l

/-- **A sequence of lines.** -/
def linesI (q K ρ f : ℕ) (g : Geo) : List Line → IStmt HS 1
  | [] => .skip
  | l :: ls => .seq (lineI q K ρ f g l) (linesI q K ρ f g ls)

/-- Lines with admissible parameters. -/
def LineOK (l : Line) : Prop := (l.sgn = 1 ∨ l.sgn = -1) ∧ (l.b = 0 ∨ l.b = 1)

/-- Per-line cost bound of the rotation (both values of `b`). -/
def cRotG (K f : ℕ) (g : Geo) (Bb : ℕ) : ℕ :=
  NLogN.KFL.cRot K f 0 (g.N0 * g.Q * g.N1 * g.Q * g.N2) g.Q g.S2 1 Bb +
    NLogN.KFL.cRot K f 1 (g.N0 * g.Q * g.N1 * g.Q * g.N2) g.Q g.S2 1 Bb
end NLogN.KFL
end