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

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

import Definitions.Def_IntMul_Kappa_K_Defs20
import Definitions.Def_IntMul_HvdH_BApprox
import Definitions.Def_IntMul_HvdH_DiagApprox
import Definitions.Def_IntMul_HvdH_EApprox
import Definitions.Def_IntMul_HvdH_FPApprox
import Definitions.Def_IntMul_HvdH_JApprox
import Definitions.Def_IntMul_HvdH_SApprox

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.ResampleParams
/-!
# §9.1–9.2 Fixed rational choice and sizes

* `constants_ok`: the layer and dimension conditions of (38), `margins_ok`: the seven margins
  of (45) with `min gᵢ = 2⁻¹⁸¹ = 2κ`.
* Axis lengths `t_i ∈ {r/2, r}` with `∏ t_i = T`, last axis `r` (`axExp_sum`, `axLen_prod`,
  `axExp_last`), `p/(12d) ≤ ℓ ≤ p/(3d)` (`ell_bounds`), `24n ≤ T p < 48 n` (`Tp_bounds`).
* Source lengths: `(1-2η)t_i < s_i ≤ (1-η)t_i` ⇒ `S > T/2` (`S_gt_half`), `θ_i > η`
  (`theta_gt_eta`), `α⁴θ_i > p` (`alpha4_theta_gt_p`), `γ < 28 d² √b` (`gamma_bound`);
  existence of pairwise coprime odd `s_i` (`sources_exist`, from `RCop.coprime_params`).
* Radix digits (§9.2.3): `2q - 2 < S` (`degree_lt_S`) and `c_j < 2^{3b}` (`coeff_lt`).
-/

open Finset

namespace IntMul.OpenAI.RP

/-! ### (38) and (45) -/

/-! ### Axis lengths of the padded box -/

/-- Number of axes of length `r/2`: `dℓ - L`. -/
def mcnt (L d : ℕ) : ℕ := d * ell L d - L

/-! ### Source lengths, resampling margins -/

/-! ### §9.2.3 radix-digit polynomials -/

end IntMul.OpenAI.RP

end

section
-- Source module: Solutions.KFHDefs
/-! # KFH: Lemma 13 address definitions (moved out of KFH_WIP so stable files avoid the WIP import) -/
set_option linter.unusedVariables false

namespace NLogN.KFH

open IntMul RS AModel KIL KFD KFB Complex

/-! ## Lemma 13 -/

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

/-- The rectangle `[N0] × [2^L] × [N1] × [2^L] × [N2] × [2^L] × [N3]`. -/
def rect (N0 N1 N2 N3 L : ℕ) : Shape := [N0, 2 ^ L, N1, 2 ^ L, N2, 2 ^ L, N3]

/-- Field index of chunk slot `t : Fin 3`. -/
def slotF (t : Fin 3) : ℕ := 2 * t.val + 1

/-- The address map of one packed row addition: `y ← y ⊕ (x ∧ mask)`. -/
def rowAddAddr (K ρ f : ℕ) (px py : Fin 3) (a : List ℕ) : List ℕ :=
  a.set (slotF py) ((a.getD (slotF py) 0) ^^^ ((a.getD (slotF px) 0) &&& selMask K ρ f))

/-- Records permuted by an address map (records of `R` Gaussian entries). -/
def permRecs (sh : Shape) (π : List ℕ → List ℕ) (R : ℕ) (zs : List GI) : List GI :=
  (List.range (vol sh)).flatMap fun q => (zs.drop (encode sh (π (decode sh q)) * R)).take R

end NLogN.KFH

end

section
-- Source module: Solutions.KFLAddr
/-!
# KFL §17: `permRecs (rowAddAddr …)` as a triple map on the bit array

The Gaussian records (`R` entries, `2 (w+1)` bits per entry after `toBitsI`) of the rectangle
`rect N0 N1 N2 N3 L` form the bit array of the geometry `geoB` (`Q = 2^L`, `B = 2 R (w+1)`).
`bits_permRecs`: converting `permRecs rect (rowAddAddr px py) R zs` to bits gives
`rixT (geoB …) (Fpp px py)` of the bits of `zs`, where `Fpp px py` replaces slot `py` by
`y ⊕ (x ∧ selM f)` (`x` = slot `px`).  `Fpp_conj`: `Fpp px py` is the canonical row addition
`RAc` conjugated by at most two slot swaps.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD KFB KFH AModel Butterfly

/-! ## Chunks of a list -/

/-! ## The bit records -/

/-- Record `m` of `zs` in bits. -/
def recB (R w : ℕ) (zs : List GI) (m : ℕ) : List ℤ := bitsL w (gflat ((zs.drop (m * R)).take R))

/-- Padded to all `m`. -/
def recsB (R w M : ℕ) (zs : List GI) (m : ℕ) : List ℤ :=
  if m < M then recB R w zs m else List.replicate (R * (2 * (w + 1))) 0

/-! ## The geometry of the bit array -/

/-- The bit-array geometry: `Q = 2^L`, records of `B = 2 R (w + 1)` bits. -/
def geoB (N0 N1 N2 N3 L R w : ℕ) : Geo := ⟨2 ^ L, N0, N1, N2, N3, R * (2 * (w + 1))⟩

section Geom

variable {N0 N1 N2 N3 L R w : ℕ}

/-- Slot access on triples. -/
def slotG (c : Tr) : Fin 3 → ℕ
  | 0 => c.1
  | 1 => c.2.1
  | 2 => c.2.2

def slotS (c : Tr) : Fin 3 → ℕ → Tr
  | 0, v => (v, c.2.1, c.2.2)
  | 1, v => (c.1, v, c.2.2)
  | 2, v => (c.1, c.2.1, v)

/-- The stride of a slot. -/
def strd (g : Geo) : Fin 3 → ℕ
  | 0 => g.S0
  | 1 => g.S1
  | 2 => g.S2

end Geom

/-! ## The slot row addition as a triple map -/

/-- Slot `py` gets `y ⊕ (x ∧ selM f)`, `x` = slot `px`. -/
def Fpp (K ρ f : ℕ) (px py : Fin 3) (c : Tr) : Tr :=
  slotS c py (slotG c py ^^^ (slotG c px &&& selM K ρ f))

/-! ## Records permuted by `rowAddAddr`, in bits -/

/-! ## Conjugation by slot swaps -/

def swT : Fin 3 → Tr → Tr
  | 0 => sw01
  | 1 => sw12
  | 2 => sw02

def compL : List (Fin 3) → Tr → Tr
  | [] => id
  | i :: l => swT i ∘ compL l

/-- The swaps bringing `(px, py)` to the canonical slots `(0, 1)`. -/
def preL : Fin 3 → Fin 3 → List (Fin 3)
  | 1, 0 => [0]
  | 0, 2 => [1]
  | 2, 1 => [2]
  | 2, 0 => [2, 1]
  | 1, 2 => [0, 1]
  | _, _ => []

end NLogN.KFL

end

section
-- Source module: Solutions.KFLMain
/-!
# KFL §18: Lemma 13 in IL

`lem13_il` (the corrected statement of `KFH_WIP.lem13_iruns`, see `KFL_WIP`): integers → bits
(`toBitsI`), the slot canonicalisation by at most two Lemma-12 swaps (`swSeq (preL px py)`), the
packed lines (`linesI`), the guard-set correction (`corrI`), the last selected bit (`xorStep`),
the inverse swaps and bits → integers (`fromBitsI`).  The composite is `permRecs (rowAddAddr …)`
(`bits_permRecs`, `pipeline_rixT`, `Fpp_conj`); the cost is of the form (19) with
`τ = 1 - 2^{-50}`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFL

open IntMul RS KIL KFD KFB KFH AModel Butterfly

/-! ## Sequences of slot swaps -/

def swCode (q u : ℕ) (g : Geo) : Fin 3 → IStmt HS 1
  | 0 => swI q g.N0 g.N1 (vol [g.N2, g.Q, g.N3, g.B]) u
  | 1 => swI q (vol [g.N0, g.Q, g.N1]) g.N2 (vol [g.N3, g.B]) u
  | 2 => swI q g.N0 (vol [g.N1, g.Q, g.N2]) (vol [g.N3, g.B]) u

def swSeq (q u : ℕ) (g : Geo) : List (Fin 3) → IStmt HS 1
  | [] => .skip
  | i :: l => .seq (swCode q u g i) (swSeq q u g l)

/-- The swap property delivered by `swI_geo` for one prime `q`. -/
def SwOK (q : ℕ) (Ksw : ℝ) : Prop :=
  ∀ (i : Fin 3) (g : Geo), g.Pos → ∀ u : ℕ, 1 ≤ u → g.Q = 2 ^ u →
    ∀ (I : List Sym) (t : IState HS) (zs : List ℤ) (R : List Item),
    Bits zs → t.G 0 = ints zs ++ R → zs.length = g.V →
    ∃ c : ℕ, IRuns I l12prog (swCode q u g i) t (t.upd 0 (ints (rixT g (swT i) zs) ++ R)) c ∧
      (c : ℝ) ≤ Ksw * ((g.V : ℝ) + 1) * (u : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50)

/-! ## Cost bounds of the rotations -/

/-! ## The program -/

/-- The last-bit step on the geometry. -/
def xorG (K ρ f : ℕ) (g : Geo) : IStmt HS 1 :=
  xorStep (g.N1 * 2 ^ (f * K - (ρ + (f - 1) * K) - 1)) g.Q (ρ + (f - 1) * K)
    (g.N0 * g.Q * g.N1 * 2 ^ (f * K - (ρ + (f - 1) * K) - 1)) (2 ^ (ρ + (f - 1) * K) * g.S1)

/-- Lines, correction, last bit: the canonical row addition `RAc` on the bit array. -/
def coreB (q K ρ f : ℕ) (g : Geo) : IStmt HS 1 :=
  .seq (linesI q K ρ f g lines) (.seq (corrI K ρ f g) (xorG K ρ f g))

