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

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

import Definitions.Def_IntMul_Kappa_K_Defs15

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFLXor
/-!
# KFL §8: the last selected bit (two-bit XOR)

`xorStep`: table `b(c₀) = bit j of c₀` (`j = ρ + (f-1)K`, `xorGen`), then `rotF` on the view
`[N0·Q·N1·E, 2, D·S1]` (`D = 2^j`, `Q = 2 D E`): the bit `j` of the `y` chunk is flipped iff bit
`j` of `x` is set.  `xorStep_geo`: the result is `rixT g (XLc K ρ f)`.
-/

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

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly AModel

/-- `((p / d) % Q / 2^j) % 2` for `p = N - 1`. -/
def bE (d Q j : ℕ) : Ex HS := .nmod (uE d Q j) (.lit 2)

/-- One table entry. -/
def xorBody (d Q j : ℕ) : RS :=
  .seq (.push 14 (bE d Q j)) (.asg NLogN.KFL.rP (.add (.reg NLogN.KFL.rP) (.lit (-1))))

/-- The bit table: `bit_j ((p / d) % Q)` for `p = P-1, …, 0`. -/
def xorGen (d Q j P : ℕ) : RS :=
  .seq (.push NLogN.KFL.rP (.lit (P : ℤ))) (.seq (.rep NLogN.KFL.rq (.lit (P : ℤ)) (xorBody d Q j)) (.pop NLogN.KFL.rP))

/-- Table entries. -/
def xorA (d Q j : ℕ) (p : ℕ) : ℕ := p / d % Q / 2 ^ j % 2

def cXorGen (P Bb : ℕ) : ℕ := cLp P (20 * 10 ^ 7 * (Bb + 2) ^ 2) Bb + 3 * 10 ^ 7 * (Bb + 2) ^ 2

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

/-- **The last-bit step**: bit table, then `rotF` on `[P', 2, B']`. -/
def xorStep (d Q j P B : ℕ) : IStmt HS Pp :=
  KFD.withPushI (rotL P 2 B) (.seq (fcode (xorGen d Q j P)) KFC.rotF)

def cXor (P B w Bb : ℕ) : ℕ := cXorGen P Bb + KFC.cRotF P 2 B w + 10 ^ 8 * (Bb + 2) ^ 2

end

/-! ## The pipeline on the canonical layout -/

/-- The correction map: position `c` receives the record at `S0c (T0c c)`. -/
def corrc (K ρ f : ℕ) (c : Tr) : Tr := S0c K ρ f (T0c K ρ f c)
end NLogN.KFL
end

section
-- Source module: Solutions.KFLRec
/-!
# KFL §9: records, bad records and their number

* `Geo.recG g` is the record geometry (`B = 1`): record ranks `m < g.recG.V`, bit index `m B + b`.
  `phi_mul`: `phi g F (m B + b) = phi g.recG F m · B + b`; `rixT_flR`: a triple map acts on the
  bit array `flR recs M` record-wise.
* `badR`, the records with a failing guard; `corrR`/`keyR` (the correction and its inverse on
  ranks), `badR_corrR`, `corrR_keyR`, `keyR_corrR`, `corrR_good`.
* `card_badR_le`: `#bad · 2^K ≤ 80 (f-1) M` (from `card_bad_le`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open AModel Butterfly

/-- The record geometry. -/
def Geo.recG (g : Geo) : Geo := { g with B := 1 }

namespace Geo
variable (g : Geo)
@[simp] theorem recG_Q : g.recG.Q = g.Q := rfl
@[simp] theorem recG_N0 : g.recG.N0 = g.N0 := rfl
@[simp] theorem recG_N1 : g.recG.N1 = g.N1 := rfl
@[simp] theorem recG_N2 : g.recG.N2 = g.N2 := rfl
@[simp] theorem recG_N3 : g.recG.N3 = g.N3 := rfl
@[simp] theorem recG_B : g.recG.B = 1 := rfl

end Geo

/-! ## Record arrays -/

/-- The bit array of `M` records. -/
def flR (recs : ℕ → List ℤ) (M : ℕ) : List ℤ := (List.range M).flatMap recs

/-! ## Bad records -/

section Bad

variable {K ρ f : ℕ}

/-- Record `m` is bad. -/
def badR (K ρ f : ℕ) (g : Geo) (m : ℕ) : Prop := ¬ good K ρ f (trOf g.recG m)

/-- The correction on ranks: destination `m` receives source `corrR m`. -/
def corrR (K ρ f : ℕ) (g : Geo) (m : ℕ) : ℕ := phi g.recG (corrc K ρ f) m

/-- The destination of source `m`. -/
def keyR (K ρ f : ℕ) (g : Geo) (m : ℕ) : ℕ := phi g.recG (keyc K ρ f) m

variable {g : Geo} (hg : g.Pos) (hQe : g.Q = 2 ^ (f * K))
include hg hQe

variable (hK : 6 ≤ K) (hρK : ρ < K) (hf : 1 ≤ f)
include hK hρK hf

end Bad
end NLogN.KFL
end

section
-- Source module: Solutions.KFLKey
/-!
# KFL §10: the destination key of a record, in registers

* `segStep`: push `u, v, w, a` on `rU, rV, rW, rA`, run the segment loop, assign a function of
  the accumulator to a target register, pop.  `segStep_fsp` is the generic specification.
* `lineR l`: one packed line (22) on the registers `rX, rY, rZ` (`lineP` on the triple).
* `t0R`: the ideal map `T0` on `rY` (an XOR with `x ∧ selM (f-1)`).
* `keyProg`: from the rank `m` on `rm` push the key `keyR m` on `rK`
  (`x, y, z` by division, the eight lines in `foldr` order, `T0`, and the rank update).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly

/-! ## Expression specifications -/

/-- An expression is executable on `t`, bounded by `Bb` bits and has value `v`. -/
structure ESp (e : Ex HS) (t : IState HS) (Bb : ℕ) (v : ℤ) : Prop where
  clean : e.Clean
  ok : e.ok t
  bnd : e.Bnd t Bb
  ev : e.eval t = v

/-- Reads no register of the segment loop. -/
def NoSeg (e : Ex HS) : Prop := ∀ i ∈ e.regs, i ≠ rA ∧ i ≠ NLogN.KFL.rU ∧ i ≠ rV ∧ i ≠ rW

/-! ## The generic segment step -/

def segStep (K n : ℕ) (w : ℤ) (eu ev ea T : Ex HS) (tgt : Fin HS) (efin : Ex HS) : RS :=
  .seq (.push NLogN.KFL.rU eu) (.seq (.push rV ev) (.seq (.push rW (.lit w)) (.seq (.push rA ea)
  (.seq (segLoop K n T) (.seq (.asg tgt efin) (.seq (.pop rA) (.seq (.pop rW)
  (.seq (.pop rV) (.pop NLogN.KFL.rU)))))))))

def cStep (K n : ℕ) (T : Ex HS) (Bb c : ℕ) : ℕ := cLp n (cSegIt T Bb) Bb + (c + 20) * 10 ^ 7 * (Bb + 2) ^ 2

/-! ## Expression combinators -/

section ESpC

variable {t : IState HS} {Bb : ℕ}

end ESpC

/-! ## Registers of the key computation -/

abbrev rm : Fin HS := 33
abbrev rX : Fin HS := 34
abbrev rY : Fin HS := 35
abbrev rZ : Fin HS := 36
abbrev rK : Fin HS := 37
abbrev rY0 : Fin HS := 38
abbrev rZ0 : Fin HS := 39

/-- The triple `c` on top of `rX, rY, rZ` over the base `H`. -/
def trF (H : Fin HS → List ℤ) (c : Tr) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update H rX ((c.1 : ℤ) :: H rX))
    rY ((c.2.1 : ℤ) :: H rY)) rZ ((c.2.2 : ℤ) :: H rZ)

/-! ## Size facts -/

/-! ## One packed line on the registers -/

def tgtR (l : Line) : Fin HS := if l.tgtY then rY else rZ
def othR (l : Line) : Fin HS := if l.tgtY then rZ else rY

def lineR (K ρ f : ℕ) (l : Line) : RS :=
  segStep K (f - 1) (2 ^ ρ) (.ndiv (.reg rX) (.lit (2 ^ ρ))) (.ndiv (.reg (othR l)) (.lit (2 ^ ρ)))
    (.lit 0) (offE l.b) (tgtR l)
    (.nmod (.add (.add (.reg (tgtR l)) (.lit (2 ^ (f * K)))) (.neg (.mul (.lit l.sgn) (.reg rA))))
      (.lit (2 ^ (f * K))))

def cLine (K f : ℕ) (b : ℤ) (Bb : ℕ) : ℕ := cStep K (f - 1) (offE b) Bb 17

/-! ## The ideal map `T0` on the registers -/