/-- **The code of Lemma 13** (for the prime `q` of Lemma 12). -/
def l13code (q N0 N1 N2 N3 K f ρ R w : ℕ) (px py : Fin 3) : IStmt HS 1 :=
  .seq (toBitsI (2 * (vol (rect N0 N1 N2 N3 (f * K)) * R)) w)
    (.seq (swSeq q (f * K) (geoB N0 N1 N2 N3 (f * K) R w) (preL px py))
    (.seq (coreB q K ρ f (geoB N0 N1 N2 N3 (f * K) R w))
    (.seq (swSeq q (f * K) (geoB N0 N1 N2 N3 (f * K) R w) (preL px py).reverse)
      (fromBitsI (2 * (vol (rect N0 N1 N2 N3 (f * K)) * R)) w))))

/-- The lines property delivered by `linesI_spec` for one prime `q`. -/
def LinesOK (q : ℕ) (Ks : ℝ) : Prop :=
  ∀ (K ρ f : ℕ), 1 ≤ K → ρ < K → 1 ≤ f →
    ∀ (g : Geo), g.Pos → g.Q = 2 ^ (f * K) →
    ∀ (Bb : ℕ), 2 ≤ Bb → 4 * 2 ^ (f * K) < 2 ^ Bb → zsz (g.V : ℤ) ≤ Bb →
    ∀ (ls : List Line), (∀ l ∈ ls, LineOK l) →
    ∀ (I : List Sym) (t : IState HS) (zs : List ℤ) (R : List Item),
    Bits zs → t.G 0 = ints zs ++ R → zs.length = g.V →
    ∃ c : ℕ, IRuns I l12prog (linesI q K ρ f g ls) t
        (t.upd 0 (ints (rixT g (fun c => ls.foldr (fun l c => lineP K ρ f l c) c) zs) ++ R)) c ∧
      (c : ℝ) ≤ 1 + ls.length * (2 * Ks * ((g.V : ℝ) + 1) * ((f * K : ℕ) : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50)
        + cRotG K f g Bb)

/-- The correction property delivered by `corrI_spec`. -/
def CorrOK (Kc : ℝ) : Prop :=
  ∀ (K ρ f : ℕ), 6 ≤ K → ρ < K → 1 ≤ f →
    ∀ (g : Geo), g.Pos → g.Q = 2 ^ (f * K) →
    ∀ (I : List Sym) (iprog : Fin 1 → IStmt HS 1) (t : IState HS) (zs : List ℤ) (R : List Item),
    Bits zs → t.G 0 = ints zs ++ R → zs.length = g.V →
    ∃ c : ℕ, IRuns I iprog (corrI K ρ f g) t
        (t.upd 0 (ints (rixT g (corrc K ρ f) zs) ++ R)) c ∧
      (c : ℝ) ≤ Kc * (((g.V / g.B : ℕ) : ℝ) * ((zsz (g.V : ℤ) : ℝ) + 1) ^ 3 + (g.V : ℝ) +
        min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * ((g.V / g.B : ℕ) : ℝ) *
          ((zsz (g.V : ℤ) : ℝ) + 1) * ((g.B : ℝ) + (zsz (g.V : ℤ) : ℝ) + 1))

/-- The cost bound of the core. -/
noncomputable def coreCost (Ks Kcr : ℝ) (K ρ f : ℕ) (g : Geo) : ℝ :=
  (1 + 8 * (2 * Ks * ((g.V : ℝ) + 1) * ((f * K : ℕ) : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50)
        + cRotG K f g (zsz (g.V : ℤ) + 3))) +
    Kcr * (((g.V / g.B : ℕ) : ℝ) * ((zsz (g.V : ℤ) : ℝ) + 1) ^ 3 + (g.V : ℝ) +
      min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * ((g.V / g.B : ℕ) : ℝ) *
        ((zsz (g.V : ℤ) : ℝ) + 1) * ((g.B : ℝ) + (zsz (g.V : ℤ) : ℝ) + 1)) +
    cXor (g.N0 * g.Q * g.N1 * 2 ^ (f * K - (ρ + (f - 1) * K) - 1)) (2 ^ (ρ + (f - 1) * K) * g.S1)
      1 (zsz (g.V : ℤ) + 3)

/-! ## Lemma 13: denotation and the cost pieces -/

/-! ## The cost in the form (19) -/

end NLogN.KFL

end

section
-- Source module: Solutions.KFMBaseTw
/-! Copy of base `Sol_twiddle_bound.lean`; the two `Theorems.*` imports are replaced by the
proved copies `KFMBasePi`/`KFMBaseExp` through the shims below. -/

open Complex ComplexConjugate

namespace IntMul.HvdH.TwPf
open IntMul.HvdH

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

end IntMul.HvdH.TwPf

end

section
-- Source module: Solutions.KFRBits
/-!
# KFR: bit arithmetic for the packed row addition

* `bitZ z l ∈ {0, 1}` (as an integer) and `bitsum`: `z = Σ_{l<N} bitZ z l · 2^l` for `z < 2^N`.
* `testBit_selMask`: the mask `Σ_{u<f} 2^{ρ + uK}` has exactly the bits `ρ + uK`, `u < f`.
* **`xor_mask_sub`**: `(Y ⊕ (X ∧ mask)) - Y = Σ_{u<f} (bit_{ρ+uK}(Y ⊕ X) - bit_{ρ+uK} Y) 2^{ρ+uK}`.
* `bitZ_digit`: bit `l` of the field `q / W mod 2^L` is `q / (W 2^l) mod 2`.
-/

namespace NLogN.KFR

/-- Bit `l` of `z`, as an integer. -/
def bitZ (z l : ℕ) : ℤ := if z.testBit l then 1 else 0

end NLogN.KFR

end

section
-- Source module: Solutions.KFRLay
/-!
# KFR: the packed layout of three rows and the row addition as an address map

For a stride-`K` row family (`pos (r, c) + K (c + f r) = e`, `f = m^k`, `AffK2`), the rows
`a < b < c` (`σ : Fin 3 → Fin m` strictly monotone) are the three chunk slots of the rectangle
`rect N0 N1 N2 N3 (fK)` (`rectN`), with offset `ρ' = e mod K`: row `r` occupies the address bits
`[sR r, sR r + fK)` of the record index divided by `S₀`, its bit `f-1-u` at `sR r + ρ' + uK`
(`pos_eq_sR`).

* `fsetZ_rowAddr`: the record moved by `rowAddr i j` as a sum over the bits of row `j`.
* **`encode_rowAddAddr`**: `encode (rowAddAddr K ρ' f px py (decode q)) = fset q (rowAddr i j …)`
  for `σ px = i`, `σ py = j`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFR

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly KFN

/-! ## The moved record as a sum over the bits of row `j` -/

section Sum

variable {m f : ℕ}

/-- The contribution of bit `c` of row `j`. -/
def rowTerm (st : Fin m × Fin f → ℕ) (q : ℕ) (i j : Fin m) (c : ℕ) : ℤ :=
  if h : c < f then
    ((((q / st (j, ⟨c, h⟩) % 2 + q / st (i, ⟨c, h⟩) % 2) % 2 : ℕ) : ℤ) -
      ((q / st (j, ⟨c, h⟩) % 2 : ℕ) : ℤ)) * (st (j, ⟨c, h⟩) : ℤ)
  else 0

end Sum

/-! ## Rows of a stride-`K` family -/

section Rows

/-- Start of the field of row `r`. -/
def sR (K f e r : ℕ) : ℕ := e - K * (f * (r + 1) - 1) - e % K

variable {K m k e : ℕ} {pos : Fin m × Fin (m ^ k) → ℕ}

end Rows

end NLogN.KFR

end

section
-- Source module: Solutions.KFRAddr
/-!
# KFR: the packed rectangle of three rows, and the row addition as `rowAddAddr`

For a stride-`K` family `pos` (`pos (r, c) + K (c + f r) = e`, `f = m^k`) and three rows
`σ 0 < σ 1 < σ 2`, `shR` is the rectangle `rect N0 N1 N2 N3 (fK)` whose chunk slot `τ` is the
field of row `σ τ` (weight `2^{sR (σ τ)} S₀`).  **`encode_rowAddAddr`**: on this rectangle, the
address map `rowAddAddr K (e mod K) f px py` is `fset (rowAddr (σ px) (σ py))`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFR

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly KFN

variable {K m k e : ℕ} {pos : Fin m × Fin (m ^ k) → ℕ}

/-- The rectangle of the rows `σ 0 < σ 1 < σ 2`. -/
def shR (K f e t S0 B : ℕ) (σ : Fin 3 → Fin m) : Shape :=
  rect (B * 2 ^ (t - sR K f e (σ 0) - f * K)) (2 ^ (sR K f e (σ 0) - sR K f e (σ 1) - f * K))
    (2 ^ (sR K f e (σ 1) - sR K f e (σ 2) - f * K)) (2 ^ sR K f e (σ 2) * S0) (f * K)

section Shape

variable {t S0 B : ℕ} {σ : Fin 3 → Fin m}
  (he : ∀ p, pos p + K * (p.2.val + m ^ k * p.1.val) = e) (htop : e - e % K + K ≤ t)
  (hσ : σ 0 < σ 1 ∧ σ 1 < σ 2)
include he htop hσ

end Shape

end NLogN.KFR

end

section
-- Source module: Solutions.KFRPerm
/-!
# KFR: `permRecs` realises the row operator `Pm i j`, and the statement `Lem13ILP`

* `permRecs_recRep`, `permRecs_GB`: permuting records by an address map `π` represents
  `q ↦ xs[encode (π (decode q))]` with the same scale and the same entry bound.
* `selOpF_Pm`: `selOpF st n (Pm i j) xs = q ↦ xs[fset q (rowAddr i j (fbit q))]`.
* `slots3`: two distinct rows `i ≠ j` of `m ≥ 3` are two of three sorted rows `σ 0 < σ 1 < σ 2`.
* `Lem13ILP iprog`: the statement of `NLogN.KFL.lem13_il` (`Solutions/KFL_WIP.lean`), copied with
  `l12prog` replaced by a parameter `iprog` (so `Lem13ILP l12prog` is `lem13_il`).
* **`rowImpl_of_lem13`**: under `Lem13Body`, the row operator `Pm i j` of a stride-`K` family has an
  `SImplI` implementation (scale and bound unchanged) whose cost is the Lemma 13 bound `bnd`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFR

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly KFN

/-! ## Records permuted by an address map -/

/-! ## Three sorted rows -/

/-! ## The Lemma 13 statement -/

/-- The body of `NLogN.KFL.lem13_il` for a given program table, code family and constant. -/
def Lem13Body (iprog : Fin 1 → IStmt HS 1)
    (code : ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → Fin 3 → Fin 3 → IStmt HS 1) (Kc : ℝ) : Prop :=
  ∀ {I : List Sym} (N0 N1 N2 N3 K f ρ R w : ℕ) (px py : Fin 3), px ≠ py → 6 ≤ K → ρ < K →
    1 ≤ f → 1 ≤ N0 → 1 ≤ N1 → 1 ≤ N2 → 1 ≤ N3 → 1 ≤ R →
    ∀ (s : IState HS) (g : Fin HS → List ℤ) (zs : List GI) (l : List ℤ),
    g 0 = gflat zs ++ l → zs.length = vol (rect N0 N1 N2 N3 (f * K)) * R → NLogN.KFB.GB w zs →
    ∃ c, IRuns I iprog (code N0 N1 N2 N3 K f ρ R w px py) (encK s g)
        (encK s (Function.update g 0 (gflat (permRecs (rect N0 N1 N2 N3 (f * K))
          (rowAddAddr K ρ f px py) R zs) ++ l))) c ∧
      (c : ℝ) ≤ Kc * ((zs.length : ℝ) * (w + 1) * (((f * K : ℕ) : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50) + 1) +
        (vol (rect N0 N1 N2 N3 (f * K)) : ℝ) * ((zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1) ^ 3 +
        min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * (vol (rect N0 N1 N2 N3 (f * K)) : ℝ) *
          ((zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1) * ((R : ℝ) * (w + 1) +
            (zsz (2 * zs.length * (w + 1) : ℤ) : ℝ) + 1))

/-- **The statement of `NLogN.KFL.lem13_il`**, with `l12prog` as the parameter `iprog`. -/
def Lem13ILP (iprog : Fin 1 → IStmt HS 1) : Prop :=
  ∃ (code : ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → ℕ → Fin 3 → Fin 3 → IStmt HS 1) (Kc : ℝ), 0 < Kc ∧
    Lem13Body iprog code Kc

/-- The Lemma 13 cost for `n` records of `r` entries, `f` bits per row. -/
noncomputable def bnd (Kc : ℝ) (K f n r w : ℕ) : ℝ :=
  Kc * (((n * r : ℕ) : ℝ) * (w + 1) * (((f * K : ℕ) : ℝ) ^ (1 - (2 : ℝ)⁻¹ ^ 50) + 1) +
    (n : ℝ) * ((zsz (2 * ((n * r : ℕ) : ℤ) * ((w : ℤ) + 1)) : ℝ) + 1) ^ 3 +
    min 1 (80 * ((f : ℝ) - 1) * 2⁻¹ ^ K) * (n : ℝ) *
      ((zsz (2 * ((n * r : ℕ) : ℤ) * ((w : ℤ) + 1)) : ℝ) + 1) * ((r : ℝ) * (w + 1) +
        (zsz (2 * ((n * r : ℕ) : ℤ) * ((w : ℤ) + 1)) : ℝ) + 1))

end NLogN.KFR

end

section
-- Source module: Solutions.KFRBridge
/-!
# KFR: the bridge `Lem13ILP iprog → RowHyp … τ50`

* `PsiR r Z w = 2 (Z + zsz (w+1) + 2)^3 + ⌊r (w+1) zsz (w+1) / (Z+2)⌋ + 2` (monotone, `≥ Z+2`).
* `ceil_bnd_le`: when `80 Z (Z+2) ≤ 2^K` and `f ≤ Z`, the Lemma 13 cost `⌈bnd⌉` is at most
  `(Kc (K+2) + 1) · n · (r (w+1) (f+1)^τ50 + PsiR r Z w)`.  The guard term
  `min 1 (80 (f-1) 2^{-K})` is at most `1/(Z+2)`, which absorbs the `zsz` factors of the
  correction cost into `O(r (w+1))` plus `r (w+1) zsz(w+1)/(Z+2)` (inside `PsiR`).
* **`rowHyp_of_lem13`**: `Lem13ILP iprog → ∃ Kc > 0, ∀ K m r Z, 6 ≤ K → 3 ≤ m → 1 ≤ r →
  80 Z (Z+2) ≤ 2^K → RowHyp Isym iprog K m r Z (PsiR r Z) (Kc (K+2) + 1) τ50`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFR

open IntMul RS AModel KIL KFD KFB KFG KFH Complex P8Lift Butterfly KFN

/-- The additive per-record term of the row hypothesis obtained from Lemma 13. -/
def PsiR (r Z w : ℕ) : ℕ :=
  2 * (Z + zsz ((w : ℤ) + 1) + 2) ^ 3 + r * (w + 1) * zsz ((w : ℤ) + 1) / (Z + 2) + 2

end NLogN.KFR

namespace NLogN.KFR

open IntMul RS KIL KFH KFN

end NLogN.KFR

end

section
-- Source module: Solutions.KFRL12
/-!
# KFR: the bridge for the Lemma 12 program table `l12prog`

`Lem13IL` is the statement of `NLogN.KFL.lem13_il` (`Solutions/KFL_WIP.lean`, still `sorry`
there; not imported), i.e. `Lem13ILP KFL.l12prog`.  `rowHyp_l12` gives KFN's `RowHyp … τ50`
from it, and `closedHyps_l12` the full hypothesis set of `KFN.prop14_il_complex_closed`.
-/

namespace NLogN.KFR

open IntMul RS KIL KFH KFN

/-- **The statement of `NLogN.KFL.lem13_il`.** -/
def Lem13IL : Prop := Lem13ILP KFL.l12prog

end NLogN.KFR

end

section
-- Source module: Solutions.KFMTable
/-!
# KFM 7: one procedure table for Lemma 12 / Prop 14 and the FM multiplier

Procedure tables are merged by renaming procedure indices (`IStmt.ren id g`, nlogn `HRen`).

* **`IRuns.relabel`**, **`SImplI.relabel`**, **`RowHyp.relabel`**: runs, `SImplI` implementations
  and the row hypothesis move to any table `iprog'` that contains the old one at `g`
  (`iprog' (g j) = (iprog j).ren id g`), with the same costs.
* **`kTable ext`**: the table `[l12, mulWrap, mulProg 0..4, ext 0..E-1]` (`7 + E` procedures);
  `ext` are further procedures (pointwise passes, line maps, …) that may call everything.
* **`kTable_mulOK`**: procedure `1` satisfies `MulOK` with cost `mulCostH (zsz a + zsz b)`.
* **`prop14Ph_kTable`**: `prop14Ph_l12` (unconditional Prop 14, `λ = τ50`) for `kTable ext`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFM

open IntMul RS AModel KIL KFD KFB KFG KFH KFN Complex P8Lift Butterfly

section Relabel

variable {S P P' : ℕ} {I : List Sym}

end Relabel

section RelabelHS

variable {P P' : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P} {iprog' : Fin P' → IStmt HS P'}
  {g : Fin P → Fin P'}

end RelabelHS

/-! ## The merged table -/

section Table

variable {E : ℕ}

/-- Index of the Lemma 12 procedure. -/
def gL (E : ℕ) : Fin 1 → Fin (7 + E) := fun _ => ⟨0, by omega⟩

/-- Indices of the renamed `mulProg`. -/
def gM (E : ℕ) : Fin 5 → Fin (7 + E) := fun j => ⟨j.val + 2, by omega⟩

/-- Index of the `MulOK` wrapper. -/
def pMul (E : ℕ) : Fin (7 + E) := ⟨1, by omega⟩

/-- The merged table `[l12, mulWrap, mulProg 0..4, ext 0..E-1]`. -/
noncomputable def kTable (ext : Fin E → IStmt HS (7 + E)) : Fin (7 + E) → IStmt HS (7 + E) :=
  fun i =>
    if h0 : i.val = 0 then (KFL.l12prog 0).ren id (gL E)
    else if h1 : i.val = 1 then mulWrap (gM E)
    else if h7 : i.val < 7 then (mulProg ⟨i.val - 2, by omega⟩).ren ssE (gM E)
    else ext ⟨i.val - 7, by omega⟩

end Table

end NLogN.KFM

end

section
-- Source module: Solutions.KFMP18
/-!
# KFM 14: Proposition 18 as an explicit IL hypothesis `P18IL`

`P18IL` is the interface between the assembly (item 14) and Prop 18 in IL (item 13).

* The procedure table is **one fixed table** `kTable ext` (`ext`, `entry` and the descriptor stack
  `dS` are chosen before `n`). The final machine must not depend on `n`, so every `n`-dependent
  quantity (`p, ℓ, K, D`, the axis lengths `t_i`) is run-time data on stack `dS`, not a literal
  in the code.
* Data: two arrays `f, g` of `N = ∏ t_i` records in row-major order over `t_1 × … × t_D`
  (last axis fastest). Each record is `r = 2^ℓ` Gaussian coefficients at scale `2^p` (`RecRep`),
  in the unit disk, width `W = p + 2`. They sit on stack `0` (`f` on top). All other managed
  stacks are empty and the rest of the state (`s`) is arbitrary and preserved.
* Result: one array `h` (`RecRep r p`, width `W`) replacing `f, g` on stack `0`, within
  `2 N (3ℓ + 1) 2^{-p}` of `M_r(f, g) = (f ∗ g)/(r N)` in every coefficient (the bound of
  `GlueProp18.prop18_scaled_full`, with `√2 ≤ 2`, transported to the row-major order).
* Cost `≤ C₀ · N r (W + 1) · Ncost(6 lg n, d, K, ℓ, α) · (log(6 lg n) + 1)^A` for `n ≥ n₁`,
  with `D + 1 = d = dP n`, `ℓ = ℓP n`, `K = KP n`, `p = pP n`, `t_i ∈ {r/2, r}`.
-/

namespace NLogN.KFM

open IntMul KIL KFB KFH IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-- Negacyclic product of two coefficient vectors (`C[y]/(y^r + 1)`). -/
def ncm {r : ℕ} (u v : Fin r → ℂ) : Fin r → ℂ := fun k =>
  ∑ i : Fin r, ∑ j : Fin r,
    if i.val + j.val = k.val then u i * v j
    else if i.val + j.val = k.val + r then -(u i * v j) else 0

/-- Row-major index of `j - k` (componentwise, mod `t_i`), last axis fastest. -/
def rmSub : List ℕ → ℕ → ℕ → ℕ
  | [], _, _ => 0
  | t :: ts, j, k =>
    ((j / ts.prod % t + t - k / ts.prod % t) % t) * ts.prod + rmSub ts (j % ts.prod) (k % ts.prod)

/-- The normalized cyclic convolution `M_r(f, g) = (f ∗ g)/(r N)` over `∏ ℤ/t_i` with values in
`C[y]/(y^r + 1)`, on row-major lists of coefficient vectors. -/
def mrL (r : ℕ) (ts : List ℕ) (xf xg : List (Fin r → ℂ)) : List (Fin r → ℂ) :=
  (List.range ts.prod).map fun j =>
    ((r * ts.prod : ℕ) : ℂ)⁻¹ • ∑ k ∈ Finset.range ts.prod, ncm (xf.getD k 0) (xg.getD (rmSub ts j k) 0)

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMBridge
/-!
# KFM 14: `P18IL`'s row-major lists vs. the box of the error chain

`P18IL` speaks about row-major lists of records (`mrL`, `rmSub`); `chain_concrete` about functions
on the box `∏_{i ≤ D} ℤ/t_i` (`mrX`, last axis = record coefficient).

* `digL ts x i`: digit `i` of the row-major index `x` (first axis most significant);
  `rmSub_lt`, **`digL_rmSub`** (digits of `rmSub` are the differences mod `t_i`), `digL_lt`,
  `digL_inj`.
* `fullIdx t j c`: the box point of record `j`, coefficient `c`; **`fullIdx_bij`** (a bijection
  `Fin N × Fin r ≃ box`), `recL t f` (box function → row-major list of records),
  `boxOf t xs` (its inverse), `recL_boxOf`, `boxOf_recL`.
* **`mrL_box`**: `mrL r ts (recL f) (recL g) = recL (mrX t f g)`.
* **`mrt_err_of_list`**: a coefficientwise list bound against `mrL` (the form of `P18IL`) gives
  the sup-norm bound `‖boxOf xh - mrX f g‖ ≤ E` that `chain_concrete` needs (`hMrt`).
-/

namespace NLogN.KFM

open Complex

noncomputable section

/-! ## Row-major digits -/

/-- Digit `i` of the row-major index `x` over the axis lengths `ts`. -/
def digL : List ℕ → ℕ → ℕ → ℕ
  | [], _, _ => 0
  | t :: ts, x, 0 => x / ts.prod % t
  | _ :: ts, x, i + 1 => digL ts (x % ts.prod) i

lemma prod_pos_of {ts : List ℕ} (h : ∀ t ∈ ts, 0 < t) : 0 < ts.prod :=
  List.prod_pos h

theorem digL_lt : ∀ (ts : List ℕ), (∀ t ∈ ts, 0 < t) → ∀ x i (hi : i < ts.length),
    digL ts x i < ts[i]
  | [], _, _, _, hi => by simp at hi
  | t :: ts, h, x, i, hi => by
    have ht : 0 < t := h t (by simp)
    have h' : ∀ u ∈ ts, 0 < u := fun u hu => h u (by simp [hu])
    rcases i with _ | i
    · simp only [digL, List.getElem_cons_zero]; exact Nat.mod_lt _ ht
    · simp only [digL, List.getElem_cons_succ]; exact digL_lt ts h' _ i (by simpa using hi)

theorem digL_inj : ∀ (ts : List ℕ), (∀ t ∈ ts, 0 < t) → ∀ x y, x < ts.prod → y < ts.prod →
    (∀ i (hi : i < ts.length), digL ts x i = digL ts y i) → x = y
  | [], _, x, y, hx, hy, _ => by simp at hx hy; omega
  | t :: ts, h, x, y, hx, hy, hd => by
    have h' : ∀ u ∈ ts, 0 < u := fun u hu => h u (by simp [hu])
    have hR := prod_pos_of h'
    simp only [List.prod_cons] at hx hy
    have hx' : x / ts.prod < t := Nat.div_lt_of_lt_mul (by rwa [Nat.mul_comm])
    have hy' : y / ts.prod < t := Nat.div_lt_of_lt_mul (by rwa [Nat.mul_comm])
    have h0 := hd 0 (by simp)
    simp only [digL] at h0
    rw [Nat.mod_eq_of_lt hx', Nat.mod_eq_of_lt hy'] at h0
    have hm := digL_inj ts h' (x % ts.prod) (y % ts.prod) (Nat.mod_lt _ hR) (Nat.mod_lt _ hR)
      (fun i hi => by have := hd (i + 1) (by simpa using hi); simpa [digL] using this)
    rw [← Nat.div_add_mod x ts.prod, ← Nat.div_add_mod y ts.prod, h0, hm]

/-! ## The box and the record lists -/

variable {D : ℕ}

/-- The axis lengths of the first `D` axes, as a list. -/
def tsOf (t : Fin (D + 1) → ℕ) : List ℕ := List.ofFn fun i : Fin D => t i.castSucc

lemma tsOf_length (t : Fin (D + 1) → ℕ) : (tsOf t).length = D := by simp [tsOf]

lemma tsOf_get (t : Fin (D + 1) → ℕ) (i : ℕ) (hi : i < (tsOf t).length) :
    (tsOf t)[i] = t ⟨i, by rw [tsOf_length] at hi; omega⟩ := by
  simp [tsOf]

lemma tsOf_pos (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)] : ∀ u ∈ tsOf t, 0 < u := by
  intro u hu
  simp only [tsOf, List.mem_ofFn] at hu
  obtain ⟨i, rfl⟩ := hu
  exact Nat.pos_of_ne_zero (NeZero.ne _)

lemma tsOf_prod (t : Fin (D + 1) → ℕ) : (tsOf t).prod = ∏ i : Fin D, t i.castSucc := by
  simp [tsOf, List.prod_ofFn]

/-- The box point of record `j`, coefficient `c`. -/
def fullIdx (t : Fin (D + 1) → ℕ) (j c : ℕ) : (i : Fin (D + 1)) → ZMod (t i) :=
  fun i => ((if i.val < D then digL (tsOf t) j i.val else c : ℕ) : ZMod (t i))

lemma val_cast_lt {m x : ℕ} [NeZero m] (h : x < m) : ((x : ℕ) : ZMod m).val = x := by
  rw [ZMod.val_natCast, Nat.mod_eq_of_lt h]

theorem fullIdx_inj (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)] {j j' c c' : ℕ}
    (hj : j < (tsOf t).prod) (hj' : j' < (tsOf t).prod) (hc : c < t (Fin.last D))
    (hc' : c' < t (Fin.last D)) (h : fullIdx t j c = fullIdx t j' c') : j = j' ∧ c = c' := by
  have hpos := tsOf_pos t
  constructor
  · refine digL_inj _ hpos _ _ hj hj' fun i hi => ?_
    have hiD : i < D := by rw [tsOf_length] at hi; exact hi
    have e := congrFun h ⟨i, by omega⟩
    simp only [fullIdx, hiD, if_true] at e
    have l1 := digL_lt _ hpos j i hi
    have l2 := digL_lt _ hpos j' i hi
    rw [tsOf_get] at l1 l2
    have := congrArg ZMod.val e
    rwa [val_cast_lt l1, val_cast_lt l2] at this
  · have e := congrArg ZMod.val (congrFun h (Fin.last D))
    simp only [fullIdx, Fin.val_last, lt_irrefl, if_false] at e
    rwa [val_cast_lt hc, val_cast_lt hc'] at e

/-- The map `Fin N × Fin r → box`. -/
def phiB (t : Fin (D + 1) → ℕ) : Fin (tsOf t).prod × Fin (t (Fin.last D)) → ((i : Fin (D + 1)) → ZMod (t i)) :=
  fun x => fullIdx t x.1 x.2

theorem phiB_bij (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)] : Function.Bijective (phiB t) := by
  rw [Fintype.bijective_iff_injective_and_card]
  refine ⟨fun x y h => ?_, ?_⟩
  · obtain ⟨h1, h2⟩ := fullIdx_inj t x.1.2 y.1.2 x.2.2 y.2.2 h
    exact Prod.ext (Fin.ext h1) (Fin.ext h2)
  · simp only [Fintype.card_prod, Fintype.card_fin, Fintype.card_pi, ZMod.card, tsOf_prod]
    rw [Fin.prod_univ_castSucc]

/-- Box function → row-major list of records. -/
def recL (t : Fin (D + 1) → ℕ) (f : ((i : Fin (D + 1)) → ZMod (t i)) → ℂ) :
    List (Fin (t (Fin.last D)) → ℂ) :=
  (List.range (tsOf t).prod).map fun j c => f (fullIdx t j c)

/-- Row-major list of records → box function. -/
def boxOf (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)] (xs : List (Fin (t (Fin.last D)) → ℂ)) :
    ((i : Fin (D + 1)) → ZMod (t i)) → ℂ := fun K =>
  let x := (Equiv.ofBijective _ (phiB_bij t)).symm K
  xs.getD x.1 0 x.2

/-! ## The ring convolution on both sides -/

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMNeed
/-!
# KFM: the uniform interfaces `PhiNeed` and `SwapNeed` (hypotheses of the assembly)

These are stated in the style of `P18IL`. Each one asks for a fixed table `kTable ext`, with
`ext`, the entry points and the descriptor stack `dS` all chosen before any parameter. Every
size-dependent quantity is run-time data on `dS`. The data sit on stack `0` (array on top, then
an arbitrary tail `l`). Every other managed stack is empty (`f j = []` for `j ≠ 0, dS`), so a
procedure may use any stack as scratch as long as it leaves it as it found it. The unmanaged
part `s` underneath is arbitrary and preserved. On exit, `dS` still holds its descriptor.

* **`PhiNeed`**: uniform `Φ` / `Φ⁻¹` of §9.3, with the denotations `KFU.PhiD` / `KFU.PhiInvD`.
  These are KFU's `phi_spec` / `phiInv_spec` with the literals `w, h, d, bL, s_j, μ_j` moved into
  the descriptor `descPhi`. Cost: a constant times KFU's `cPhi` (Lemma 12 constant `1`), plus a
  descriptor term `(d+1)(h+w+size V_S+1)²`.
* **`SwapNeed`**: uniform Lemma 12 on Gaussian (record) arrays: interchange two axes `p < q` of
  equal size `2^u` in an array of any positive shape `sh`. This is KFS's `swAxG_spec`, with
  `w, u, p, q, sh` in the descriptor `descSw`. Records of `r` entries use the shape `sh ++ [r]`.
  Cost: `K (w+1)(V+1) u^τ`, plus a descriptor term `(|sh|+1)(size V+w+u+1)²`. The assembly
  builds the axis exposure (moving an axis innermost before a line map, and back after it)
  from these swaps. A `2^{h+1}` axis is split as `2 × 2^h` by a free reshape.

* **`TransNeed`**: uniform axis rotation (`transposeL t (V/t)`) on padded Gaussian arrays whose
  axes are all `2^h` or `2^{h+1}`. This is the exposure step that replaces nlogn's `transP 2`
  inside `GL.axisP`. Equal-size swaps cannot do it when the sizes are mixed (see the docstring).

The final theorem may take `P18IL ∧ PhiNeed ∧ TransNeed` (and `SwapNeed` if it is used). The tables are merged by relabelling
(`IRuns.relabel`).
-/

namespace NLogN.KFM

open IntMul KIL KFB AModel AModelL12

noncomputable section

/-! ## `Φ` / `Φ⁻¹` -/

/-! ## Axis swaps (Lemma 12 on Gaussian record arrays) -/

/-- The descriptor of a swap call: `[w, u, p, q, |sh|, sh_0, …]`. -/
def descSw (w u p q : ℕ) (sh : Shape) : List ℤ :=
  (w : ℤ) :: (u : ℤ) :: (p : ℤ) :: (q : ℤ) :: (sh.length : ℤ) :: sh.map (fun a => (a : ℤ))

/-- **Lemma 12 on Gaussian arrays in IL, uniform** (hypothesis of the assembly). It is
`KFS.swAxG_spec` with the code fixed and `w, u, p, q, sh` read from `dS`. -/
def SwapNeed : Prop :=
  ∃ (E : ℕ) (ext : Fin E → IStmt HS (7 + E)) (entry : Fin (7 + E)) (dS : Fin HS) (K : ℝ),
    dS ≠ 0 ∧ 0 ≤ K ∧
  ∀ (Isym : List Sym) (u w : ℕ) (sh : Shape) (p q : ℕ), 1 ≤ u → (∀ a ∈ sh, 1 ≤ a) →
    p < q → q < sh.length → sh.getD p 0 = 2 ^ u → sh.getD q 0 = 2 ^ u →
  ∀ (s : IState HS) (f : Fin HS → List ℤ) (zs : List GI) (l : List ℤ),
    (∀ j, j ≠ 0 → j ≠ dS → f j = []) →
    f dS = descSw w u p q sh →
    f 0 = gflat zs ++ l → zs.length = vol sh → NLogN.KFB.GB w zs →
    ∃ c : ℕ, IRuns Isym (kTable ext) (.call entry) (encK s f)
        (encK s (Function.update f 0 (gflat (pull sh sh (swapPos p q) 0 zs) ++ l))) c ∧
      (c : ℝ) ≤ K * (((w : ℝ) + 1) * ((vol sh : ℝ) + 1) * (u : ℝ) ^ KFS.τS +
        (((sh.length + 1) * (Nat.size (vol sh) + w + u + 1) ^ 2 : ℕ) : ℝ))

/-! ## Axis rotation on padded arrays (the exposure step around the line maps) -/

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMPw
/-!
# KFM 9a: the pointwise Gaussian product pass `Q_Δ(θ · f)` with `MulOK` products

`gmulP pM Δ N`: stack `0` holds `gflat zs ++ R0` (`N` Gaussian entries on top), stack `pW` holds
`gflat ts ++ RW` (the multipliers, numerators at scale `2^Δ`), stack `pR` holds the ruler `2^Δ`.
The pass replaces the entries by `gmulT Δ z t = ((z·t).re.tdiv 2^Δ, (z·t).im.tdiv 2^Δ)`
(four `MulOK` products per entry, `split` by the ruler for `tdiv`) and consumes the multipliers:

* **`entry_runs`**: one entry; result pushed (imaginary part last) on `pO`.
* **`gmulP_runs`**: the pass, `0 = gflat (zipWith (gmulT Δ) zs ts) ++ R0`, `pW = RW`, at cost
  `N (4 φ + 10^5 (w + Δ + 1)) + 10^4 (N + 1)` when all products cost `≤ φ` on `w`-bit factors.

On numerators (`RecRep`) this is the twist / chirp step `Q_p(θ̃ · f)` of §9.4–9.5
(`gmulT_eq`: `gmulT Δ z t = KFH.gtr Δ (gmul z t)`, and `toC_gmul`).
**`gmulP_kTable`**: the FM instance in the merged table, cost `N·4·fmB(2w) + 10^6 (N+1)(w+Δ+1)`.

Generic per-item pass **`itemP body N`** (`itemP_runs`: `x ↦ f x` on `N` integer items of
stack `0`) with three instances for §9.3/§9.6: **`shlP_runs`** (`× 2^γ`), **`roundP_runs`**
(`(x + h).tdiv 2^Δ`, i.e. rounding when `h = 2^{Δ-1}`, `x + h ≥ 0`), **`smulP_runs`**
(`(S x).tdiv 2^Δ`, one `MulOK` product per item).
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT KIL

/-! ## Stacks -/

abbrev pW : Fin HS := 43
abbrev pO : Fin HS := 44
abbrev pR : Fin HS := 45
abbrev pX1 : Fin HS := 46
abbrev pX2 : Fin HS := 47
abbrev pX3 : Fin HS := 48
abbrev pX4 : Fin HS := 49
abbrev pP1 : Fin HS := 50
abbrev pP2 : Fin HS := 51
abbrev pT : Fin HS := 53
abbrev pA : Fin HS := 60
abbrev pC : Fin HS := 61
abbrev pD : Fin HS := 62

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

/-! ## One product through `MulOK` -/

/-- Push `top j * top i` on `k` via the multiplication procedure (the `RM.mul` code). -/
def mulC (pM : Fin P) (k i j : Fin HS) : IStmt HS P :=
  .seq (.op (.dup i ARGM)) (.seq (.op (.dup j ARGM)) (.seq (.call pM) (.op (.mv ARGM k))))

/-- The cost of `mulC` on factors `a` (from `i`) and `c` (from `j`). -/
def cMulC (M : MulSpec) (a c : ℤ) : ℕ :=
  (dupCost (zsz a) + 17) + (dupCost (zsz c) + 17) + (M.Φ c a + 2) + (2 * moveCost (zsz (c * a)) + 19)

/-! ## Combining two products: `(u ± v).tdiv 2^Δ` pushed on `pO` -/

/-- `neg` (if `sub`) then `add`, `split` by the ruler on `pR`, drop the remainder, push on `pO`. -/
def combC (sub : Bool) : IStmt HS P :=
  .seq (if sub then .op (.neg pP2) else .skip) (.seq (.op (.add pP1 pP2 pP1))
    (.seq (.op (.split pP1 pR pT)) (.seq (.op (.drop pT)) (.op (.mv pP1 pO)))))

/-! ## One entry -/

/-- The Gaussian product with truncation, `Q_Δ`. -/
def gmulT (Δ : ℕ) (z t : GI) : GI :=
  ((z.1 * t.1 - z.2 * t.2).tdiv (2 ^ Δ), (z.1 * t.2 + z.2 * t.1).tdiv (2 ^ Δ))

/-- One entry: pop `a, b` from `0` and `c, d` from `pW`, push `Q_Δ((a+bi)(c+di))` on `pO`. -/
def entryB (pM : Fin P) : IStmt HS P :=
  .seq (.seq (.op (.mv 0 pX1)) (.seq (.op (.mv 0 pX2)) (.seq (.op (.mv pW pX3)) (.op (.mv pW pX4)))))
    (.seq (mulC pM pP1 pX1 pX3) (.seq (mulC pM pP2 pX2 pX4) (.seq (combC true)
    (.seq (mulC pM pP1 pX1 pX4) (.seq (mulC pM pP2 pX2 pX3) (.seq (combC false)
    (.seq (.op (.drop pX1)) (.seq (.op (.drop pX2)) (.seq (.op (.drop pX3))
      (.op (.drop pX4)))))))))))

/-! ## Moving the results back -/

/-- Move one item from `pO` to `0`. -/
def mvB : IStmt HS P := .op (.mv pO 0)

/-! ## The pass -/

/-- **The pointwise Gaussian product pass** on `N` entries. -/
def gmulP (pM : Fin P) (N : ℕ) : IStmt HS P :=
  .seq (rpt pA pC pD N (entryB pM)) (rpt pA pC pD (2 * N) mvB)

/-- The cost of `gmulP` on `N` entries of `w` bits with products `≤ φ`. -/
def cGmul (N w Δ φ : ℕ) : ℕ :=
  N * (4 * φ + 100000 * (w + Δ + 1) + 64) + 60 + cRpt N +
    ((2 * N) * (110 * (4 * w + 2) + 64) + 60 + cRpt (2 * N))

/-! ## Generic per-item pass -/

/-- Stacks that a per-item body may change: `0`, `pO`, and the `rpt` counters. -/
def Fixd (s st : IState HS) : Prop :=
  ∀ j : Fin HS, j ≠ 0 → j ≠ pO → j ≠ pA → j ≠ pC → j ≠ pD → st.G j = s.G j

/-- **Per-item pass**: apply `body` (top of `0` ↦ pushed on `pO`) to `N` items, move back. -/
def itemP (body : IStmt HS P) (N : ℕ) : IStmt HS P :=
  .seq (rpt pA pC pD N body) (rpt pA pC pD N mvB)

/-! ## Instances: `× 2^γ`, `Q_Δ(· × S)`, rounding -/

/-- `x ↦ x · 2^γ` with the ruler `2^γ` on `pR`. -/
def shlB : IStmt HS P := .seq (.op (.mv 0 pP1)) (.seq (.op (.shl pP1 pR)) (.op (.mv pP1 pO)))

/-- `x ↦ (x + h).tdiv 2^Δ` with `h` on `pX2` and the ruler `2^Δ` on `pR` (rounding when
`h = 2^{Δ-1}` and `x + h ≥ 0`). -/
def roundB : IStmt HS P := .seq (.op (.mv 0 pP1)) (.seq (.op (.dup pX2 pP2)) (combC false))

/-- `x ↦ (x · S).tdiv 2^Δ` with `S` on `pX3` and the ruler `2^Δ` on `pR` (one `MulOK` call). -/
def smulB (pM : Fin P) : IStmt HS P :=
  .seq (.op (.mv 0 pX1)) (.seq (mulC pM pP1 pX1 pX3) (.seq (.op (.const pP2 0))
    (.seq (combC false) (.op (.drop pX1)))))

/-! ## The three passes -/

/-- Cost of a per-item pass with per-item cost `c` and output width `B`. -/
def cItem (N c B : ℕ) : ℕ :=
  N * (c + 64) + 60 + cRpt N + (N * (110 * (B + 1) + 64) + 60 + cRpt N)

/-! ## Denotation and the FM instance -/

/-- The closed FM cost bound as a function of the total bit length. -/
def fmB (n : ℕ) : ℕ :=
  (8 * fmA + 10 ^ 7 + 10 ^ 5) * (n + 1) * (zsz ((n + 1 : ℕ) : ℤ) + 2) *
    (zsz (((zsz ((n + 1 : ℕ) : ℤ) + 2) ^ 4 : ℕ) : ℤ) + 1)

/-- The FM multiplication spec of `KFMMul`. -/
def fmSpec : MulSpec := ⟨fun a b => mulCostH (zsz a + zsz b), fun _ _ => True⟩

end NLogN.KFM

end

section
-- Source module: Solutions.KFMRound
/-!
# KFM 4c: the closed per-round bound (MACHINE-PLAN item 4, §5 (C))

The cost of one Proposition 14 round in the merged table (`prop14Ph_kTable`) is
`kC w + cWrapPh …`, with the row constant
`F1 = 2((k+1)m(10^5 + k a + 2 k b CR) D^{τ50} + 1) + 10^6(71k+2)`, `k = ⌊log_m D⌋`, `m = 100^3`,
`CR = Kc(K+2)+1`, and the width `w' = w + Ew(D)` inside `Xf r Ψ`.

**`round_cost_closed`**: for fixed `A`, `Kc` there is `C₀` such that, whenever the three bounds of
`prop14Ph_kTable` hold for `E`, `dC`, `kC w`, and

* `K ≤ D^{2^{-56}}` (the layer chunk width `KP ≈ d^c`),
* `Ew(D) ≤ w + 1` (`p ≫ D^{4.6} G²`),
* `w + Ew(D) + 2 < 2^Z` and `2(2Z+4)^3 + 2 ≤ r` (`r ≫ Z³`),
* `(K D + P + 10)(Z + 2) ≤ vol·r`,

the round costs at most `C₀ · vol·r · (w+1) · D^{λ'}`, `λ' = 1 − 2^{-54}`.

The key steps are `F1_le` (`F1 ≤ c (k+1)²(K+3) D^{τ50}`), `logsq_le`
(`(⌊log_m D⌋+1)² ≤ (4/g²+2) D^g` for any `g > 0`), and `Xf_le` (`Xf r Ψ w' ≤ 6 r (w+1)`).
The exponent gap is `g = λ' − τ50 − 2^{-56} = 59·2^{-56}`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false

namespace NLogN.KFM

open IntMul KFN KFR AModel

/-! ## `(log D)²` against a small power of `D` -/

/-! ## The row constant -/

/-! ## The width term -/

/-! ## The closed bound -/

/-- `Ew(D)`: the width growth `w' − w` inside `Xf` in `prop14Ph_kTable`. -/
noncomputable def EwB (A D : ℕ) : ℕ :=
  2 * D + (Nat.log (100 ^ 3) D + 1) * 100 ^ 3 * (4 * (Nat.log (100 ^ 3) D + 1) ^ 2 *
      NLogN.KFN.Gb P8Lift.complexIface.net.length P8Lift.sc P8Lift.Wc A ^ (Nat.log (100 ^ 3) D + 2)) +
    20 * (Nat.log (100 ^ 3) D + 1) ^ 3 * NLogN.KFN.Gb P8Lift.complexIface.net.length P8Lift.sc P8Lift.Wc A ^ (Nat.log (100 ^ 3) D + 2) +
    2

/-- The exponent gap `λ' − τ50 − 2^{-56}`. -/
noncomputable def gapR : ℝ := (1 - (2 : ℝ)⁻¹ ^ 54) - τ50 - (2 : ℝ)⁻¹ ^ 56

end NLogN.KFM

end

section
-- Source module: Solutions.KFMTabM
/-!
# KFM 9b: fixed-point root tables by FM product recurrences — denotations and error bounds

All values are Gaussian numerators at scale `2^P` (`toC z / 2^P`); products are
`gmulT P z w = Q_P(z · w)` (the per-entry operation of `KFMPw.gmulP`).

* `geoL P c x n = [x, x c, x c², …]` (`n` entries, each step `Q_P(· c)`).
* `chirpL P c h g n`: `g_0 = g`, `g_{j+1} = Q_P(g_j h_j)`, `h_{j+1} = Q_P(h_j c)` (the quadratic
  phase `e^{-πi s j²/t}` with `h_0 ≈ e^{-πi s/t}`, `c ≈ e^{-2πi s/t}`).
* `prodF P n x = Q_P(⋯Q_P(1·x₀)⋯x_{n-1})`; `chirpEnt P s t j` is the `prodF` of the `d` axis
  factors `axTab P (s i) (t i)` at `j i` (`axTab` = `chirpL` from the `twiddle` roots `hRoot`,
  `cRoot`); `twTab P k j n = geoL` from `twiddle P k j` (twist table `≈ e^{2πi j m/k}`).

Error bounds (unit `u = 2^{-P}`):
* **`step_err`**: `‖Q_P(z w) - Z W‖ ≤ ‖z - Z‖ + 2‖w - W‖ + 2u` for `‖Z‖, ‖W‖ ≤ 1`, `‖z - Z‖ ≤ 1`.
* **`geoL_err`**: entry `k` of `geoL` is within `e₀ + k (2 e_c + 2u)` of `X C^k`.
* **`chirpL_err`**: entry `j` is within `4 j² u` of `G H … ` (`e(g₀) = 0`, `e(h₀), e(c) ≤ u`),
  provided `4 n² u ≤ 1`.
* **`chirp_axis_err`**: with `h₀`, `c` within `u` of `e^{-πi s/t}`, `e^{-2πi s/t}`, entry `j` of
  `chirpL` is within `4 t² u` of `gch s t j` (`RC.gch`, the chirp factor of `RC.chirp`).
* **`prodF_err`**: a product of `n` factors within `e` of unit-disc values is within
  `n (2e + 2u)` of the exact product.
* **`axTab_err`** (`≤ 4t²u` vs `RC.gZ s t k`), **`chirpEnt_err`** (`≤ d(8T²+2)u` vs `RC.chirp`),
  **`twTab_err`** (`≤ 4nu`), **`gtr_err`** (final rounding to precision `p`: `+ 2·2^{-p}`).
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFM

open Complex KFB

noncomputable section

/-! ## The value of a numerator -/

/-- The complex value of a numerator at scale `2^P`. -/
def vP (P : ℕ) (z : GI) : ℂ := NLogN.KFB.toC z / 2 ^ P

/-! ## Geometric tables -/

/-- `[x, x c, x c², …]` with `Q_P` after every product. -/
def geoL (P : ℕ) (c : GI) : GI → ℕ → List GI
  | _, 0 => []
  | x, n + 1 => x :: geoL P c (gmulT P x c) n

theorem length_geoL (P : ℕ) (c : GI) : ∀ (x : GI) (n : ℕ), (geoL P c x n).length = n
  | _, 0 => rfl
  | x, n + 1 => by simp [geoL, length_geoL P c _ n]

/-! ## Quadratic-phase (chirp) tables -/

/-- `g_0 = g`, `g_{j+1} = Q_P(g_j h_j)`, `h_{j+1} = Q_P(h_j c)`. -/
def chirpL (P : ℕ) (c : GI) : GI → GI → ℕ → List GI
  | _, _, 0 => []
  | h, g, n + 1 => g :: chirpL P c (gmulT P h c) (gmulT P g h) n

theorem length_chirpL (P : ℕ) (c : GI) : ∀ (h g : GI) (n : ℕ), (chirpL P c h g n).length = n
  | _, _, 0 => rfl
  | h, g, n + 1 => by simp [chirpL, length_chirpL P c _ _ n]

/-! ## Products of `d` factors -/

/-- The numerator of `1` at scale `2^P`. -/
def oneP (P : ℕ) : GI := (2 ^ P, 0)

/-- `Q_P(⋯Q_P(Q_P(1·x₀)·x₁)⋯ x_{n-1})`. -/
def prodF (P : ℕ) : (n : ℕ) → (Fin n → GI) → GI
  | 0, _ => oneP P
  | n + 1, x => gmulT P (prodF P n (fun i => x i.castSucc)) (x (Fin.last n))

/-! ## Base roots from `twiddle` -/

/-- `h₀ ≈ e^{-πi s/t}`. -/
def hRoot (P s t : ℕ) : GI := IntMul.HvdH.twiddle P (2 * t) (2 * t - s)
/-- `c ≈ e^{-2πi s/t}`. -/
def cRoot (P s t : ℕ) : GI := IntMul.HvdH.twiddle P t (t - s)

/-! ## The per-axis chirp table -/

/-- Axis table `[ĝ_0, …, ĝ_{t-1}]`, `ĝ_j ≈ e^{-πi s j²/t}` (`chirpL` from `1`, `h₀`, `c`). -/
def axTab (P s t : ℕ) : List GI := chirpL P (cRoot P s t) (hRoot P s t) (oneP P) t

theorem length_axTab (P s t : ℕ) : (axTab P s t).length = t := length_chirpL _ _ _ _ _

/-! ## Chirp entries -/

/-- The table value of `RC.chirp s t j`: the `Q_P`-product of the `d` axis-table factors. -/
def chirpEnt {d : ℕ} (P : ℕ) (s t : Fin d → ℕ) (j : (i : Fin d) → ZMod (t i)) : GI :=
  prodF P d (fun i => (axTab P (s i) (t i)).getD (j i).val 0)

/-! ## Twist tables -/

/-- `[1, ω, ω², …, ω^{n-1}]`, `ω ≈ e^{2πi j/k}`. -/
def twTab (P k j n : ℕ) : List GI := geoL P (IntMul.HvdH.twiddle P k j) (oneP P) n

theorem length_twTab (P k j n : ℕ) : (twTab P k j n).length = n := length_geoL _ _ _ _

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMTabP
/-!
# KFM 9c: the per-axis root-table program (`chirpL` recurrence in IL)

The recurrence `g_{k+1} = Q_Δ(g_k h_k)`, `h_{k+1} = Q_Δ(h_k c)` of `KFMTabM.chirpL` is run by
`chirpG pM n`:

* state: `0 = [g.re, g.im] ++ R0`, `pH = [h.re, h.im] ++ RH`, `pK = [c.re, c.im] ++ RK`,
  ruler `2^Δ` on `pR`;
* one iteration `chirpB` copies `g` to the output stack `pV` and replaces `(h, g)` by
  `(Q_Δ(h c), Q_Δ(g h))` with two `entryB` calls (eight `MulOK` products);
* after `n` iterations `pV` holds the entries in reverse order; `revP` moves them in order onto
  `pW`, and the state pairs are dropped.

`chirpG_runs`: `pW = gflat (chirpL Δ c h g n) ++ RW`, everything else restored, at cost
`n (8 φ + 3·10^5 (w + Δ + 1)) + 10^4 (n + 1)(w + 1)` when all states stay within `w` bits.

With `c = oneP Δ` the same program produces the geometric table `geoL` (`geoL_eq_chirpL`).
The state errors (`cst_err`) give the size bound `zsz ≤ Δ + 2` (`zsz_of_near`).
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT KIL

/-! ## Stacks -/

abbrev pH : Fin HS := 56
abbrev pK : Fin HS := 57
abbrev pU : Fin HS := 58
abbrev pV : Fin HS := 59

/-! ## Recurrence states -/

/-- One step `(h, g) ↦ (Q_Δ(h c), Q_Δ(g h))`. -/
def cstep (Δ : ℕ) (c : GI) (x : GI × GI) : GI × GI := (gmulT Δ x.1 c, gmulT Δ x.2 x.1)

/-- The state after `k` steps. -/
def cst (Δ : ℕ) (c h g : GI) (k : ℕ) : GI × GI := (cstep Δ c)^[k] (h, g)

/-! ## State errors and sizes -/

/-! ## Pair moves -/

section IL

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

/-- Copy the top pair of `X` onto `Y` (through `pU`). -/
def cpyB (X Y : Fin HS) : IStmt HS P :=
  .seq (.op (.mv X pU)) (.seq (.op (.dup X Y)) (.seq (.op (.dup pU Y)) (.op (.mv pU X))))

/-- Move the top pair of `X` onto `Y`, order kept (through `pU`). -/
def mvPB (X Y : Fin HS) : IStmt HS P :=
  .seq (.op (.mv X pU)) (.seq (.op (.mv X Y)) (.op (.mv pU Y)))

def cPair (a b : ℤ) : ℕ := 1000 * (zsz a + zsz b + 2)

/-! ## One iteration -/

/-- Move four items from `pO`: two onto `pH`, two onto `0`. -/
def mv4B : IStmt HS P :=
  .seq (.op (.mv pO pH)) (.seq (.op (.mv pO pH)) (.seq (.op (.mv pO 0)) (.op (.mv pO 0))))

/-- One iteration: output `g`, then `(h, g) ↦ (Q_Δ(h c), Q_Δ(g h))`. -/
def chirpB (pM : Fin P) : IStmt HS P :=
  .seq (cpyB 0 pV) (.seq (cpyB pH pW) (.seq (entryB pM) (.seq (cpyB pK pW) (.seq (mvPB pH 0)
    (.seq (entryB pM) mv4B)))))

def cChirpB (φ w Δ : ℕ) : ℕ := 2 * (4 * φ + 100000 * (w + Δ + 1)) + 10000 * (w + 1)

def GBnd (w : ℕ) (z : GI) : Prop := zsz z.1 ≤ w ∧ zsz z.2 ≤ w

/-! ## The loop, the reversal and the clean-up -/

/-- The output after `k` iterations. -/
def outL (Δ : ℕ) (c h g : GI) (k : ℕ) : List GI := (List.range k).map (fun i => (cst Δ c h g i).2)

/-- Drop the two state pairs and `c`. -/
def drop6B : IStmt HS P :=
  .seq (.op (.drop 0)) (.seq (.op (.drop 0)) (.seq (.op (.drop pH)) (.seq (.op (.drop pH))
    (.seq (.op (.drop pK)) (.op (.drop pK))))))

/-- **The table program**: `n` iterations, reversal onto `pW`, clean-up. -/
def chirpG (pM : Fin P) (n : ℕ) : IStmt HS P :=
  .seq (rpt pA pC pD n (chirpB pM)) (.seq (rpt pA pC pD n (mvPB pV pW)) drop6B)

def cChirpG (n φ w Δ : ℕ) : ℕ := n * (8 * φ + 300000 * (w + Δ + 1)) + 10000 * (n + 1) * (w + 1)

/-! ## Base roots by `twiddleP` -/

/-- Push `n, k, j` on `28, 29, 53`, run `twiddleP`, pop. -/
def twRS (n k j : ℕ) : RS :=
  .seq (.push 28 (.lit n)) (.seq (.push 29 (.lit k)) (.seq (.push 53 (.lit j)) (.seq twiddleP
    (.seq (.pop 53) (.seq (.pop 29) (.pop 28))))))

/-- Run `twRS` as IL code and move the root onto `X` (`re` on top). -/
def twC (n k j : ℕ) (X : Fin HS) : IStmt HS P :=
  .seq (RS.code (twRS n k j)) (.seq (.op (.mv 52 X)) (.op (.mv 51 X)))

def cTw (n k j : ℕ) : ℕ :=
  twiddleCost n (zsz (j : ℤ)) (zsz (k : ℤ)) +
    7 * 10 ^ 7 * (zsz (n : ℤ) + zsz (k : ℤ) + zsz (j : ℤ) + 2) ^ 2 + 220 * (n + 4)

/-! ## Size bounds along the recurrence -/

/-! ## The axis-table and twist-table programs -/

/-- **Axis table** `axTab Δ s t` (`≈ e^{-πi s j²/t}`, `j < t`) onto `pW`. -/
def axG (pM : Fin P) (Δ s t : ℕ) : IStmt HS P :=
  .seq (.op (.const 0 0)) (.seq (.op (.const 0 (2 ^ Δ))) (.seq (twC Δ (2 * t) (2 * t - s) pH)
    (.seq (twC Δ t (t - s) pK) (chirpG pM t))))

def cAx (Δ s t φ : ℕ) : ℕ :=
  cTw Δ (2 * t) (2 * t - s) + cTw Δ t (t - s) + cChirpG t φ (Δ + 2) Δ + (Δ + 6)

/-- **Twist table** `twTab Δ k j n` (`≈ e^{2πi j m/k}`, `m < n`) onto `pW`. -/
def twG (pM : Fin P) (Δ k j n : ℕ) : IStmt HS P :=
  .seq (.op (.const 0 0)) (.seq (.op (.const 0 (2 ^ Δ))) (.seq (twC Δ k j pH)
    (.seq (.op (.const pK 0)) (.seq (.op (.const pK (2 ^ Δ))) (chirpG pM n)))))

def cTwG (Δ k j n φ : ℕ) : ℕ := cTw Δ k j + cChirpG n φ (Δ + 2) Δ + (2 * Δ + 12)

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMTabQ
/-!
# KFM 9d: the `d`-axis chirp array — outer products of the axis tables

`chirpArr Δ d s t` is the row-major array (last axis innermost) of the chirp entries
`chirpEnt Δ s t j`, built by `d` outer products `oprod Δ A B = [Q_Δ(a b) | a ∈ A, b ∈ B]`
starting from `[1]`:

* **`chirpArr_get`**: entry `flatIx j` of `chirpArr` is `chirpEnt Δ s t j` (the `Q_Δ`-product of
  the `d` axis-table factors, `KFMTabM`);
* **`chirpArr_err`**: it is within `d (8T² + 2) 2^{-Δ}` of `RC.chirp s t j`;
* **`chirpArr_near`**: every entry is within that distance of a point of the unit disc, so all
  entries have `Δ + 2`-bit components (`NearU.bnd`).

The IL program is `chirpP` (below), one level `levP` per axis: `axG` (axis table), then for each
entry `a` of the current array a replicate / copy / `gmulP` / move-out round (`obB`).
-/

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

namespace NLogN.KFM

open IntMul KFC KFB KFT KIL

/-! ## Outer products -/

/-- `[Q_Δ(a b) | a ∈ A, b ∈ B]`, `B` innermost. -/
def oprod (Δ : ℕ) (A B : List GI) : List GI := A.flatMap (fun a => B.map (gmulT Δ a))

/-! ## The chirp array -/

/-- The row-major chirp array (last axis innermost). -/
def chirpArr (Δ : ℕ) : (d : ℕ) → (Fin d → ℕ) → (Fin d → ℕ) → List GI
  | 0, _, _ => [oneP Δ]
  | d + 1, s, t => oprod Δ (chirpArr Δ d (fun i => s i.castSucc) (fun i => t i.castSucc))
      (axTab Δ (s (Fin.last d)) (t (Fin.last d)))

/-- The row-major flat index. -/
def flatIx : (d : ℕ) → (t : Fin d → ℕ) → ((i : Fin d) → ZMod (t i)) → ℕ
  | 0, _, _ => 0
  | d + 1, t, j => flatIx d (fun i => t i.castSucc) (fun i => j i.castSucc) * t (Fin.last d) +
      (j (Fin.last d)).val

/-! ## Errors and sizes -/

/-- `z` is within `e` of a point of the unit disc. -/
def NearU (Δ : ℕ) (e : ℝ) (z : GI) : Prop := ∃ Z : ℂ, ‖Z‖ ≤ 1 ∧ ‖vP Δ z - Z‖ ≤ e

/-! ## Generic list loops (counters `pA, pC, pD`) -/

section IL

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

/-- Stack `X` is not a counter of the inner `rpt`. -/
abbrev NC (X : Fin HS) : Prop := X ≠ pA ∧ X ≠ pC ∧ X ≠ pD

/-- `dup X Y; mv X Z` per item. -/
def dmB (X Y Z : Fin HS) : IStmt HS P := .seq (.op (.dup X Y)) (.op (.mv X Z))

/-! ## One outer-product round -/

abbrev pQ1 : Fin HS := 40
abbrev pQ2 : Fin HS := 41
abbrev pQ3 : Fin HS := 42

/-- Round for one entry `a` of the array on `pQ1` against the table on `pQ2` (length `b`):
the products `Q_Δ(a b_j)` are pushed (reversed) on `pQ3`. -/
def obB (pM : Fin P) (b : ℕ) : IStmt HS P :=
  .seq (mvPB pQ1 pH) (.seq (rpt pA pC pD b (cpyB pH 0)) (.seq (rpt pA pC pD (2 * b) (.op (.mv pQ2 pU)))
    (.seq (rpt pA pC pD (2 * b) (dmB pU pW pQ2)) (.seq (gmulP pM b)
    (.seq (rpt pA pC pD (2 * b) (.op (.mv 0 pQ3))) (.seq (.op (.drop pH)) (.op (.drop pH))))))))

def cOb (b w Δ φ : ℕ) : ℕ :=
  2000 * (w + 1) + (b * (1000 * (2 * w + 2) + 64) + 60 + cRpt b) +
    2 * ((2 * b) * (110 * (w + 1) + 64) + 60 + cRpt (2 * b)) +
    ((2 * b) * (410 * (w + 1) + 64) + 60 + cRpt (2 * b)) + cGmul b w Δ φ + 4 * (w + 1)

/-! ## The outer-product loop and one level -/

abbrev oA : Fin HS := 37
abbrev oC : Fin HS := 38
abbrev oD : Fin HS := 39

/-- `oprod Δ A B` (reversed) onto `pQ3`, consuming `A` from `pQ1`. -/
def oprodP (pM : Fin P) (n b : ℕ) : IStmt HS P := rpt NLogN.KFM.oA NLogN.KFM.oC NLogN.KFM.oD n (obB pM b)

/-- **One level**: the axis table of `(s, t)`, then `A ↦ oprod Δ A tab` on `pQ1`. -/
def levP (pM : Fin P) (Δ s t n : ℕ) : IStmt HS P :=
  .seq (axG pM Δ s t) (.seq (rpt pA pC pD (2 * t) (.op (.mv pW pU)))
    (.seq (rpt pA pC pD (2 * t) (.op (.mv pU pQ2))) (.seq (oprodP pM n t)
    (.seq (rpt pA pC pD (2 * (n * t)) (.op (.mv pQ3 pQ1)))
      (rpt pA pC pD (2 * t) (.op (.drop pQ2)))))))

def cLev (Δ s t n φ : ℕ) : ℕ :=
  cAx Δ s t φ + 2 * ((2 * t) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * t)) +
    (n * (cOb t (Δ + 2) Δ φ + 64) + 60 + cRpt n) +
    ((2 * (n * t)) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * (n * t))) +
    ((2 * t) * (2 * (Δ + 2) + 2 + 64) + 60 + cRpt (2 * t))