/-- Segment term of `T0`: `[u odd] (1 - 2 [v odd])`. -/
def tf0 (u v : ℤ) : ℤ := (u % 2) * (1 - 2 * (v % 2))

def t0E : Ex HS :=
  .mul (.nmod (.reg NLogN.KFL.rU) (.lit 2)) (.add (.lit 1) (.neg (.mul (.lit 2) (.nmod (.reg rV) (.lit 2)))))

def t0R (K ρ f : ℕ) : RS :=
  segStep K (f - 1) (2 ^ ρ) (.ndiv (.reg rX) (.lit (2 ^ ρ))) (.ndiv (.reg rY) (.lit (2 ^ ρ)))
    (.reg rY) t0E rY (.reg rA)

def cT0 (K f Bb : ℕ) : ℕ := cStep K (f - 1) t0E Bb 8

/-! ## The eight lines -/

def progL (K ρ f : ℕ) : List Line → RS
  | [] => .skip
  | l :: ls => .seq (progL K ρ f ls) (lineR K ρ f l)

/-! ## The key of a rank -/

/-- `(m / S) % Q` from the rank register. -/
def gcE (S Q : ℕ) : Ex HS := .nmod (.ndiv (.reg NLogN.KFL.rm) (.lit (S : ℤ))) (.lit (Q : ℤ))

def keyE (S1 S2 : ℕ) : Ex HS :=
  .add (.reg NLogN.KFL.rm) (.add (.mul (.add (.reg rY) (.neg (.reg rY0))) (.lit (S1 : ℤ)))
    (.mul (.add (.reg rZ) (.neg (.reg rZ0))) (.lit (S2 : ℤ))))

def keyProg (K ρ f : ℕ) (g : Geo) : RS :=
  .seq (.push rY0 (gcE g.S1 g.Q)) (.seq (.push rZ0 (gcE g.S2 g.Q)) (.seq (.push rX (gcE g.S0 g.Q))
  (.seq (.push rY (gcE g.S1 g.Q)) (.seq (.push rZ (gcE g.S2 g.Q)) (.seq (progL K ρ f lines)
  (.seq (t0R K ρ f) (.seq (.push NLogN.KFL.rK (keyE g.S1 g.S2)) (.seq (.pop rX) (.seq (.pop rY)
  (.seq (.pop rZ) (.seq (.pop rY0) (.pop rZ0))))))))))))

def cKey (K f Bb : ℕ) : ℕ := 1 + 8 * cLine K f 0 Bb + cT0 K f Bb + 200 * 10 ^ 7 * (Bb + 2) ^ 2

end NLogN.KFL

end

section
-- Source module: Solutions.KFLSort
/-!
# KFL §11: the list model of the radix sort

Entries are pairs `(key, record)`.  `rpass i` is one stable LSD pass on bit `i`, *ones first*
(this is what two stacks give: the 0-bucket is moved back first, the 1-bucket on top of it);
`rsort A` runs the passes `0, …, A-1`.  `rsort_eq`: with keys below `2^A` and distinct, the
result is the unique list sorted by strictly decreasing key.

`sort_moved`: the entries extracted by the first pass of the correction (the moved records
`κ m ≠ m`, keyed by their destination `κ m`) sort into destination order, i.e. the record of
source `γ d` for every moved destination `d`, in decreasing `d`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

/-- An entry: destination key and record. -/
abbrev Ent := ℕ × List ℤ

def bitP (i : ℕ) (e : Ent) : Bool := decide (e.1 / 2 ^ i % 2 = 1)

/-- One radix pass: the entries with bit `i` set first, both parts in their old order. -/
def rpass (i : ℕ) (E : List Ent) : List Ent := E.filter (bitP i) ++ E.filter (fun e => !bitP i e)

def rsort : ℕ → List Ent → List Ent
  | 0, E => E
  | i + 1, E => rpass i (rsort i E)

/-- Order on the low `i` bits of the keys (decreasing). -/
def RO (i : ℕ) (a b : Ent) : Prop := b.1 % 2 ^ i ≤ a.1 % 2 ^ i

/-! ## The moved records -/

section Moved

variable (κ γ : ℕ → ℕ) (M : ℕ)

/-- The moved ranks below `M`, in decreasing order. -/
def mvL : List ℕ := (List.range M).reverse.filter fun m => decide (κ m ≠ m)

/-- The extracted entries (top first): destination key and the (reversed) record. -/
def ent0 (recs : ℕ → List ℤ) : List Ent := (NLogN.KFL.mvL κ M).map fun m => (κ m, (recs m).reverse)

/-- The sorted entries: for every moved destination `d` (decreasing), the record of `γ d`. -/
def ent1 (recs : ℕ → List ℤ) : List Ent := (NLogN.KFL.mvL κ M).map fun d => (d, (recs (γ d)).reverse)

variable {κ γ M}

end Moved

end NLogN.KFL

end

section
-- Source module: Solutions.KFLCorr1
/-!
# KFL §12: the correction, pass 1 (extraction of the moved records)

The bit array `flR recs M` (record `m` = `recs m`, `B` bits) is on stack `0`, record `0` on top.
For `m = 0, …, M-1` the pass computes the key `κ m = keyR m` (`keyProg`), and moves the record
* to `sKB` (and pushes `κ m` on `sKK`) if `κ m ≠ m` (a *moved* record),
* to `sT` otherwise,
and pushes the flag `[κ m ≠ m]` on `sFL`; `rNB` counts the moved records.
`p1_fsp`: the whole pass, as the list state `P1 … M`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly

/-! ## Stacks and registers of the correction -/

abbrev sT : Fin HS := 16
abbrev sKB : Fin HS := 17
abbrev sKK : Fin HS := 18
abbrev sFL : Fin HS := 19
abbrev sZ0B : Fin HS := 20
abbrev sZ0K : Fin HS := 21
abbrev sZ1B : Fin HS := 22
abbrev sZ1K : Fin HS := 23
abbrev rNB : Fin HS := 40
abbrev cMv : Fin HS := 41
abbrev c1 : Fin HS := 42
abbrev c2 : Fin HS := 43
abbrev c3 : Fin HS := 44
abbrev c4 : Fin HS := 47
abbrev rn0 : Fin HS := 48
abbrev rn1 : Fin HS := 49
abbrev rn0B : Fin HS := 53
abbrev rn1B : Fin HS := 54

/-! ## List helpers -/

/-- The records `F m` of the ranks `m < k` with `p m`, the largest on top. -/
def stkF (p : ℕ → Bool) (F : ℕ → List ℤ) (k : ℕ) : List ℤ :=
  ((List.range k).reverse.filter p).flatMap F

/-- The keys `G m` of the ranks `m < k` with `p m`, the largest on top. -/
def stkK (p : ℕ → Bool) (G : ℕ → ℤ) (k : ℕ) : List ℤ := ((List.range k).reverse.filter p).map G

/-- The rest of the input after `k` records. -/
def restR (recs : ℕ → List ℤ) (M k : ℕ) : List ℤ := ((List.range M).drop k).flatMap recs

/-! ## The pass-1 state -/

section P1

variable (κ : ℕ → ℕ) (recs : ℕ → List ℤ) (M : ℕ)

/-- `κ m ≠ m`: record `m` is moved. -/
def mvP (m : ℕ) : Bool := decide (κ m ≠ m)

def flagV (m : ℕ) : ℤ := if κ m ≠ m then 1 else 0

/-- The flags of the ranks `< k`, the largest on top. -/
def flL (k : ℕ) : List ℤ := (List.range k).reverse.map (flagV κ)

/-- The number of moved ranks `< k`. -/
def nMv (k : ℕ) : ℕ := ((List.range k).filter (mvP κ)).length

/-- The state after `k` records of pass 1. -/
def P1 (G0 : Fin HS → List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
  (Function.update (Function.update G0
    0 (restR recs M k))
    sT (stkF (fun m => !mvP κ m) (fun m => (recs m).reverse) k ++ G0 sT))
    sKB (stkF (mvP κ) (fun m => (recs m).reverse) k ++ G0 sKB))
    sKK (stkK (mvP κ) (fun m => (κ m : ℤ)) k ++ G0 sKK))
    sFL (flL κ k ++ G0 sFL))
    NLogN.KFL.rm ((k : ℤ) :: G0 NLogN.KFL.rm))
    rNB ((nMv κ k : ℤ) :: G0 rNB)

end P1

/-! ## The branches -/

def badB (B : ℕ) : RS :=
  .seq (NLogN.mvK 0 sKB NLogN.KFL.cMv (.lit (B : ℤ))) (.seq (.push sKK (.reg NLogN.KFL.rK))
    (.seq (.push sFL (.lit 1)) (.asg rNB (.add (.reg rNB) (.lit 1)))))

def goodB (B : ℕ) : RS := .seq (NLogN.mvK 0 sT NLogN.KFL.cMv (.lit (B : ℤ))) (.push sFL (.lit 0))

def cBr (B Bb : ℕ) : ℕ := cLp B (250 * (1 + 2)) Bb + 10 ^ 9 * (Bb + 2) ^ 2

/-! ## One record of pass 1 -/

def e1E : Ex HS := .add (.reg NLogN.KFL.rK) (.neg (.reg NLogN.KFL.rm))
def e2E : Ex HS := .add (.reg NLogN.KFL.rm) (.neg (.reg NLogN.KFL.rK))

def brB (B : ℕ) : RS := .ite e1E (badB B) (.ite e2E (badB B) (goodB B))

def p1B (K ρ f : ℕ) (g : Geo) : RS :=
  .seq (keyProg K ρ f g.recG) (.seq (brB g.B) (.seq (.pop NLogN.KFL.rK) (.asg NLogN.KFL.rm (.add (.reg NLogN.KFL.rm) (.lit 1)))))

def cIte (Bb : ℕ) : ℕ := 5 * 2000 * (Bb + 2) + (2 * Bb + 2) + 25

def cP1 (K f B Bb : ℕ) : ℕ := cKey K f Bb + 2 * cIte Bb + cBr B Bb + 10 ^ 4 * (Bb + 2) ^ 2

def p1Prog (K ρ f : ℕ) (g : Geo) : RS := .rep NLogN.KFL.c1 (.lit (g.recG.V : ℤ)) (p1B K ρ f g)

end NLogN.KFL

end

section
-- Source module: Solutions.KFLCorr2
/-!
# KFL §13: the correction, the radix sort on the stacks

The extracted entries `E` (keys on `sKK`, records on `sKB`, entry `0` on top) are sorted by
`A` LSD passes.  A pass pops every entry and pushes it on the bucket of its low key bit
(`sZ0K/sZ0B` or `sZ1K/sZ1B`, the key halved), then moves bucket 0 and then bucket 1 back:
the result is `rpass i E` (`passB_fsp`).  `radix_fsp`: `A` passes give `rsort A E`.
All per-entry work is linear in the word size (`Sp.push`, `Sp.asg` on linear expressions).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly

/-! ## One entry -/

def bucket (ZK ZB n nB : Fin HS) (B : ℕ) : RS :=
  .seq (.push ZK (.quo (.reg sKK) (.lit 2))) (.seq (.pop sKK) (.seq (NLogN.mvK sKB ZB NLogN.KFL.cMv (.lit B))
    (.seq (.asg n (.add (.reg n) (.lit 1))) (.asg nB (.add (.reg nB) (.lit (B : ℤ)))))))

def cBk (B Bb : ℕ) : ℕ := cLp B (250 * (1 + 2)) Bb + 30000 * (Bb + 2)

def condE : Ex HS := .neg (.rem (.reg sKK) (.lit 2))

def entB (B : ℕ) : RS :=
  .ite condE (bucket sZ1K sZ1B rn1 rn1B B) (bucket sZ0K sZ0B rn0 rn0B B)

def cEnt (B Bb : ℕ) : ℕ := cBk B Bb + 2 * 10 ^ 4 * (Bb + 2)

/-- The key part on `sKK` during pass `i`. -/
def kh (i : ℕ) (e : Ent) : ℤ := ((e.1 / 2 ^ i : ℕ) : ℤ)

/-- The pass state: the rest `X` of the entries on `sKK/sKB`, the buckets `Y1, Y0` (last pushed
first) with their counts. -/
def RP (G : Fin HS → List ℤ) (B i : ℕ) (X Y1 Y0 : List Ent) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update
  (Function.update (Function.update (Function.update (Function.update (Function.update G
    sKK (X.map (kh i) ++ G sKK))
    sKB (X.flatMap Prod.snd ++ G sKB))
    sZ1K (Y1.map (kh (i + 1)) ++ G sZ1K))
    sZ1B (Y1.flatMap (fun e => e.2.reverse) ++ G sZ1B))
    sZ0K (Y0.map (kh (i + 1)) ++ G sZ0K))
    sZ0B (Y0.flatMap (fun e => e.2.reverse) ++ G sZ0B))
    rn1 ((Y1.length : ℤ) :: G rn1))
    rn1B (((Y1.length * B : ℕ) : ℤ) :: G rn1B))
    rn0 ((Y0.length : ℤ) :: G rn0))
    rn0B (((Y0.length * B : ℕ) : ℤ) :: G rn0B)

/-! ## The entry loop of a pass -/

def Yb (i : ℕ) (E : List Ent) (k : ℕ) (b : Bool) : List Ent :=
  ((E.take k).filter (fun e => bitP i e == b)).reverse

def entLoop (B : ℕ) : RS := .rep c3 (.reg rNB) (entB B)

/-! ## One pass -/

def passB (B : ℕ) : RS :=
  .seq (.push rn0 (.lit 0)) (.seq (.push rn1 (.lit 0)) (.seq (.push rn0B (.lit 0))
  (.seq (.push rn1B (.lit 0)) (.seq (entLoop B)
  (.seq (NLogN.mvK sZ0B sKB NLogN.KFL.cMv (.reg rn0B)) (.seq (NLogN.mvK sZ0K sKK NLogN.KFL.cMv (.reg rn0))
  (.seq (NLogN.mvK sZ1B sKB NLogN.KFL.cMv (.reg rn1B)) (.seq (NLogN.mvK sZ1K sKK NLogN.KFL.cMv (.reg rn1))
  (.seq (.pop rn0) (.seq (.pop rn1) (.seq (.pop rn0B) (.pop rn1B))))))))))))

def cPass (B N Bb : ℕ) : ℕ :=
  cLp N (cEnt B Bb) Bb + 2 * cLp (N * B) (250 * (1 + 2)) Bb + 2 * cLp N (250 * (Bb + 2)) Bb +
    10 ^ 5 * (Bb + 2)

/-! ## The sort -/

def radixProg (B A : ℕ) : RS := .rep c2 (.lit (A : ℤ)) (passB B)

end NLogN.KFL

end

section
-- Source module: Solutions.KFLCorr3
/-!
# KFL §14: the correction, pass 3 (merge) and the whole correction

Pass 3 rebuilds stack `0` from the bottom: for `d = M-1, …, 0` it pops the flag of `d` and moves
the top record of `sKB` (a moved destination: the sorted entries are in decreasing destination
order) or of `sT` (an unmoved record) to stack `0`.

`corrRS` = pass 1 + radix sort + pass 3; `corrRS_fsp`: on `flR recs M` it produces
`flR (recs ∘ corrR) M`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly

/-! ## Pass 3 -/

def p3B (B : ℕ) : RS :=
  .seq (.ite (.neg (.reg sFL)) (.seq (NLogN.mvK sKB 0 NLogN.KFL.cMv (.lit (B : ℤ))) (.pop sKK))
    (NLogN.mvK sT 0 NLogN.KFL.cMv (.lit (B : ℤ)))) (.pop sFL)

def cP3 (B Bb : ℕ) : ℕ := cLp B (250 * (1 + 2)) Bb + 10 ^ 4 * (Bb + 2)

section P3

variable (κ γ : ℕ → ℕ) (recs : ℕ → List ℤ) (M A : ℕ)

/-- The state of pass 3 when the destinations `≥ k` are done. -/
def P3 (G0 : Fin HS → List ℤ) (k : ℕ) : Fin HS → List ℤ :=
  Function.update (Function.update (Function.update (Function.update (Function.update G0
    0 (restR (fun d => recs (γ d)) M k ++ G0 0))
    sT (stkF (fun m => !mvP κ m) (fun m => (recs m).reverse) k ++ G0 sT))
    sKB (stkF (mvP κ) (fun d => (recs (γ d)).reverse) k ++ G0 sKB))
    sKK (stkK (mvP κ) (fun d => ((d / 2 ^ A : ℕ) : ℤ)) k ++ G0 sKK))
    sFL (flL κ k ++ G0 sFL)

end P3

def p3Prog (B M : ℕ) : RS := .rep c4 (.lit (M : ℤ)) (p3B B)

/-! ## The whole correction -/

def corrRS (K ρ f : ℕ) (g : Geo) (A : ℕ) : RS :=
  .seq (.push NLogN.KFL.rm (.lit 0)) (.seq (.push rNB (.lit 0)) (.seq (p1Prog K ρ f g) (.seq (.pop NLogN.KFL.rm)
    (.seq (radixProg g.B A) (.seq (p3Prog g.B g.recG.V) (.pop rNB))))))

def cCorr (K f : ℕ) (g : Geo) (A nB Bb : ℕ) : ℕ :=
  10 ^ 5 * (Bb + 2) + cLp g.recG.V (cP1 K f g.B Bb) Bb + cLp A (cPass g.B nB Bb) Bb +
    cLp g.recG.V (cP3 g.B Bb) Bb