/-! ## All levels -/

/-- **The chirp-array program**: `chirpArr Δ d s t` onto `pQ1`. -/
def chirpP (pM : Fin P) (Δ : ℕ) : (d : ℕ) → (Fin d → ℕ) → (Fin d → ℕ) → IStmt HS P
  | 0, _, _ => .seq (.op (.const pQ1 0)) (.op (.const pQ1 (2 ^ Δ)))
  | d + 1, s, t => .seq (chirpP pM Δ d (fun i => s i.castSucc) (fun i => t i.castSucc))
      (levP pM Δ (s (Fin.last d)) (t (Fin.last d)) (∏ i : Fin d, t i.castSucc))

def cChirpP (Δ φ : ℕ) : (d : ℕ) → (Fin d → ℕ) → (Fin d → ℕ) → ℕ
  | 0, _, _ => Δ + 6
  | d + 1, s, t => cChirpP Δ φ d (fun i => s i.castSucc) (fun i => t i.castSucc) +
      cLev Δ (s (Fin.last d)) (t (Fin.last d)) (∏ i : Fin d, t i.castSucc) φ

/-! ## Moving the array to `pW`, and cost bounds -/

/-- `chirpArr` onto `pW` (for `gmulP`), `pQ1` restored. -/
def chirpTabP (pM : Fin P) (Δ d : ℕ) (s t : Fin d → ℕ) : IStmt HS P :=
  .seq (chirpP pM Δ d s t) (.seq (rpt pA pC pD (2 * ∏ i, t i) (.op (.mv pQ1 pU)))
    (rpt pA pC pD (2 * ∏ i, t i) (.op (.mv pU pW))))

def cChirpTab (Δ φ d : ℕ) (s t : Fin d → ℕ) : ℕ :=
  cChirpP Δ φ d s t + 2 * ((2 * ∏ i, t i) * (110 * (Δ + 2 + 1) + 64) + 60 + cRpt (2 * ∏ i, t i))

end IL

end NLogN.KFM

end

section
-- Source module: Solutions.KFMU1
/-!
# KFM U1: uniform counted repetition and uniform pointwise passes

`UNIFORM.md`: the count of a repetition is run-time data. `rptX X A C D b` is KFT's `rpt` with the
literal count replaced by a copy (`dup`) of the top of stack `X`.

* **`rptX_runs`**: as `rpt_runs`, from a state whose `X` holds `n`; cost `+ dupCost (zsz n) + 17`.
* **`mvBackU_runs`**, **`gmulPU_runs`**, **`itemPU_runs`**: KFM's `mvBack`/`gmulP`/`itemP` with the
  counts on `nX = 30` (`N`) and `nX2 = 31` (`2N`), and the instances **`shlPU_runs`**,
  **`roundPU_runs`**, **`smulPU_runs`**, **`gmulPU_kTable`**. The bodies (`entryB`, `mvB`, `shlB`,
  `roundB`, `smulB`) are KFM's, unchanged (stacks 43–53, 60–62; leaf procedures).