section Lists

variable (κ γ : ℕ → ℕ) (recs : ℕ → List ℤ) (M A : ℕ)

end Lists

end NLogN.KFL

end

section
-- Source module: Solutions.KFLCorr4
/-!
# KFL §15: the number of moved records and the IL correction step

`card_moved`: the records moved by the correction (`keyR m ≠ m`) are bad, and the bad records
number at most `80 (f-1) 2^{-K}` of all (two guards per index, twenty bad values of `2^{K-1}`
each, by `count_digit` / `count_bad_one`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD Butterfly Finset

section Count

variable {K ρ f : ℕ} {g : Geo}

end Count

/-! ## Cost bookkeeping -/

section Cost

end Cost

/-! ## Cutting a bit array into records -/

section Chunks

/-- Record `m` of an array of `M` records of `B` items (padding beyond `M`). -/
def chunk (B M : ℕ) (zs : List ℤ) (m : ℕ) : List ℤ :=
  if m < M then (zs.drop (m * B)).take B else List.replicate B 0

end Chunks

/-! ## The correction step in IL -/

/-- The IL code of the correction. -/
def corrI (K ρ f : ℕ) (g : Geo) : IStmt HS 1 := fcode (corrRS K ρ f g (zsz (g.V : ℤ)))

end NLogN.KFL

end

section
-- Source module: Solutions.KFMB0
/-!
# The permutation-free `B̃₀ = D̃' J̃' C̃`

`B0approx` is the base `Bapprox` (HvdH Proposition 4.7(ii)) without its two data permutations
`P̃_t` and `P̃_s⁻¹`: `Bapprox = P̃_s⁻¹ ∘ B0approx ∘ P̃_t` (`Bapprox_eq`). Lemma 19 (κ, §8.1) keeps
both permutations in the transform, so the machine applies `B0approx` (see `KFMR19`, `KFMBap0`).
-/

namespace NLogN.KFM

/-- `B̃₀ = D̃' J̃' C̃`: the base `Bapprox` without its two data permutations. -/
def B0approx (p s t α : ℕ) (W : ℕ → ℤ × ℤ) : ℕ → ℤ × ℤ :=
  IntMul.HvdH.Dapprox p s t α (IntMul.HvdH.Japprox p s t α (IntMul.HvdH.cSel s t W))

/-- `B̃₀` read periodically: index `l` is reduced mod `s` (the machine reads `l < s` only). -/
def B0P (p s t α : ℕ) (W : ℕ → ℤ × ℤ) (l : ℕ) : ℤ × ℤ := B0approx p s t α W (l % s)

end NLogN.KFM

end

section
-- Source module: NLogN.SSAlg
/-!
# Algebra of the Fermat-ring FFT: inversion and the convolution property
-/

namespace NLogN

open IntMul

section ring

variable (n : ℕ)

/-! ## The forward transform evaluates at roots of unity -/

/-- `∑_{i<m} A_i x^i`. -/
def peval (A : List ℤ) (x : (ZMod (2 ^ n + 1))) (m : ℕ) : (ZMod (2 ^ n + 1)) := ∑ i ∈ Finset.range m, (A.getD i 0 : (ZMod (2 ^ n + 1))) * x ^ i

end ring

section ring2

variable (n : ℕ)

end ring2

/-! ## Cyclic convolution and pointwise products -/

/-- Cyclic convolution of length `K`, reduced mod `F`. -/
def cycZ (F : ℤ) (K : ℕ) (A B : List ℤ) : List ℤ :=
  (List.range K).map fun m => (∑ i ∈ Finset.range K, ∑ l ∈ Finset.range K,
    if (i + l) % K = m then A.getD i 0 * B.getD l 0 else 0) % F

/-- Pointwise products, reduced mod `F`. -/
def pmulZ (F : ℤ) : List ℤ → List ℤ → List ℤ
  | a :: A, b :: B => (a * b) % F :: pmulZ F A B
  | _, _ => []

section ring3

variable (n : ℕ)

end ring3

end NLogN

end

section
-- Source module: Solutions.KFMBap0
/-!
# Item 12 (machine side, part 1): the permutation-free `B̃₀` slice program

Lemma 19 (κ, §8.1) keeps both permutations `P_s`, `P_t` in the transform. So the machine
applies `B̃₀ = D̃' J̃' C̃` (`B0approx`), not the base `B̃ = P̃_s⁻¹ B̃₀ P̃_t` that nlogn's `bapP`
computes. This file drops the two `permP` passes from `bapP`:

* `b0L`: the list form (`(List.range s).map (B0approx …)`), with `b0L_eq` (the pipeline
  `dapL ∘ japL ∘ cselL`) and the bound `b0L_LB`.
* `bapP0 = C̃-stage; JP; D̃'-stage` with `bapP0_run` and `bapP0_okc` (cost `bap0Cost`).
* `bap0SlP` (relocated by nlogn's `σP`, then transferred `bI → bO`) and `bap0Sl`, a
  `GL.LProg` in the same context `CtxB` as nlogn's `bapSl`.

The proofs follow nlogn's `HRspBap`/`HAx5` with the two permutation steps removed.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KFM

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.Rsp NLogN.AX

/-! ## The list form -/

/-- `B̃₀` on a slice of length `t`, producing length `s`. -/
def b0L (p s t α : ℕ) (V : List (ℤ × ℤ)) : List (ℤ × ℤ) :=
  (List.range s).map (B0approx p s t α fun j => V.getD j (0, 0))

/-! ## The slice program on `CU` -/

/-- **`B̃₀` on the slice stack**: `C̃` stage, `J̃'`, `D̃'` stage (no permutations). -/
def bapP0 : RM := RM.seqL [(stageP Ac).ren σc, JP, (stageP Ad).ren σd]

/-- Cost bound of `bapP0`: the `C̃` and `D̃'` stages and the `J̃'` loop. -/
def bap0Cost (p s t α mm Φ B : ℕ) : ℕ :=
  NLogN.Rsp.stCost t 1 s mm Φ B (t / s + 1) +
  jCost s (2 * HvdH.eWidth p α + 1) (HvdH.jIter p s t α) mm Φ B +
  NLogN.Rsp.stCost s 1 s mm Φ B 1

/-! ## The slice program on the `B̃` axis layout -/

/-- The `B̃₀` slice program: `bapP0` in place on `bI`, then transfer to `bO`. -/
def bap0SlP : RM := .seq (RM.ren σP bapP0) (xferP bI bO (cnt2 bS))

/-- Cost bound of one `B̃₀` slice. -/
def bap0SlC (p s t α mm Φ Bacc Q : ℕ) : ℕ :=
  bap0Cost p s t α mm Φ Bacc + (16000 * (Q + 2) + s * 60000 * (Q + 2) + 64)

end NLogN.KFM

end

section
-- Source module: Solutions.KFMBap0Ax
/-!
# Item 12 (machine side, part 2): the `B̃₀` axis program `bap0AxP`

nlogn's `bapAxP` (`HAx8`) with the slice program `bapSlP` replaced by the permutation-free
`bap0SlP` (`KFMBap0`). The table building is unchanged; the two key tables (`ptKR`, `psKR`) are
still built and removed but no longer read. That costs `O((t+s)·Bq²)` per axis call.
On `gD = gints L ++ R` (`L` of dims `t :: rest`) it computes
`gints (tensorGo [(s, b0L p s t α)] (t :: rest) L) ++ R`. The proof is nlogn's `bapAxP_all`
with the slice lemmas renamed. The declarations live in `NLogN.AX` (as in `HAx8`) so
that names resolve exactly as there.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.AX

open IntMul NLogN IntMul.HvdH.Ref NLogN.GL NLogN.Rsp NLogN.KFM

/-- The axis core (= `GL.axisP (bap0Ax …)`, written out). -/
def bap0CoreP : RM :=
  .seq (RM.seqL [.rs (.push 1 (.reg bT)), .rs (.push 2 (.reg bM)), .rs (.push 3 (.reg bKB)),
      .rs (transP 2), .rs (.pop 1), .rs (.pop 2), .rs (.pop 3)])
    (.seq (.rep gCC (.reg bM) (RM.seqL [.rs (mvList NLogN.GL.gD gZT gCV (cnt2 bT)),
        .rs (mvList gZT bI gCV (cnt2 bT)), bap0SlP, .rs (mvList bO gZB gCV (cnt2 bS))]))
      (.rs (mvList gZB NLogN.GL.gD gCV (cnt2 bTot))))

/-- **The `B̃₀` axis program.** -/
def bap0AxP : RM :=
  RM.seqL [.rs tmpP, .rs eTabR, .rs eDelR, .rs cTabR, .rs cDelR, .rs dTabR', .rs ptKR, .rs psKR,
    .rs tmpQ, bap0CoreP, .rs (popK bW gCC (.reg bNW)), .rs (popK bD gCC (.reg bS)),
    .rs (popK bW1 gCC (.reg bS)), .rs (popK bD1 gCC (.reg bS)), .rs (popK bW3 gCC (.reg bS)),
    .rs (popK bK3 gCC (.reg bT)), .rs (popK bKS gCC (.reg bS))]

/-- Cost bound of `bap0AxP`. -/
def bap0AxCost (p s t α m kb Bi mm Φ Bacc Q Bq : ℕ) : ℕ :=
  16 * 10 ^ 7 * (Bq + 2) ^ 2 +
  (tabCost (2 * HvdH.eWidth p α + 1) s (eWCost p Bq) +
  ((tabCost 1 s ((eDelE.size + 1) * 10 ^ 7 * (s + 4) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s (2 * 10 ^ 7 * (1 + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s ((cDelE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s (dTWCost p (2 * α ^ 2 - 2) Bq) + 10 ^ 9) +
  ((tabCost 1 t ((ptKE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((tabCost 1 s ((psKE.size + 1) * 10 ^ 7 * (Bq + 2) ^ 2) + 10 ^ 9) +
  ((6 * Bq + 6) +
  ((3 * 10 ^ 6 * (kb + 1) * (t * m + 1) * (Bi + kb + 6) +
      m * (bap0SlC p s t α mm Φ Bacc Q + 100000 * (t + s + 1) * (Q + 2)) +
      100000 * (m * s + 1) * (Q + 2)) +
  ((2 * 2000 * (Bq + 2) + s * (2 * HvdH.eWidth p α + 1) * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32) +
  ((2 * 2000 * (Bq + 2) + t * (12000 * (Bq + 2)) + 32) +
  (2 * 2000 * (Bq + 2) + s * (12000 * (Bq + 2)) + 32))))))))))))))))

end NLogN.AX

end

section
-- Source module: Solutions.KFMBaseDiag
open Complex

namespace IntMul.HvdH.DiagPf
open IntMul.HvdH

/-- A pair of integers as a Gaussian integer in `ℂ`. -/
def toC (p : ℤ × ℤ) : ℂ := (p.1 : ℂ) + (p.2 : ℂ) * I

end IntMul.HvdH.DiagPf

end

section
-- Source module: Solutions.KFMBaseP47
set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

section Resample

open Complex Real

namespace IntMul.HvdH.Res

/-! ### Definitions (HvdH §2.4, §4.1) -/

/-- Normalised DFT `(F_n u)_j = n⁻¹ ∑_k e^{-2πi jk/n} u_k` on `ℤ/n`. -/
noncomputable def dftN (n : ℕ) [NeZero n] (u : ZMod n → ℂ) : ZMod n → ℂ :=
  fun j => (n : ℂ)⁻¹ * ∑ k : ZMod n, ZMod.stdAddChar (-(j * k)) * u k

/-- Resampling map `S : ℂ^s → ℂ^t`,
`(S u)_k = α⁻¹ ∑_{j ∈ ℤ} exp(-π α⁻² s² (k/t - j/s)²) u_{j mod s}`. -/
noncomputable def resS (s t : ℕ) (α : ℝ) (u : ZMod s → ℂ) : ZMod t → ℂ :=
  fun k => (α : ℂ)⁻¹ * ∑' j : ℤ,
    ((Real.exp (-π * α⁻¹ ^ 2 * (s : ℝ) ^ 2 * ((k.val : ℝ) / t - j / s) ^ 2) : ℝ) : ℂ) * u j

/-- Resampling map `T : ℂ^s → ℂ^t`,
`(T u)_k = ∑_{j ∈ ℤ} exp(-π α² t² (k/t - j/s)²) u_{j mod s}`. -/
noncomputable def resT (s t : ℕ) (α : ℝ) (u : ZMod s → ℂ) : ZMod t → ℂ :=
  fun k => ∑' j : ℤ,
    ((Real.exp (-π * α ^ 2 * (t : ℝ) ^ 2 * ((k.val : ℝ) / t - j / s) ^ 2) : ℝ) : ℂ) * u j

/-- `(P_s u)_j = u_{t j}`. -/
def permS (s t : ℕ) (u : ZMod s → ℂ) : ZMod s → ℂ := fun j => u ((t : ZMod s) * j)

/-- `(P_t u)_k = u_{-s k}`. -/
def permT (s t : ℕ) (u : ZMod t → ℂ) : ZMod t → ℂ := fun k => u (-((s : ZMod t) * k))

/-! ### Gaussian sums -/

lemma summable_gauss {a : ℝ} (ha : 0 < a) (c : ℝ) (hc : c ≠ 0) (b : ℝ) :
    Summable (fun j : ℤ => Real.exp (-a * (c * j - b) ^ 2)) := by
  have hτ : 0 < (I * ((a * c ^ 2 / π : ℝ) : ℂ)).im := by
    have : 0 < c ^ 2 := by positivity
    rw [mul_im, I_re, I_im, ofReal_re, ofReal_im]
    simp only [zero_mul, one_mul, zero_add]
    positivity
  have h := ((summable_jacobiTheta₂_term_iff (I * ((-a * c * b / π : ℝ) : ℂ))
    (I * ((a * c ^ 2 / π : ℝ) : ℂ))).mpr hτ).norm
  have h2 := h.mul_left (Real.exp (-a * b ^ 2))
  refine h2.congr (fun n => ?_)
  rw [norm_jacobiTheta₂_term, ← Real.exp_add]
  congr 1
  simp only [mul_im, I_re, I_im, ofReal_re, ofReal_im]
  field_simp
  ring

lemma summable_gauss_mul {a : ℝ} (ha : 0 < a) (c : ℝ) (hc : c ≠ 0) (b : ℝ) (f : ℤ → ℂ)
    (B : ℝ) (hf : ∀ j, ‖f j‖ ≤ B) :
    Summable (fun j : ℤ => ((Real.exp (-a * (c * j - b) ^ 2) : ℝ) : ℂ) * f j) := by
  refine Summable.of_norm_bounded ((summable_gauss ha c hc b).mul_right B) (fun j => ?_)
  rw [norm_mul, Complex.norm_real, Real.norm_eq_abs, abs_of_pos (Real.exp_pos _)]
  exact mul_le_mul_of_nonneg_left (hf j) (Real.exp_pos _).le

/-! ### Characters -/

/-- The character `j ↦ e^{2πi ρ j / n}`. -/
noncomputable def ch (n : ℕ) [NeZero n] (ρ : ZMod n) : ZMod n → ℂ := fun j => ZMod.stdAddChar (ρ * j)

/-! ### Linearity -/

section lin
variable {s t : ℕ} [NeZero s] [NeZero t]

lemma bound_pi (u : ZMod s → ℂ) (j : ℤ) : ‖u j‖ ≤ ∑ k, ‖u k‖ :=
  Finset.single_le_sum (f := fun k => ‖u k‖) (fun _ _ => norm_nonneg _) (Finset.mem_univ _)

lemma summable_resS_term (α : ℝ) (hα : 0 < α) (u : ZMod s → ℂ) (x : ℝ) :
    Summable (fun j : ℤ =>
      ((Real.exp (-π * α⁻¹ ^ 2 * (s : ℝ) ^ 2 * (x - j / s) ^ 2) : ℝ) : ℂ) * u j) := by
  have hs : (s : ℝ) ≠ 0 := by exact_mod_cast NeZero.ne s
  have := summable_gauss_mul (a := π * α⁻¹ ^ 2) (by positivity) 1 one_ne_zero (s * x)
    (fun j => u j) _ (bound_pi u)
  refine this.congr (fun j => ?_)
  congr 3
  field_simp
  ring

lemma summable_resT_term (α : ℝ) (hα : 0 < α) (u : ZMod s → ℂ) (x : ℝ) :
    Summable (fun j : ℤ =>
      ((Real.exp (-π * α ^ 2 * (t : ℝ) ^ 2 * (x - j / s) ^ 2) : ℝ) : ℂ) * u j) := by
  have hs : (s : ℝ) ≠ 0 := by exact_mod_cast NeZero.ne s
  have ht : (t : ℝ) ≠ 0 := by exact_mod_cast NeZero.ne t
  have := summable_gauss_mul (a := π * α ^ 2) (by positivity) ((t : ℝ) / s)
    (div_ne_zero ht hs) (t * x) (fun j => u j) _ (bound_pi u)
  refine this.congr (fun j => ?_)
  congr 3
  field_simp
  ring

end lin

/-! ### The resampling identity (HvdH Theorem 4.2) -/

section ident
variable {s t : ℕ} [NeZero s] [NeZero t]

end ident

end IntMul.HvdH.Res

end Resample

section Bounds

open Real

namespace IntMul.HvdH.Res

end IntMul.HvdH.Res

end Bounds

section Maps

open Complex Real

namespace IntMul.HvdH.Res

/-- nearest integer `[x] = ⌊x + 1/2⌋` -/
noncomputable def rnd (x : ℝ) : ℤ := ⌊x + 1 / 2⌋

section maps
variable {s t : ℕ} [NeZero s] [NeZero t]

variable (s t) in
/-- `β_j = ⟨t j / s⟩` -/
noncomputable def beta (j : ℤ) : ℝ := (t : ℝ) * j / s - rnd ((t : ℝ) * j / s)

variable (s t) in
/-- Row selection `C : ℂ^t → ℂ^s`, `(C v)_l = v_{[t l / s]}`. -/
noncomputable def rowSel (v : ZMod t → ℂ) : ZMod s → ℂ :=
  fun l => v ((rnd ((t : ℝ) * l.val / s) : ℤ) : ZMod t)

variable (s t) in
/-- Diagonal map `(D_c u)_l = e^{c π α² β_l²} u_l` (`c = 1`: the map `D` of §4.2). -/
noncomputable def diagD (α c : ℝ) (u : ZMod s → ℂ) : ZMod s → ℂ :=
  fun l => ((Real.exp (c * π * α ^ 2 * beta s t l.val ^ 2) : ℝ) : ℂ) * u l

end maps

end IntMul.HvdH.Res

end Maps

section Prop47

open Complex Real

namespace IntMul.HvdH.Res

section clm
variable {s t : ℕ} [NeZero s] [NeZero t]

/-- Bundle a linear map between finite coordinate spaces. -/
noncomputable def clmOf {a b : ℕ} [NeZero a] (f : (ZMod a → ℂ) → (ZMod b → ℂ))
    (hadd : ∀ u v, f (u + v) = f u + f v) (hsmul : ∀ (c : ℂ) u, f (c • u) = c • f u) :
    (ZMod a → ℂ) →L[ℂ] (ZMod b → ℂ) :=
  LinearMap.toContinuousLinearMap ⟨⟨f, hadd⟩, hsmul⟩

@[simp] lemma clmOf_apply {a b : ℕ} [NeZero a] (f : (ZMod a → ℂ) → (ZMod b → ℂ)) (hadd hsmul)
    (u : ZMod a → ℂ) : IntMul.HvdH.Res.clmOf f hadd hsmul u = f u := rfl

lemma resS_add (α : ℝ) (hα : 0 < α) (u v : ZMod s → ℂ) :
    resS s t α (u + v) = resS s t α u + resS s t α v := by
  ext k; simp only [resS, Pi.add_apply, mul_add]
  rw [Summable.tsum_add (summable_resS_term α hα u _) (summable_resS_term α hα v _), mul_add]

lemma resS_smul (α : ℝ) (c : ℂ) (u : ZMod s → ℂ) :
    resS s t α (c • u) = c • resS s t α u := by
  ext k; simp only [resS, Pi.smul_apply, smul_eq_mul]
  rw [← tsum_mul_left, ← tsum_mul_left, ← tsum_mul_left]
  exact tsum_congr (fun j => by ring)

lemma resT_add (α : ℝ) (hα : 0 < α) (u v : ZMod s → ℂ) :
    resT s t α (u + v) = resT s t α u + resT s t α v := by
  ext k; simp only [resT, Pi.add_apply, mul_add]
  rw [Summable.tsum_add (summable_resT_term α hα u _) (summable_resT_term α hα v _)]

lemma resT_smul (α : ℝ) (c : ℂ) (u : ZMod s → ℂ) :
    resT s t α (c • u) = c • resT s t α u := by
  ext k; simp only [resT, Pi.smul_apply, smul_eq_mul]
  rw [← tsum_mul_left]; congr 1; ext j; ring

lemma diagD_add (α c : ℝ) (u v : ZMod s → ℂ) :
    diagD s t α c (u + v) = diagD s t α c u + diagD s t α c v := by
  ext l; simp [diagD, mul_add]

lemma diagD_smul (α c : ℝ) (a : ℂ) (u : ZMod s → ℂ) :
    diagD s t α c (a • u) = a • diagD s t α c u := by
  ext l; simp [diagD]; ring

variable (s t) in
noncomputable def Smap (α : ℝ) (hα : 0 < α) : (ZMod s → ℂ) →L[ℂ] (ZMod t → ℂ) :=
  IntMul.HvdH.Res.clmOf (resS s t α) (resS_add α hα) (resS_smul α)

variable (s t) in
noncomputable def Tmap (α : ℝ) (hα : 0 < α) : (ZMod s → ℂ) →L[ℂ] (ZMod t → ℂ) :=
  IntMul.HvdH.Res.clmOf (resT s t α) (resT_add α hα) (resT_smul α)

variable (s t) in
noncomputable def Cmap : (ZMod t → ℂ) →L[ℂ] (ZMod s → ℂ) :=
  IntMul.HvdH.Res.clmOf (rowSel s t) (fun _ _ => rfl) (fun _ _ => rfl)

variable (s t) in
noncomputable def Dmap (α : ℝ) : (ZMod s → ℂ) →L[ℂ] (ZMod s → ℂ) :=
  IntMul.HvdH.Res.clmOf (diagD s t α 1) (diagD_add α 1) (diagD_smul α 1)

variable (s t) in
noncomputable def PTmap : (ZMod t → ℂ) →L[ℂ] (ZMod t → ℂ) :=
  IntMul.HvdH.Res.clmOf (IntMul.HvdH.Res.permT s t) (fun _ _ => rfl) (fun _ _ => rfl)

variable (s t) in
/-- `(P_s⁻¹ w)_j = w_{t⁻¹ j}`. -/
noncomputable def PSinv : (ZMod s → ℂ) →L[ℂ] (ZMod s → ℂ) :=
  IntMul.HvdH.Res.clmOf (fun w j => w ((t : ZMod s)⁻¹ * j)) (fun _ _ => rfl) (fun _ _ => rfl)

end clm

end IntMul.HvdH.Res

end Prop47

open Complex Real

namespace IntMul.HvdH.P47
open IntMul.HvdH IntMul.HvdH.Res

/-- A pair of integers as a Gaussian integer in `ℂ`. -/
def toC (p : ℤ × ℤ) : ℂ := (p.1 : ℂ) + (p.2 : ℂ) * I

/-- `x = α²θ`. -/
noncomputable def xθ (s t α : ℕ) : ℝ := ((α ^ 2 * (t - s) : ℕ) : ℝ) / s

section
variable {s t : ℕ} [NeZero s] [NeZero t]

end

end IntMul.HvdH.P47

end

section
-- Source module: Solutions.KFMBaseTGo
/-!
# Base `ref`: `tensorGo` computes `tensorMapZ` (copies)

Verbatim excerpts of the base thread's `ref/` development (not in this workspace):
`Spec.ixf`, all of `ListBox`, the `lists` section and the `boxIdxs`/`ixf` cons lemmas of `MDif`,
the snoc lemmas of `FermatW`, and all of `TGo`. Main result: **`tensorGo_eq`**, the list
pipeline `tensorGo` (nlogn's machine reference) on row-major `listOf` arrays equals the base
fixed-point tensor `tensorMapZ` read on the output box.
-/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

namespace IntMul.HvdH.RefPf
open IntMul.HvdH IntMul.HvdH.Ref

/-- the index function of a list index -/
def ixf (l : List ℕ) : ℕ → ℕ := fun i => l.getD i 0

section lists
variable {α : Type*}

end lists

/-- the list of steps fed to `tensorGo` -/
def tgSteps (K : ℕ) (ns' : List ℕ) (A : ℕ → (ℕ → ℤ × ℤ) → (ℕ → ℤ × ℤ)) :
    List (ℕ × (List (ℤ × ℤ) → List (ℤ × ℤ))) :=
  (List.range K).map fun i => (ns'.getD i 1, fun X : List (ℤ × ℤ) =>
    (List.range (ns'.getD i 1)).map (A i fun j => X.getD j (0, 0)))

end IntMul.HvdH.RefPf

end

section
-- Source module: Solutions.ResampleHvdH
/-! Copy (renamed namespace `IntMul.OpenAI.RH`) of the base thread's HvdH Prop 4.7(i)
(`intmul/base/Sol_prop_4_7_i.lean`): Gaussian resampling identity, maps S, T, C, D, P_s, P_t. -/

set_option linter.unusedSectionVars false
set_option linter.unusedSimpArgs false

section Resample

open Complex Real

namespace IntMul.OpenAI.RH

/-! ### Definitions (HvdH §2.4, §4.1) -/

/-- `(P_s u)_j = u_{t j}`. -/
def permS (s t : ℕ) (u : ZMod s → ℂ) : ZMod s → ℂ := fun j => u ((t : ZMod s) * j)

/-- `(P_t u)_k = u_{-s k}`. -/
def permT (s t : ℕ) (u : ZMod t → ℂ) : ZMod t → ℂ := fun k => u (-((s : ZMod t) * k))

/-! ### Gaussian sums -/

/-! ### Characters -/

/-! ### Linearity -/

section lin
variable {s t : ℕ} [NeZero s] [NeZero t]

end lin

/-! ### The resampling identity (HvdH Theorem 4.2) -/

section ident
variable {s t : ℕ} [NeZero s] [NeZero t]

end ident

end IntMul.OpenAI.RH

end Resample

section Bounds

open Real

namespace IntMul.OpenAI.RH

end IntMul.OpenAI.RH

end Bounds

section Maps

open Complex Real

namespace IntMul.OpenAI.RH

section maps
variable {s t : ℕ} [NeZero s] [NeZero t]

end maps

end IntMul.OpenAI.RH

end Maps

section Prop47

open Complex Real

namespace IntMul.OpenAI.RH

section clm
variable {s t : ℕ} [NeZero s] [NeZero t]

/-- Bundle a linear map between finite coordinate spaces. -/
noncomputable def clmOf {a b : ℕ} [NeZero a] (f : (ZMod a → ℂ) → (ZMod b → ℂ))
    (hadd : ∀ u v, f (u + v) = f u + f v) (hsmul : ∀ (c : ℂ) u, f (c • u) = c • f u) :
    (ZMod a → ℂ) →L[ℂ] (ZMod b → ℂ) :=
  LinearMap.toContinuousLinearMap ⟨⟨f, hadd⟩, hsmul⟩

@[simp] lemma clmOf_apply {a b : ℕ} [NeZero a] (f : (ZMod a → ℂ) → (ZMod b → ℂ)) (hadd hsmul)
    (u : ZMod a → ℂ) : IntMul.OpenAI.RH.clmOf f hadd hsmul u = f u := rfl

end clm

end IntMul.OpenAI.RH

end Prop47

end

section
-- Source module: Solutions.ResampleTensor
/-!
# §8.2 Lemma 20: tensor interface, padded box, linewise implementation

* `lemma20_tensor`: tensoring the one-dimensional identities (33) gives (35),
  `R F_s = 2^γ B₀ Q F_t A` with `γ = 2dα²`, `R = ⊗ P_{s_i}`, `Q = ⊗ P_{t_i}`, `‖A‖, ‖B₀‖ ≤ 1`.
* `axisOp` / `tensOpI`: the padded box keeps its full shape; applying the padded line kernels of
  the axes one at a time (in any order) equals the tensor kernel (`axis_fold_eq_tensOp`), and
  with zero padding outside `∏ [0, s_i)` this is exactly the tensor map `A`
  (`tensOp_pad_eq_tens`).
* `approx_iterate`: `d` contraction stages with errors `< E` each give error `< d E`.
* cost bookkeeping: `∑_i (T/t_i)·(t_i X) = d T X`, line count `∑ T/t_i ≤ 2dT/r`.
-/

open Finset Complex

namespace IntMul.OpenAI.RT

open IntMul.OpenAI.RH

variable {d : ℕ}

/-! ### Kernel tensor (copied from the base thread's `Thm41.lean`, renamed) -/

/-! ### Multidimensional transforms and the retained permutations -/

/-- `Q = ⊗ P_{t_i}`: `(Q w)_k = w_{(-s_i k_i)_i}`. -/
def Qperm (s t : Fin d → ℕ) (w : ((i : Fin d) → ZMod (t i)) → ℂ) :
    ((i : Fin d) → ZMod (t i)) → ℂ :=
  fun k => w (fun i => -((s i : ZMod (t i)) * k i))

/-! ### The padded box: axis maps applied one at a time -/

section padded
variable {n : Fin d → ℕ}

end padded

/-! ### Errors and costs of the linewise implementation -/

end IntMul.OpenAI.RT

end

section
-- Source module: Solutions.ResampleChirp
/-!
# §8.3 Lemma 21: chirp identity (36) and permutation cancellation (37)

* `chirp_periodic`: `a_j = exp(-π i Σ s_i j_i²/t_i)` has period `t_i` in coordinate `i`
  (each `t_i` even).
* `lemma21_chirp` (36): `Q F_t u = ā · (a ∗ (ā · u)) / T`.
* `lemma21_cancel` (37): `R F_s⁺ ((R F_s u)·(R F_s v)) = (u ∗ v)/S²` for `t_i` units mod `s_i`.
* `conj_dftM`: conjugating input and output of `F_s` gives `F_s⁺` (also after `R`).
* `chirp_phase_numerator` (§9.5): the chirp phase is `exp(2πi N/(2r))` with
  `N = -s j² (r/t)`.
-/

open Finset Complex Real

namespace IntMul.OpenAI.RC

open IntMul.OpenAI.RT

variable {d : ℕ}

/-! ### One-dimensional chirp factor -/

/-- `g(x) = exp(-π i s x² / t)` for an integer `x`. -/
noncomputable def gZ (s t : ℕ) (x : ℤ) : ℂ := Complex.exp (-(π * I * s * x ^ 2 / t))

/-- Chirp read on `ZMod t` through the representative `val`. -/
noncomputable def gch (s t : ℕ) (j : ZMod t) : ℂ := gZ s t (j.val : ℤ)

/-! ### Multidimensional chirp and identity (36) -/

/-- `a_j = ∏_i exp(-π i s_i j_i² / t_i)`. -/
noncomputable def chirp (s t : Fin d → ℕ) (j : (i : Fin d) → ZMod (t i)) : ℂ :=
  ∏ i, gch (s i) (t i) (j i)

/-! ### Permutation cancellation (37) -/

section cancel
variable (s : Fin d → ℕ) [∀ i, NeZero (s i)]

variable (t : Fin d → ℕ)

end cancel

/-! ### §9.5: chirp phase numerators modulo `2r` -/

end IntMul.OpenAI.RC

end

section
-- Source module: Solutions.KFMR19
/-!
# Item 12 (math side): Lemma 19/20 with the base witnesses and a permutation-free `B̃₀`

The resampling identities of `Resample19`/`ResampleTensor`/`ResampleFinal` are existential in
the maps `A`, `B₀`. The machine computes the base thread's fixed-point maps `Sapprox` and
`Bapprox`. This file pins the witnesses to those maps.

* `lemma19_base`: for one axis, the base maps `A`, `B` of Proposition 4.7 give
  `P_s F_s = 2^{2α²} B₀ P_t F_t A` with `B₀ := P_s ∘ B ∘ P_t⁻¹` (`reidx` maps), and the
  **permutation-free** procedure `B0approx = D̃' ∘ J̃' ∘ C̃` approximates `2^p B₀`. The error
  bound is the base `Bapprox` bound, because `Bapprox = P̃_s⁻¹ ∘ B0approx ∘ P̃_t`.
* `tensor_of_lines_expl`: `tensor_of_lines` with the explicit witnesses `tensCLM (kerA A)` and
  `tensCLM (kerB B)`.
* `lemma19_fp`: the same in the `FPApprox`/`SupLip` form of base `prop_4_7_fp`, with the
  periodic read-out `B0P` (index mod `s`) and the `ℕ`-lifts `liftL`.
* `lemma20_base`, `source_transform_base`: Lemma 20 and the chirp form of the source transform,
  with tensor witnesses built from per-axis maps whose `Sapprox` / `B0approx` bounds hold.

The §9 glue does not consume the existential forms (only `ResampleAll` prints their axioms),
so no bridge lemma is needed.
-/

open Complex

namespace NLogN.KFM

open IntMul.OpenAI.RT

/-! ### Reindexing contractions -/

/-- `(reidx σ u)_j = u_{σ j}`. -/
noncomputable def reidx {α β : Type} [Fintype α] [Fintype β] (σ : β → α) :
    (α → ℂ) →L[ℂ] (β → ℂ) :=
  ContinuousLinearMap.pi fun j => ContinuousLinearMap.proj (R := ℂ) (φ := fun _ : α => ℂ) (σ j)

@[simp] lemma reidx_apply {α β : Type} [Fintype α] [Fintype β] (σ : β → α) (u : α → ℂ) (j : β) :
    NLogN.KFM.reidx σ u j = u (σ j) := rfl

/-! ### The permutation-free procedure and the single-axis statement -/

/-- `P_t⁻¹` as an index map: `k ↦ -(s⁻¹ k)`. -/
noncomputable def tauT (s t : ℕ) (k : ZMod t) : ZMod t := -((s : ZMod t)⁻¹ * k)

/-- `B₀ = P_s ∘ B ∘ P_t⁻¹`. -/
noncomputable def B0of (s t : ℕ) [NeZero s] [NeZero t] (B : (ZMod t → ℂ) →L[ℂ] (ZMod s → ℂ)) :
    (ZMod t → ℂ) →L[ℂ] (ZMod s → ℂ) :=
  (NLogN.KFM.reidx (fun l : ZMod s => (t : ZMod s) * l)).comp (B.comp (NLogN.KFM.reidx (tauT s t)))

/-- `W ∘ P_t⁻¹` on `ℕ` indices: `(unPt W)_i = W_{-(s⁻¹ i) mod t}`. -/
noncomputable def unPt (s t : ℕ) (W : ℕ → ℤ × ℤ) : ℕ → ℤ × ℤ :=
  fun i => W (tauT s t (i : ZMod t)).val

/-! ### The `FPApprox` form (base Proposition 4.7, `prop_4_7_fp`) -/

/-- A line map on `ℤ/a` lifted to `ℕ`-indexed arrays (reads `[0,a)`, periodic output). -/
noncomputable def liftL {a b : ℕ} (L : (ZMod a → ℂ) →L[ℂ] (ZMod b → ℂ)) :
    (ℕ → ℂ) → (ℕ → ℂ) :=
  fun x j => L (fun k => x k.val) (j : ZMod b)

/-! ### Explicit tensor witnesses -/

variable {d : ℕ}

/-- Kernel of a family of line maps `A i : ℂ^{s_i} → ℂ^{t_i}`. -/
noncomputable def kerA {s t : Fin d → ℕ} [∀ i, NeZero (s i)] [∀ i, NeZero (t i)]
    (A : (i : Fin d) → (ZMod (s i) → ℂ) →L[ℂ] (ZMod (t i) → ℂ)) :
    (i : Fin d) → ZMod (t i) → ZMod (s i) → ℂ :=
  fun i l j => A i (Pi.single j 1) l

/-- Kernel of a family of line maps `B i : ℂ^{t_i} → ℂ^{s_i}`. -/
noncomputable def kerB {s t : Fin d → ℕ} [∀ i, NeZero (s i)] [∀ i, NeZero (t i)]
    (B : (i : Fin d) → (ZMod (t i) → ℂ) →L[ℂ] (ZMod (s i) → ℂ)) :
    (i : Fin d) → ZMod (s i) → ZMod (t i) → ℂ :=
  fun i k m => B i (Pi.single m 1) k

/-- Per-axis base data: the maps of `lemma19_fp` with their identities, `FPApprox` bounds and
Lipschitz facts. -/
structure AxisMaps (p α : ℕ) (s t : Fin d → ℕ) [∀ i, NeZero (s i)] [∀ i, NeZero (t i)] where
  A : (i : Fin d) → (ZMod (s i) → ℂ) →L[ℂ] (ZMod (t i) → ℂ)
  B : (i : Fin d) → (ZMod (t i) → ℂ) →L[ℂ] (ZMod (s i) → ℂ)
  hA : ∀ i, ‖A i‖ ≤ 1
  hB : ∀ i, ‖B i‖ ≤ 1
  hid : ∀ i u, IntMul.OpenAI.RH.permS (s i) (t i) ((s i : ℂ)⁻¹ • ZMod.dft u) =
    (2 : ℂ) ^ (2 * α ^ 2) • B i (IntMul.OpenAI.RH.permT (s i) (t i) ((t i : ℂ)⁻¹ • ZMod.dft (A i u)))
  hS : ∀ i, IntMul.HvdH.FPApprox p (IntMul.HvdH.Sapprox p (s i) (t i) α) (NLogN.KFM.liftL (A i))
    (6 * IntMul.HvdH.sWidth p α + 4)
  hB0 : ∀ i, IntMul.HvdH.FPApprox p (B0P p (s i) (t i) α) (NLogN.KFM.liftL (B i))
    (IntMul.HvdH.jIter p (s i) (t i) α * (10 * IntMul.HvdH.eWidth p α + 12) + 4)
  hLA : ∀ i, IntMul.HvdH.SupLip (NLogN.KFM.liftL (A i))
  hLB : ∀ i, IntMul.HvdH.SupLip (NLogN.KFM.liftL (B i))

end NLogN.KFM

end

section
-- Source module: Solutions.KFMLines
/-!
# Item 12 (math side, part 2): line maps applied axis by axis

The machine applies the per-axis fixed-point maps (`Sapprox`, `B0P`) slice by slice. That is
the base `tensorMapZ` (via `tensorGo_eq`, `KFMBaseTGo`). This file identifies the exact
counterpart, the base `tensorMap` of the lifted line maps `liftL (L i)`, with the explicit tensor
witness `tensCLM (kerL L)` of `KFMR19`. Then base `fp_tensor` gives the error of the whole pass.

* `liftL_sum`: `liftL L x m = ∑_{n<a} K(m,n) x_n` with `K(m,n) = kerL L (m mod b) (n mod a)`.
* `sliceComp_sum`: after the axes of a duplicate-free list `l`, the array is the partial kernel
  sum over the box `∏_{i∈l} [0,a_i)`.
* **`tensorMap_liftL`**: `tensorMap (liftL ∘ L) x j = tens (kerL L) x̃ ĵ`.
* **`tensor_fp`**: `tensorMapZ Aa` keeps the disc and is within `∑ ε_i` of
  `2^p · tens (kerL L) (2^{-p} U)` at every box address.
-/

open Finset Complex

namespace NLogN.KFM

open IntMul.OpenAI.RT IntMul.HvdH

variable {d : ℕ}

/-- Kernel of a family of line maps `L i : ℂ^{a_i} → ℂ^{b_i}` (output index first). -/
noncomputable def kerL {a b : Fin d → ℕ} [∀ i, NeZero (a i)] [∀ i, NeZero (b i)]
    (L : (i : Fin d) → (ZMod (a i) → ℂ) →L[ℂ] (ZMod (b i) → ℂ)) :
    (i : Fin d) → ZMod (b i) → ZMod (a i) → ℂ :=
  fun i l j => L i (Pi.single j 1) l

section slices
variable (a : Fin d → ℕ) (K : Fin d → ℕ → ℕ → ℂ) (A : Fin d → (ℕ → ℂ) → (ℕ → ℂ))

/-- The index box after the axes of `l`: `[0, a_i)` on `l`, the point `j_i` elsewhere. -/
def boxF (l : List (Fin d)) (j : Fin d → ℕ) : Finset (Fin d → ℕ) :=
  Fintype.piFinset fun i => if i ∈ l then range (a i) else {j i}

end slices

/-! ### The list pipeline -/

end NLogN.KFM

end

section
-- Source module: NLogN.SSMath
/-!
# Arithmetic around the Fermat-ring FFT: pieces, twisting, recomposition and carries
-/

namespace NLogN

open IntMul

/-- `∑ A_i B^i`. -/
def zev (B : ℤ) : List ℤ → ℤ
  | [] => 0
  | a :: A => a + B * zev B A

/-- The `K` signed `M`-bit pieces of `x` (least significant first). -/
def pcs (M : ℕ) : ℕ → ℤ → List ℤ
  | 0, _ => []
  | K + 1, x => sg x (x.natAbs % 2 ^ M) :: pcs M K (sg x (x.natAbs / 2 ^ M))

/-- Twisting by powers of `2^Q`. -/
def wtZ (F : ℤ) (Q : ℕ) : ℕ → List ℤ → List ℤ
  | _, [] => []
  | i, a :: A => (a * 2 ^ (i * Q)) % F :: wtZ F Q (i + 1) A

/-- Negacyclic convolution of length `K`. -/
def ncZ (K : ℕ) (A B : List ℤ) : List ℤ :=
  (List.range K).map fun m => ∑ i ∈ Finset.range K, ∑ l ∈ Finset.range K,
    ((if i + l = m then A.getD i 0 * B.getD l 0 else 0) -
      (if i + l = m + K then A.getD i 0 * B.getD l 0 else 0))

section ring

variable (n : ℕ)

end ring

/-- Carry propagation into `M`-bit digits. -/
def carZ (M : ℕ) : List ℤ → ℤ → List ℤ × ℤ
  | [], c => ([], c)
  | e :: E, c => ((e + c) % 2 ^ M :: (carZ M E ((e + c) / 2 ^ M)).1,
      (carZ M E ((e + c) / 2 ^ M)).2)

end NLogN

end

section
-- Source module: NLogN.SSUtil
/-!
# Helper procedures: bit length and powers of two
-/

namespace NLogN

open IntMul

abbrev RMs : Fin NS := 18
abbrev RQ : Fin NS := 19
abbrev QS : Fin NS := 20
abbrev XA : Fin NS := 21
abbrev W2 : Fin NS := 22
abbrev HB : Fin NS := 23
abbrev CYS : Fin NS := 24
abbrev DG : Fin NS := 25

abbrev PW : Fin NS := 27
abbrev PB : Fin NS := 28

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt NS P}

/-! ## Bit length -/

end NLogN

end