**Nesting.** Every spec here is frame-preserving: the final state is the initial one with only the
output stacks changed, for arbitrary contents of all other stacks (including the scratch 43–53,
60–62, whose contents are restored). So a pass may run inside a caller that keeps data on those
stacks, e.g. KFV's parameter stacks 50–55. The counters avoid KFV6Loop's stacks 20–23.
-/

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

namespace NLogN.KFT

open IntMul KFC

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

/-- The repetition with its count read from the top of `X`. -/
def rptX (X A C D : Fin HS) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (.op (.dup X A)) (.seq (pro A C D) (.seq (uloop C D b) (.seq (epi C D) (.op (.drop A)))))

end NLogN.KFT

namespace NLogN.KFM

open IntMul KFC KFB KFT KIL

/-- Count registers of the uniform passes (KFM range 20–35). -/
abbrev nX : Fin HS := 30
abbrev nX2 : Fin HS := 31

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

/-- Uniform `gmulP`: counts `N` on `nX`, `2N` on `nX2`. -/
def gmulPU (pM : Fin P) : IStmt HS P :=
  .seq (rptX nX pA pC pD (entryB pM)) (rptX nX2 pA pC pD mvB)

/-- Uniform `itemP`: count `N` on `nX`. -/
def itemPU (body : IStmt HS P) : IStmt HS P :=
  .seq (rptX nX pA pC pD body) (rptX nX pA pC pD mvB)

/-- The cost of the counter copies of a uniform pass. -/
def cCnt (N : ℕ) : ℕ := dupCost (zsz (N : ℤ)) + 17

end NLogN.KFM

end