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

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

import Definitions.Def_IntMul_Kappa_K_Defs24

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: Solutions.KFMBig9
/-!
# KFM Big9: the read-out

After the rounding passes (`finA`), `Φ⁻¹; carries; write` outputs `bin(2n, x·y)`.

* `cN_le`, `zsz_cN`: the product coefficients are small.
* **`cconv_ubar`**: the cyclic convolution of the two digit arrays at `Φ(m)` is `c_m`.
* **`pt_val`**: from the final bound (`final_call3`), the rounded numerators at `k` are
  `(c_m, 0)`, `m = Φ⁻¹ k`.
* `zsW_GB`: all rounded numerators have `p + 2` bits.
* **`fin_phase`**: `finA; Φ⁻¹; carries; write`.
-/

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

namespace NLogN.KFM

open Finset AModel AModelL12 IntMul KIL KFB KFT KFU RS IntMul.HvdH IntMul.OpenAI.RT
  IntMul.OpenAI.RC IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-! ## Coefficients -/

variable {E : ℕ} {ix : Ix (7 + E)} {Isym : List Sym} {extM : Fin E → IStmt HS (7 + E)} {n D : ℕ}
  [instS : ∀ i, NeZero (sBD n D i)] [instT : ∀ i, NeZero (tBD n D i)]

/-- The final closeness (the conclusion of `final_call3`). -/
def Close (C : BigCtx ix Isym extM n D) (x y : ℕ) (Z : ((i : Fin (D + 1)) → ZMod (sBD n D i)) → ℂ) :
    Prop :=
  ∀ j, ‖(2 : ℂ) ^ (2 * bP n + 4) * ((∏ i, (sBD n D i : ℝ) : ℝ) : ℂ) *
      (((∏ i, (sBD n D i : ℝ) : ℝ) : ℂ) * cnj Z j) - cconv (sBD n D) (ubar C x) (ubar C y) j‖ <
    1 / 2

/-- The rounded array. -/
def zsWf (n : ℕ) (zsZ : List GI) : List GI :=
  (zsZ.map NLogN.KFM.gconj).map fun z => (rndV (pP n) (bP n) (sCmp n) z.1, rndV (pP n) (bP n) (sCmp n) z.2)

/-! ## Carries -/

/-- The carries with their registers. -/
def carW {P : ℕ} : IStmt HS P := withD rB2 uB2 (withD NLogN.KFM.rS uM carIU)

/-- The read-out program. -/
def finP (ix : Ix (7 + E)) : IStmt HS (7 + E) :=
  .seq finA (.seq (callW (descPhiP (bX ix.dSΦ) rPw) (bX ix.dSΦ) (bY ix.dSΦ) ix.dSΦ (ix.g₁ ix.eI))
    (.seq carW (writeI 0 rO)))

/-- The cost bound of the read-out. -/
def kFin (C : BigCtx ix Isym extM n D) : ℝ :=
  (kFinA n D : ℝ) + kPhi n C.BΦ C.CΦ +
    ((cCarU (bP n) (pP n + 2) (sCmp n) + cW (sCmp n : ℤ) + cW (2 ^ bP n) + outRCost (2 * n) : ℕ) : ℝ)

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMBig10
/-!
# KFM Big10: one run of `bigE`

`bigP = setP; mulP; finP`. From the state after the reads (`readSt x y`) it outputs
`bin(2n, val x · val y)`, given the context `BigCtx` of `n = |x| = |y|`.

* **`big_run`**: the run and its cost (setup cost `setC x y` as computed, plus `kMul`, `kFin`).
-/

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

namespace NLogN.KFM

open Finset AModel AModelL12 IntMul KIL KFB KFT KFU RS IntMul.HvdH IntMul.OpenAI.RT
  IntMul.OpenAI.RC IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

variable {E : ℕ} {ix : Ix (7 + E)} {Isym : List Sym} {extM : Fin E → IStmt HS (7 + E)} {n D : ℕ}
  [instS : ∀ i, NeZero (sBD n D i)] [instT : ∀ i, NeZero (tBD n D i)]

/-- **The program `bigE`.** -/
def bigP (ix : Ix (7 + E)) : IStmt HS (7 + E) := .seq setP (.seq (mulP ix) (finP ix))

/-- The setup cost as computed. -/
def setC (x y : List Bool) : ℕ :=
  cStart x.length + 12 * 10 ^ 15 * (bP x.length + 1) ^ 4 +
    parS.cost (setupSt x.length (stA x y)) +
    parQ.cost ((setupSt x.length (stA x y)).upd NLogN.KFM.rS (.int (sCmp x.length : ℤ) ::
      (setupSt x.length (stA x y)).G NLogN.KFM.rS)) + (3 + 2 * shlCost x.length)

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMBig11
/-!
# KFM Big11: the merged table and `BigSpec`

* `phi_cond`: the parameters of `n` meet `PhiNeed`'s side conditions.
* `dzB`: the canonical size bound of a descriptor.
* **`ctx_exists`**: from the three uniform tables (`PhiNeed`, `LineNeed` via `lineNeed_D`,
  `ChirpNeedR`) and `KFQ.chainSizes`, a `BigCtx` exists for every large `n` and every `Isym`, with
  the fixed constants and the canonical descriptor bounds.
* `kMul_kFin`: the cost bound of a context depends only on those.
* **`bigSpec_of`**: contexts for all `n ≥ N₁` give `BigSpec` for `bigE` with `own = bigP ix`.
-/

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

namespace NLogN.KFM

open Finset AModel AModelL12 IntMul KIL KFB KFT KFU RS IntMul.HvdH IntMul.OpenAI.RT
  IntMul.OpenAI.RC IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-! ## Side conditions of `PhiNeed` -/

/-! ## Canonical descriptor bounds -/

/-- The largest size of an entry. -/
def dzB (l : List ℤ) : ℕ := (l.map zsz).foldr max 0

/-! ## The merged table -/

section Merge

variable {E₁ E₂ E₃ : ℕ}

/-- The index data of the merged table (`E₀ = 1`, `own 0 = bigP`). -/
def ixM (E₁ E₂ E₃ : ℕ) (eF eI : Fin (7 + E₁)) (eA eB : Fin (7 + E₂)) (eC : Fin (7 + E₃))
    (dSΦ dSL dSC : Fin HS) : Ix (7 + (E₁ + E₂ + E₃ + 1)) where
  P₁ := E₁
  P₂ := E₂
  P₃ := E₃
  g₁ := NLogN.KFM.gI E₁ _ 0 (by omega)
  g₂ := NLogN.KFM.gI E₂ _ E₁ (by omega)
  g₃ := NLogN.KFM.gI E₃ _ (E₁ + E₂) (by omega)
  eF := eF
  eI := eI
  eA := eA
  eB := eB
  eC := eC
  dSΦ := dSΦ
  dSL := dSL
  dSC := dSC

end Merge

/-! ## The cost of a context -/

/-- `kIn` / `kFt` from the constants. -/
def kFtF (n D BL BC AL AC : ℕ) (CL CC : ℝ) : ℝ :=
  kLine n BL CL AL + kChirp n BC CC AC + kLine n BL CL AL + kShl n D

def kInF (n D BΦ BL BC AL AC : ℕ) (CΦ CL CC : ℝ) : ℝ :=
  ((2 * moveCost (n + 1) + 19 + kDig n : ℕ) : ℝ) + kPhi n BΦ CΦ + kFtF n D BL BC AL AC CL CC

/-- `kMul + kFin` from the constants. -/
def kBigF (n D BΦ BL BC AL AC : ℕ) (CΦ CL CC : ℝ) : ℝ :=
  (kInF n D BΦ BL BC AL AC CΦ CL CC + ((2 * cMvN (2 ^ (IntMul.OpenAI.Glue.LP n + 1)) (NLogN.KFM.wU n D) : ℕ) : ℝ) +
    kInF n D BΦ BL BC AL AC CΦ CL CC + (kGm n D : ℝ) + (kCj n (pP n + 1) : ℝ) +
    kFtF n D BL BC AL AC CL CC) +
  ((kFinA n D : ℝ) + kPhi n BΦ CΦ +
    ((cCarU (bP n) (pP n + 2) (sCmp n) + cW (sCmp n : ℤ) + cW (2 ^ bP n) + outRCost (2 * n) : ℕ) : ℝ))

/-! ## `BigSpec` -/

/-! ## Contexts -/

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMSetC
/-!
# KFM SetC: the cost of the parameter pushes

* `pushSeq`, **`pushSeq_cost`**: a chain of pushes costs `Σ (size+1) · 10^7 (B+2)^2 + 1` when every
  expression is `B`-bounded at its state (`PBnd`).
* **`parQ_cost`**: `parQ` costs `O(b^2)`.
* **`parS_cost`**: `parS` costs `O(d · b^2)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFM

open IntMul IntMul.OpenAI.Glue RS

/-! ## Chains of pushes -/

/-- A chain of pushes. -/
def pushSeq : List (Fin HS × Ex HS) → RS
  | [] => .skip
  | [(r, e)] => .push r e
  | (r, e) :: l => .seq (.push r e) (pushSeq l)

/-- Every expression is `B`-bounded at its state. -/
def PBnd (B : ℕ) : List (Fin HS × Ex HS) → IState HS → Prop
  | [], _ => True
  | (r, e) :: l, t => e.Bnd t B ∧ PBnd B l (t.upd r (.int (e.eval t) :: t.G r))

/-! ## `parQ` -/

/-- The pushes of `parQ`. -/
def Lq : List (Fin HS × Ex HS) :=
  [(NLogN.KFM.rT, .pow2 (.reg NLogN.KFM.zL)), (rT2, .pow2 (.add (.reg NLogN.KFM.zL) (.lit 1))),
   (NLogN.KFM.rM, .ndiv (.add (.reg NLogN.KFM.zN) (.add (.reg NLogN.KFM.zB) (.lit (-1)))) (.reg NLogN.KFM.zB)),
   (rP2, .pow2 (.mul (.lit 6) (.reg NLogN.KFM.zB))),
   (rG, .pow2 (.mul (.mul (.lit 2) (.reg zD)) (.mul (.reg zAl) (.reg zAl)))),
   (rE2, .pow2 (.add (.mul (.lit 5) (.reg NLogN.KFM.zB)) (.lit (-2)))),
   (rB2, .pow2 (.reg NLogN.KFM.zB)),
   (rZ, .mul (.lit 2) (.add (.reg NLogN.KFM.rS) (.neg (.reg NLogN.KFM.rM)))),
   (NLogN.KFM.rH, .pow2 (.add (.mul (.lit 6) (.reg NLogN.KFM.zB)) (.lit (-1)))),
   (rF, .pow2 (.add (.mul (.lit 2) (.reg NLogN.KFM.zB)) (.lit 4))),
   (rPw, .add (.mul (.lit 6) (.reg NLogN.KFM.zB)) (.lit 2)),
   (rOne, .lit 1)]

/-! ## `parS` -/

/-! ## The whole setup -/

end NLogN.KFM

end

section
-- Source module: Solutions.KFZFmb
/-!
# KFZ Fmb: `fmB` is `O(m log² m)`

**`fmB_le_log`**: for `1 ≤ p` and `m ≤ 64 p`,
`KFM.fmB m ≤ 32500 (8 fmA + 10^7 + 10^5) · p · (log p + 1)^2`.
-/

set_option autoImplicit false

namespace NLogN.KFZ

open IntMul NLogN.KFM

/-- The constant of `fmB`. -/
def Cfm : ℕ := 8 * fmA + 10 ^ 7 + 10 ^ 5

attribute [local irreducible] fmA Cfm

end NLogN.KFZ

end

section
-- Source module: Solutions.KFMCost
/-!
# KFM Cost: the cost of `bigE` in the unit `ZU`

`ZU n = T (p + 1 + fmB (8p)) + (b + 1)^6 + n`. Every cost of `bigE` other than the three
interfaces is `O(ZU n)`:

* `wU_le`, `wFR_le`, `sCmp_le_T`, `mN_le`: sizes;
* `kShl_le`, `cMvN_le'`, `kGm_le`, `kCj_le`, `kFinA_le`, `cCarU_le'`, `kDig_le`: the passes;
* `desc_le`: the descriptor parts of the calls;
* `setC_le`: the setup.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFM

open Finset IntMul KIL KFB KFT KFU RS IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-- The unit. -/
def ZU (n : ℕ) : ℕ := IntMul.OpenAI.Glue.TP n * (pP n + 1 + fmB (8 * pP n)) + (bP n + 1) ^ 6 + n

/-! ## Basic facts (as in `KFZCostR`, without its imports) -/

section Basic

variable {n : ℕ} (hn : n0 ≤ n)
include hn

end Basic

/-! ## Primitive costs -/

/-! ## Sizes -/

section Sizes

variable {n D : ℕ} (hn : n0 ≤ n) (hD : dP n = D + 1)
include hn hD

end Sizes

/-! ## The passes -/

section Passes

variable {n D : ℕ} (hn : n0 ≤ n) (hD : dP n = D + 1)
include hn hD

end Passes

/-! ## Descriptors -/

section Desc

variable {n : ℕ} (hn : n0 ≤ n)
include hn

end Desc

/-! ## Large `d`, small `r` -/

/-! ## `Φ` -/

/-- The `Ncost` of `n`. -/
def NcN (n : ℕ) : ℝ := Ncost (6 * (lg n : ℝ)) (IntMul.OpenAI.Glue.dP n) (KP n) (ℓP n) (αP n)

/-- The log factor of `n`. -/
def LgN (n : ℕ) : ℝ := Real.log (6 * (lg n : ℝ)) + 1

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFMFinal
/-!
# KFM Final: `PhiNeed ∧ LineNeed ∧ ChirpNeedR → KappaBound (1/2^182)`

* `ZU_le`, `TpN_le`: the unit `ZU` and `T p Ncost` are within `cLineU n A` (`A ≥ 2`).
* **`total_le`**: setup plus `kMul + kFin` of a context is at most `Ctot · cLineU n A'`,
  `A' = AL + AC + 2`.
* **`kappa_final`**: the target.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFM

open Finset IntMul KIL KFB KFT KFU RS IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-! ## The unit against the budget -/

section Unit

variable {n D : ℕ} (hn : n0 ≤ n) (hD : dP n = D + 1)
include hn hD

/-- The constant of the unit. -/
def Kz : ℝ := 32500 * (KFZ.Cfm : ℝ) + 200

end Unit

/-! ## The total cost -/

/-- The constant of the total. -/
def Ctot (CΦ CL CC : ℝ) : ℝ := (10 ^ 18 + 10 ^ 5 * CΦ) * Kz + 4 * 10 ^ 9 * CΦ + 6 * CL + 3 * CC

/-! ## The target -/

end

end NLogN.KFM

end

section
-- Source module: Solutions.KFZDesc
/-!
# KFZ Desc: reading the chirp descriptor (register code)

The chirp descriptor on `cS` is `[K, p, α, h, d, b_0..b_{d−1}, s_0..s_{d−1}]` (`h = ℓ − 1`,
`t_i = 2^{h + b_i}`). **`setupP`** (one register program, leaf registers `4, 8–13`) reads it,
leaves `K, p, h, d, Σ b_i` on the registers `rgK, rgP, rgH, rgD, rgS` (scratch stacks
`55, 50, 51, 53, 54`), builds the chirp-table descriptor `descA 1 L ++ [2 ∏ t_i]` (`KFMU2.descT`)
on stack `35`, and restores the chirp descriptor (**`setupP_sp`**).
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL NLogN.KFM RS

abbrev cS : Fin HS := 36
abbrev rgP : Fin HS := 50
abbrev rgH : Fin HS := 51
abbrev rgD : Fin HS := 53
abbrev rgS : Fin HS := 54
abbrev rgK : Fin HS := 55
abbrev lfA : Fin HS := 4
abbrev lfB : Fin HS := 8
abbrev lfS : Fin HS := 9
abbrev lfB2 : Fin HS := 10
abbrev lfC : Fin HS := 11
abbrev lfS2 : Fin HS := 12
abbrev lfE : Fin HS := 13
abbrev tD : Fin HS := 35

/-- Integer copy of a list of naturals. -/
def zl (l : List ℕ) : List ℤ := l.map fun x : ℕ => (x : ℤ)

@[simp] theorem zl_nil : NLogN.KFZ.zl [] = [] := rfl
@[simp] theorem zl_cons (a : ℕ) (l : List ℕ) : NLogN.KFZ.zl (a :: l) = (a : ℤ) :: NLogN.KFZ.zl l := rfl

/-! ## One move -/

def mv1 (Y X : Fin HS) : RS := .seq (.push Y (.reg X)) (.pop X)

/-! ## Prefix exponents -/

/-- `Σ_{i<m} (h + b_i)`. -/
def eS (h : ℕ) (bs : List ℕ) (m : ℕ) : ℕ := m * h + (bs.take m).sum

/-- The axis records `(s_i, 2^{h + b_i})`. -/
def axD (h : ℕ) (bs ss : List ℕ) : List (ℕ × ℕ) := List.zipWith (fun b s => (s, 2 ^ (h + b))) bs ss

theorem length_axD (h : ℕ) (bs ss : List ℕ) (hl : ss.length = bs.length) :
    (axD h bs ss).length = bs.length := by simp +decide [axD, hl]

/-! ## Summing the bits -/

def sumB : RS := .seq (.asg rgS (.add (.reg rgS) (.reg NLogN.KFZ.cS))) (NLogN.KFZ.mv1 lfB NLogN.KFZ.cS)

/-! ## Expression helpers on managed states -/

/-! ## One axis record of the chirp-table descriptor -/

/-- `t = 2^{h + b}`. -/
def tE : Ex HS := .pow2 (.add (.reg rgH) (.reg lfB))
/-- `E − (h + b)`. -/
def nE : Ex HS := .add (.reg lfE) (.neg (.add (.reg rgH) (.reg lfB)))

/-- Push `tupL s t n` (`n = 2^{E−h−b}`, `2nt = 2·2^E`) on `tD`, lower `E`, advance the bit and
source lists. -/
def dtB : RS :=
  .seq (.push tD (.mul (.lit 2) (.pow2 (.reg lfE))))
  (.seq (.push tD (.pow2 NLogN.KFZ.nE))
  (.seq (.push tD (.add NLogN.KFZ.tE (.neg (.reg lfS))))
  (.seq (.push tD (.add (.mul (.lit 2) NLogN.KFZ.tE) (.neg (.reg lfS))))
  (.seq (.push tD (.mul (.lit 2) NLogN.KFZ.tE))
  (.seq (.push tD NLogN.KFZ.tE)
  (.seq (.asg lfE NLogN.KFZ.nE)
  (.seq (NLogN.KFZ.mv1 lfB2 lfB) (NLogN.KFZ.mv1 lfS2 lfS))))))))

/-! ## The descriptor loop (last axis first) -/

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZSetup
/-!
# KFZ Setup: the setup program `setupP`

`setupP` (one register program) reads the chirp descriptor on `cS`, leaves `K, p, h, d, Σ b_i`
on the registers `rgK, rgP, rgH, rgD, rgS`, builds the chirp-table descriptor on `tD`, and
restores the chirp descriptor (**`setupP_sp`**).
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL NLogN.KFM RS

/-! ## The setup program -/

def hdrIn : RS :=
  .seq (NLogN.KFZ.mv1 rgK NLogN.KFZ.cS) (.seq (NLogN.KFZ.mv1 rgP NLogN.KFZ.cS) (.seq (NLogN.KFZ.mv1 lfA NLogN.KFZ.cS) (.seq (NLogN.KFZ.mv1 rgH NLogN.KFZ.cS) (NLogN.KFZ.mv1 rgD NLogN.KFZ.cS))))

def hdrOut : RS :=
  .seq (.push NLogN.KFZ.cS (.reg rgD)) (.seq (.push NLogN.KFZ.cS (.reg rgH)) (.seq (NLogN.KFZ.mv1 NLogN.KFZ.cS lfA)
    (.seq (.push NLogN.KFZ.cS (.reg rgP)) (.push NLogN.KFZ.cS (.reg rgK)))))

/-- **Read the chirp descriptor** and build the chirp-table descriptor on `tD`. -/
def setupP : RS :=
  .seq hdrIn (.seq (.push rgS (.lit 0)) (.seq (.rep lfC (.reg rgD) NLogN.KFZ.sumB)
  (.seq (NLogN.mvK NLogN.KFZ.cS lfS lfC (.reg rgD))
  (.seq (.push lfE (.add (.mul (.reg rgD) (.reg rgH)) (.reg rgS)))
  (.seq (.push tD (.mul (.lit 2) (.pow2 (.reg lfE))))
  (.seq (.rep lfC (.reg rgD) dtB) (.seq (.pop lfE)
  (.seq (NLogN.mvK lfS2 lfS lfC (.reg rgD)) (.seq (NLogN.mvK lfS NLogN.KFZ.cS lfC (.reg rgD))
  (.seq (NLogN.mvK lfB2 lfB lfC (.reg rgD)) (.seq (NLogN.mvK lfB NLogN.KFZ.cS lfC (.reg rgD)) hdrOut)))))))))))

/-- The leaf and register stacks of `setupP`. -/
def setupRegs : List (Fin HS) := [rgP, rgH, rgD, rgS, rgK, lfA, lfB, lfS, lfB2, lfC, lfS2, lfE]

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZSet
/-!
# KFZ Set: parameter pushes for the table and pass procedures

`pushes l` pushes the values of the expressions `l` (read from the registers `rgP, rgH, rgD,
rgS, rgK`) on the parameter stacks of the KFM procedures; `popsL` removes them again.

* **`pushes_sp`**, **`popsL_sp`**: the generic specs.
* The parameter expressions `eL` (`L = d h + Σb`), `eN` (`N = 2^L`), `eDl` (`Δ = 2p + 2h + d + 8`),
  `eG` (`g = Δ − p`), `eR` (`r = 2^{h+1}`), … with their values (**`Regs.ev_*`**).
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL NLogN.KFM RS

/-! ## Generic pushes and pops -/

def pushes : List (Fin HS × Ex HS) → RS
  | [] => .skip
  | (i, e) :: l => .seq (.push i e) (pushes l)

def popsL : List (Fin HS) → RS
  | [] => .skip
  | i :: l => .seq (popsL l) (.pop i)

def pushF (F : Fin HS → List ℤ) : List (Fin HS × ℤ) → Fin HS → List ℤ
  | [] => F
  | (i, v) :: l => pushF (Function.update F i (v :: F i)) l

/-! ## The registers -/

abbrev xP : Ex HS := .reg rgP
abbrev xH : Ex HS := .reg rgH
abbrev xD : Ex HS := .reg rgD
abbrev xS : Ex HS := .reg rgS
abbrev xK : Ex HS := .reg rgK

/-- The registers hold `p, h, d, σ = Σ b, K`. -/
structure Regs (F : Fin HS → List ℤ) (p h d σ K : ℕ) : Prop where
  P : F rgP = [(p : ℤ)]
  H : F rgH = [(h : ℤ)]
  D : F rgD = [(d : ℤ)]
  S : F rgS = [(σ : ℤ)]
  K : F rgK = [(K : ℤ)]

def regL : List (Fin HS) := [rgP, rgH, rgD, rgS, rgK]

/-! ## The parameter expressions -/

/-- `L = d h + σ`. -/
def eL : Ex HS := .add (.mul NLogN.KFZ.xD xH) NLogN.KFZ.xS
/-- `N = 2^L`. -/
def eN : Ex HS := .pow2 NLogN.KFZ.eL
/-- `2N`. -/
def e2N : Ex HS := .mul (.lit 2) NLogN.KFZ.eN
/-- `Δ = 2p + 2h + d + 8`. -/
def eDl : Ex HS := .add (.add (.mul (.lit 2) xP) (.mul (.lit 2) xH)) (.add NLogN.KFZ.xD (.lit 8))
/-- `g = Δ − p`. -/
def eG : Ex HS := .add (.add xP (.mul (.lit 2) xH)) (.add NLogN.KFZ.xD (.lit 8))
/-- `r = 2^{h+1}`. -/
def eR : Ex HS := .pow2 (.add xH (.lit 1))
/-- `2r`. -/
def e2R : Ex HS := .pow2 (.add xH (.lit 2))
/-- `N / r`. -/
def eNo : Ex HS := .pow2 (.add NLogN.KFZ.eL (.neg (.add xH (.lit 1))))
/-- `2^p`. -/
def eP2 : Ex HS := .pow2 xP
/-- `2^Δ`. -/
def eDl2 : Ex HS := .pow2 eDl
/-- `2^g`. -/
def eG2 : Ex HS := .pow2 eG
/-- `d' = d − σ`. -/
def eDp : Ex HS := .add NLogN.KFZ.xD (.neg NLogN.KFZ.xS)
/-- `ℓ = h + 1`. -/
def eEl : Ex HS := .add xH (.lit 1)
/-- `b = p / 6`. -/
def eB : Ex HS := .ndiv xP (.lit 6)

/-- The value of `Δ`. -/
def dlV (p h d : ℕ) : ℕ := 2 * p + 2 * h + d + 8
/-- The value of `g`. -/
def gV (p h d : ℕ) : ℕ := p + 2 * h + d + 8

section Ev

variable {s : IState HS} {F : Fin HS → List ℤ} {p h d σ K B : ℕ}

/-- Bounds on the register values: everything below `B`. -/
structure RB (p h d σ K B : ℕ) : Prop where
  σd : σ ≤ d
  hL : h + 1 ≤ d * h + σ
  big : dlV p h d + d * h + σ + 8 ≤ B
  K : K + 1 ≤ B

variable (hR : Regs F p h d σ K) (hB : RB p h d σ K B)
include hR hB

end Ev

/-! ## The parameter lists -/

/-- Chirp table: `2^Δ` on `pR`, `Δ` on `rDl`, `d` on `aX`. -/
def setT1 : List (Fin HS × Ex HS) := [(NLogN.KFM.pR, eDl2), (rDl, eDl), (aX, NLogN.KFZ.xD)]
/-- Rounding of the chirp table (`2N` integers): counts, `h = 0`, `2^g`. -/
def setRC : List (Fin HS × Ex HS) := [(nX2, e2N), (nX, e2N), (pX2, .lit 0), (NLogN.KFM.pR, eG2)]
/-- Conjugate copies and drops of the arrays: `N` on `nX`, `2N` on `nX2`. -/
def setCN : List (Fin HS × Ex HS) := [(nX, NLogN.KFZ.eN), (nX2, e2N)]
/-- Twist table: `2^Δ`, `Δ`, `k = 2r`, `j = 1`, `n = r`. -/
def setTw : List (Fin HS × Ex HS) := [(NLogN.KFM.pR, eDl2), (rDl, eDl), (KFM.rJ1, e2R), (rJ2, .lit 1), (nX, eR)]
/-- Rounding of the twist table (`2r` integers). -/
def setRT : List (Fin HS × Ex HS) := [(nX2, e2R), (nX, e2R), (pX2, .lit 0), (NLogN.KFM.pR, eG2)]
/-- Replication of the twist table: `N / r` copies of `2r` integers. -/
def setRep : List (Fin HS × Ex HS) := [(KFM.rN, eNo), (rN2, e2R)]
/-- The products: `2^p`, `N`, `2N`, one copy of `2N` integers. -/
def setC : List (Fin HS × Ex HS) := [(NLogN.KFM.pR, eP2), (nX, NLogN.KFZ.eN), (nX2, e2N), (KFM.rN, .lit 1), (rN2, e2N)]
/-- The Prop 18 descriptor registers. -/
def set18 : List (Fin HS × Ex HS) :=
  [(NLogN.KFM.zD, NLogN.KFZ.xD), (zDp, eDp), (zR, eR), (zK, NLogN.KFZ.xK), (zEl, eEl), (NLogN.KFM.zB, NLogN.KFZ.eB)]

section Setter

variable {s : IState HS} {F : Fin HS → List ℤ} {B : ℕ}

variable {p h d σ K : ℕ} (hR : Regs F p h d σ K) (hB : RB p h d σ K B)
include hR hB

end Setter

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZBox
/-!
# KFZ Box: row-major Gaussian lists as functions on the box

The arrays of the chirp step are row-major lists of Gaussian numerators (`RepFull`, last axis
fastest). This file is the dictionary between such lists and functions on the box
`∏ ZMod t_i`:

* `digL_append`, `digL_append_len`, `boxIx_castSucc`, `boxIx_last`: digits of `x` over
  `t_0 … t_D` are the digits of `x / t_D` over `t_0 … t_{D−1}`, then `x mod t_D`.
* `ptB t x`: the box point of index `x`; **`flatIx_ptB`**, **`ptB_flatIx`** (`flatIx` of
  `KFMTabQ` is its inverse).
* `decB p t zs K = toC zs[flatIx K] / 2^p`; **`repFull_decB`**, **`RepFull.eq_decB`**.
* **`decB_gmulT`** (the `gmulP` pass is `Q_p` of the pointwise product), **`decB_gconj`**,
  **`decB_gtr`** (rounding), **`decB_repl`** (a record table repeated over the records),
  **`decB_chirpArr`**.
-/

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

namespace NLogN.KFZ

open IntMul KIL KFB NLogN.KFM Complex

noncomputable section

/-! ## Row-major digits -/

variable {d : ℕ}

/-! ## Box points of indices -/

/-- The box point of the row-major index `x`. -/
def ptB (t : Fin d → ℕ) (x : ℕ) : (i : Fin d) → ZMod (t i) :=
  fun i => ((boxIx t x i : ℕ) : ZMod (t i))

/-! ## Lists as box functions -/

/-- The box function held by a row-major list at scale `2^p`. -/
def decB (p : ℕ) (t : Fin d → ℕ) (zs : List GI) : ((i : Fin d) → ZMod (t i)) → ℂ :=
  fun K => toC (zs.getD (flatIx d t K) 0) / 2 ^ p

/-! ## Pointwise operations -/

/-- Gaussian conjugation. -/
def gconj (z : GI) : GI := (z.1, -z.2)

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZPass
/-!
# KFZ Pass: list moves, conjugate copies and record replication (uniform counts)

All counts are run-time data (`rptX`, the count copied from a register). Every spec is
frame-preserving: only the named stacks change.

* **`mvX_runs`**: move `L` from `X` to `Y` (reversed); **`mv2_runs`**: move through `pU`
  (order kept); **`dropX_runs`**: drop `L`.
* **`cnjCp_runs`**: push the conjugate of the Gaussian list on `H` onto `W` (order kept), `H`
  unchanged; counts `2N` on `nX2`, `N` on `nX`.
* **`replP_runs`**: push `N` copies of the Gaussian list on `H` onto `W`; counts `N` on `rN`,
  `2|τ|` on `rN2`.
-/

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

namespace NLogN.KFZ

open IntMul KIL KFB KFC KFT NLogN.KFM

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

/-! ## Moves -/

/-- Move `n` items from `X` to `Y`, `n` on `CN`. -/
def mvX (X Y CN : Fin HS) : IStmt HS P := rptX CN NLogN.KFM.pA pC pD (.op (.mv X Y))

/-- Cost of moving `n` items of `B` bits. -/
def cMvX (n B : ℕ) : ℕ := n * (110 * (B + 1) + 64) + 60 + cRpt n + cCnt n

/-- Move through `pU`: the order is kept. -/
def mv2 (X Y CN : Fin HS) : IStmt HS P := .seq (mvX X pU CN) (mvX pU Y CN)

/-- Drop `n` items of `X`, `n` on `CN`. -/
def dropX (X CN : Fin HS) : IStmt HS P := rptX CN NLogN.KFM.pA pC pD (.op (.drop X))

/-! ## Conjugate copy -/

/-- One Gaussian of the reversed list on `pU`: push its conjugate on `W`, restore it on `H`. -/
def cnjB (H W : Fin HS) : IStmt HS P :=
  .seq (.op (.neg pU)) (.seq (.op (.dup pU W)) (.seq (.op (.neg pU)) (.seq (.op (.mv pU H))
    (.seq (.op (.dup pU W)) (.op (.mv pU H))))))

/-- Push the conjugate of the list on `H` onto `W`. -/
def cnjCp (H W : Fin HS) : IStmt HS P :=
  .seq (rptX nX2 NLogN.KFM.pA pC pD (.op (.mv H pU))) (rptX nX NLogN.KFM.pA pC pD (cnjB H W))

/-- Cost of a conjugate copy of `N` entries of `B` bits. -/
def cCnj (N B : ℕ) : ℕ :=
  (2 * N) * (110 * (B + 1) + 64) + 60 + cRpt (2 * N) + cCnt (2 * N) +
    (N * (4 * (300 * (2 * B + 1)) + 40 + 64) + 60 + cRpt N + cCnt N)

/-! ## Record replication -/

/-- Push `N` copies (`N` on `rN`) of the list on `H` (length `2|τ|` on `rN2`) onto `W`. -/
def replP (H W : Fin HS) : IStmt HS P :=
  rptX KFM.rN oA oC NLogN.KFM.oD (.seq (rptX rN2 NLogN.KFM.pA pC pD (.op (.mv H pU))) (rptX rN2 NLogN.KFM.pA pC pD (dmB pU W H)))

/-- Cost of `N` copies of `M` items of `B` bits. -/
def cRepl (N M B : ℕ) : ℕ :=
  N * (M * (110 * (B + 1) + 64) + 60 + cRpt M + cCnt M +
    (M * (410 * (B + 1) + 64) + 60 + cRpt M + cCnt M) + 64) + 60 + cRpt N + cCnt N

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZProc
/-!
# KFZ Proc: the table phases of the chirp procedure

The managed contents during the chirp procedure are described by **`St`**: the five registers
`p, h, d, σ, K`, the chirp-table descriptor `td` (stack `35`), stack `0` (`z`), the output stack
`pW` (`w`), and the four tables `tb, ta, tt, ti` on the stacks `56, 57, 59, 63`; everything
else is as in `F0`.

Each phase pushes its parameters (`pushes l`), runs one table or pass procedure, and pops the
parameters again (`wP`). The lemmas `T1_runs … T7_runs` give the exact final `St`.
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM

/-! ## The table stacks and the managed contents -/

abbrev qB : Fin HS := 56
abbrev qA : Fin HS := 57
abbrev qT : Fin HS := 59
abbrev qI : Fin HS := 63

/-- The managed contents during the chirp procedure. -/
def St (F0 : Fin HS → List ℤ) (p h d σ K : ℕ) (td z w tb ta tt ti : List ℤ) : Fin HS → List ℤ :=
  fun j => if j = 0 then z else if j = pW then w else if j = tD then td else if j = NLogN.KFZ.qB then tb
    else if j = qA then ta else if j = qT then tt else if j = qI then ti
    else if j = rgP then [(p : ℤ)] else if j = rgH then [(h : ℤ)] else if j = rgD then [(d : ℤ)]
    else if j = rgS then [(σ : ℤ)] else if j = rgK then [(K : ℤ)] else F0 j

section St

variable (F0 : Fin HS → List ℤ) (p h d σ K : ℕ) (td z w tb ta tt ti : List ℤ)

end St

/-! ## Phases with pushed parameters -/

/-- Push the parameters `l`, run `inner`, pop the parameters. -/
def wP {P : ℕ} (l : List (Fin HS × Ex HS)) (inner : IStmt HS P) : IStmt HS P :=
  .seq (pushes l).code (.seq inner (NLogN.KFZ.popsL (l.map Prod.fst)).code)

/-! ## The table phases -/

/-- Cost of pushing `n` parameters (`B`-bit registers). -/
def cPush (n B : ℕ) : ℕ := n * (31 * 10 ^ 7 * (B + 2) ^ 2) + 1
/-- Cost of popping `n` parameters. -/
def cPop (n B : ℕ) : ℕ := n * (2 * B + 2) + 1

/-- Chirp table. -/
def T1 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setT1 (chirpTabPU (pMul E))
/-- Round the chirp table (`2^g`) and move it to `qB`. -/
def T2 (E : ℕ) : IStmt HS (7 + E) :=
  NLogN.KFZ.wP setRC (.seq (NLogN.KFZ.mv2 pW 0 nX2) (.seq (itemPU roundB) (NLogN.KFZ.mv2 0 NLogN.KFZ.qB nX2)))

/-- Conjugate copy of the table `qB` to `qA`. -/
def T3 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setCN (cnjCp NLogN.KFZ.qB qA)
/-- Twist table `ζ^m`, `m < r`. -/
def T4 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setTw (twGU (pMul E))
/-- Round the twist table onto stack `0`. -/
def T5 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setRT (.seq (NLogN.KFZ.mv2 pW 0 nX2) (itemPU roundB))
/-- Replicate the rounded twist table to `qT` (`N / r` copies) and drop it from stack `0`. -/
def T6 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setRep (.seq (replP 0 qT) (dropX 0 rN2))
/-- Conjugate copy of the table `qT` to `qI`. -/
def T7 (E : ℕ) : IStmt HS (7 + E) := NLogN.KFZ.wP setCN (cnjCp qT qI)

section Tab

variable {E : ℕ} (ext : Fin E → IStmt HS (7 + E)) {I : List Sym} (s : IState HS)
  (F0 : Fin HS → List ℤ) {p h d σ K B : ℕ} (hB : RB p h d σ K B)
include hB

end Tab

/-! ## The pointwise products -/

/-- Copy the table on `H` to `pW` (one copy) and multiply it into stack `0`. -/
def Cmul (E : ℕ) (H : Fin HS) : IStmt HS (7 + E) := .seq (replP H pW) (gmulPU (pMul E))

/-- Drop the four tables (`2N` integers each). -/
def Zdrops (E : ℕ) : IStmt HS (7 + E) :=
  .seq (dropX NLogN.KFZ.qB nX2) (.seq (dropX qA nX2) (.seq (dropX qT nX2) (dropX qI nX2)))

/-- Cost of `Cmul`. -/
def cCmul (w q n : ℕ) : ℕ :=
  cRepl 1 (2 * n) w + (n * (4 * KFM.fmB (2 * w)) + 10 ^ 6 * (n + 1) * (w + q + 1) +
    2 * cCnt (2 * n))

section Prod

variable {E : ℕ} (ext : Fin E → IStmt HS (7 + E)) {I : List Sym} (s : IState HS)

end Prod

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZCall
/-!
# KFZ Call: the Prop 18 call inside the chirp procedure

**`Mcall j d18`**: mark the free leaf stack `X`, build the Prop 18 descriptor on `X` at run time
(`desc18P`, registers from `p, h, d, σ, K`), move it to the descriptor stack `d18` of the Prop 18
procedure over a mark (`xfer`), call `j`, and remove the descriptor again (`unxfer`).

**`M_runs`**: the call runs on the state `encK (bgS s G d18 l) (fC …)`: stack `0` holds exactly the
two record arrays, `d18` exactly the descriptor, every other managed stack is empty, and all other
contents (our tables, registers, the rest of stack `0`, the old contents of `d18`) are in the
unmanaged part `bgS`.
-/

set_option autoImplicit false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue

/-- The leaf stack for building the descriptor. -/
def X18 (d18 : Fin HS) : Fin HS := if d18 = 15 then 16 else 15
/-- The leaf stack for `xfer`. -/
def Y18 (d18 : Fin HS) : Fin HS := if d18 = 17 then 18 else 17

/-- Registers, descriptor, pops. -/
def rs18 (X : Fin HS) : RS := .seq (pushes set18) (.seq (desc18P X) (NLogN.KFZ.popsL (set18.map Prod.fst)))

/-- The Prop 18 call. -/
def Mcall {E : ℕ} (j : Fin (7 + E)) (d18 : Fin HS) : IStmt HS (7 + E) :=
  .seq (.op (.mark (X18 d18))) (.seq (rs18 (X18 d18)).code
    (.seq (xfer (X18 d18) (Y18 d18) d18) (.seq (.call j) (unxfer d18))))

/-- The unmanaged part during the call. -/
def bgS (s : IState HS) (G : Fin HS → List ℤ) (d18 : Fin HS) (l : List ℤ) : IState HS :=
  { s with G := fun i => if i = 0 then ints l ++ s.G 0
      else if i = d18 then .mark :: (ints (G d18) ++ s.G d18) else ints (G i) ++ s.G i }

/-- The managed part during the call. -/
def fC (d18 : Fin HS) (Z D : List ℤ) : Fin HS → List ℤ :=
  Function.update (Function.update (fun _ => []) 0 Z) d18 D

/-- Cost of the parts of `Mcall` around the call. -/
def cM (B Bd b len : ℕ) : ℕ :=
  1 + (6 * (31 * 10 ^ 7 * (B + 2) ^ 2) + 1 + (10 ^ 12 * (b + 1) ^ 3 + (6 * (2 * B + 2) + 1))) +
    (2 * len * KFM.cMv Bd + 70) + (len * KFM.cMv Bd + 31)

section Call

variable {E : ℕ} {I : List Sym} {iprog : Fin (7 + E) → IStmt HS (7 + E)}

end Call

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZBody
/-!
# KFZ Body: the product phase of the chirp procedure

With the four tables `b̃` (`qB`), `ã` (`qA`), `θ̃` (`qT`), `θ̃⁻` (`qI`) in place, `PP` computes on
stack `0` (input `x`, tail `l`):

`U = Q(x ã)`, `G = Q(U θ̃)`, `F = Q(b̃ θ̃)` (on top of `G`), the Prop 18 call `H = P18(F, G)`,
`V = Q(H θ̃⁻)`, `y = Q(V ã)`, and drops the four tables. **`PP_runs`** gives the exact result
and cost, given the run of the Prop 18 call on its state.
-/

set_option autoImplicit false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue

/-- The product phase (inside `wP setC`). -/
def PP {E : ℕ} (j : Fin (7 + E)) (d18 : Fin HS) : IStmt HS (7 + E) :=
  .seq (Cmul E qA) (.seq (Cmul E qT) (.seq (replP NLogN.KFZ.qB 0) (.seq (Cmul E qT)
    (.seq (Mcall j d18) (.seq (Cmul E qI) (.seq (Cmul E qA) (Zdrops E)))))))

/-- The state handed to the Prop 18 call. -/
def preM (F : Fin HS → List ℤ) (zF zG : List GI) (l : List ℤ) : Fin HS → List ℤ :=
  Function.update F 0 (gflat zF ++ gflat zG ++ l)

section PPs

variable {E : ℕ} (ext : Fin E → IStmt HS (7 + E)) {I : List Sym} (s : IState HS)

end PPs

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZPhase
/-!
# KFZ Phase: the rounded phase tables satisfy `Phases`

The tables are computed at precision `Δ = p + g` and rounded to precision `p` (`KFH.gtr g`).
`Phases` asks for `|θ̃ − θ| ≤ 2·2^{−p}` **and** `|θ̃| ≤ 1`. The second needs more than the
generic `gtr_err`: a grid point of `2^{−p}ℤ[i]` of modulus above `1` has modulus at least
`√(1 + 4^{−p})`, so the table error before rounding must be below about `4^{−p}`.

* **`Qp_unit`**: `‖v − X‖ ≤ 2^{−2p−2}`, `‖X‖ ≤ 1` ⟹ `‖Q_p v − X‖ ≤ 2hp` and `‖Q_p v‖ ≤ 1`.
* **`phases_QA`**: the array form.
* **`norm_chirp_le`**, `norm_twX`, `zeta_pow`: the exact phases.
-/

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

namespace NLogN.KFZ

open IntMul KIL KFB NLogN.KFM Complex

noncomputable section

section Arr

variable {ι : Type*} [Fintype ι]

end Arr

/-! ## The exact phases -/

/-- The twist root `ζ = e^{2πi/(2r)}`. -/
def zetaR (r : ℕ) : ℂ := exp (2 * Real.pi * I * ((1 : ℕ) : ℂ) / ((2 * r : ℕ) : ℂ))

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZRec
/-!
# KFZ Rec: the box arrays as `P18IL` record lists

`P18IL` takes and returns row-major lists of records of `r` coefficients (`RecRep r p`). On the
box `t` with last axis `t_D = r` these are the `RepFull` arrays (`KFMBridge.recL`, `boxOf`):

* `ptB_rec`: the index `j·r + c` is the box point `fullIdx t j c`.
* `castR`: transport of records along `t_D = r` (the identity after `subst`).
* **`rec_in`**: `RepFull p t zs w` gives `RecRep r p zs (castR (recL t w))`, length `∏_{i<D} t_i`.
* **`rec_out`**: a `RecRep` output with the `P18IL` bound gives `RepFull p t zh h` with
  `‖h − mrX t f g‖ ≤ E`.
* **`GB_of_repFull`**: arrays of norm `≤ 2^k` have `p + k + 2`-bit numerators.
* **`mrtCall`**: the per-call map `M̃r` (`h` at the actual operands, `mrX` elsewhere).
-/

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

namespace NLogN.KFZ

open IntMul KIL KFB NLogN.KFM NLogN.KFH Complex

noncomputable section

variable {D : ℕ}

/-- Records transported along `t_D = r`. -/
def castR {a b : ℕ} (h : a = b) (L : List (Fin a → ℂ)) : List (Fin b → ℂ) :=
  L.map fun v c => v (Fin.cast h.symm c)

/-! ## Bit sizes -/

/-! ## The per-call ring convolution -/

open Classical in
/-- `M̃r` for one call: the computed `h` at the actual operands, `mrX` elsewhere. -/
def mrtCall (t : Fin (D + 1) → ℕ) [∀ i, NeZero (t i)]
    (f g h : ((i : Fin (D + 1)) → ZMod (t i)) → ℂ) :
    (((i : Fin (D + 1)) → ZMod (t i)) → ℂ) → (((i : Fin (D + 1)) → ZMod (t i)) → ℂ) →
      (((i : Fin (D + 1)) → ZMod (t i)) → ℂ) :=
  fun a b => if a = f ∧ b = g then h else mrX t a b

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZMath
/-!
# KFZ Math: the four fixed tables of the chirp procedure

For a box `t : Fin (D+1) → ℕ` (every `t_i ∈ {2^h, 2^{h+1}}`, `t_D = 2^{h+1}`), sources `s ≤ t`,
the working precision `Δ = 2p + 2h + d + 8` and the shift `g = p + 2h + d + 8` (`p + g = Δ`):

* `A0 = chirpArr Δ`, `bL = A0.map (gtr g)` (`b̃ = Q_p(chirp)`), `aL = conj bL` (`ã`);
* `τ = twTab Δ (2r) 1 r` rounded by `gtr g`, `tL = τ` repeated `∏_{i<D} t_i` times (`θ̃`),
  `iL = conj tL` (`θ̃⁻`).

The precision bound `d (8T² + 2) 2^{-Δ} ≤ 2^{-(2p+2)}` makes them `Phases` (`phases_QA`).
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM Complex

noncomputable section

/-! ## Precision -/

/-! ## The tables -/

variable (p h D : ℕ) (sv t : Fin (D + 1) → ℕ)

/-- The chirp table at precision `Δ`. -/
def A0 : List GI := chirpArr (dlV p h (D + 1)) (D + 1) sv t
/-- `b̃`. -/
def bLz : List GI := (NLogN.KFZ.A0 p h D sv t).map (KFH.gtr (NLogN.KFZ.gV p h (D + 1)))
/-- `ã`. -/
def aLz : List GI := (bLz p h D sv t).map NLogN.KFZ.gconj
/-- The twist table at precision `Δ`. -/
def twz : List GI := twTab (dlV p h (D + 1)) (2 ^ (h + 2)) 1 (2 ^ (h + 1))
/-- The rounded twist table. -/
def τz : List GI := (twz p h D).map (KFH.gtr (NLogN.KFZ.gV p h (D + 1)))
/-- `θ̃` as a list. -/
def tLz : List GI := (List.replicate (∏ i : Fin D, t i.castSucc) (τz p h D)).flatten
/-- `θ̃⁻` as a list. -/
def iLz : List GI := (tLz p h D t).map NLogN.KFZ.gconj

theorem length_twz : (twz p h D).length = 2 ^ (h + 1) := length_twTab _ _ _ _

variable {p h D sv t}

section Box

variable [∀ i, NeZero (t i)]

/-- The exact twist. -/
def zt (h : ℕ) : ℂ := zetaR (2 ^ (h + 1))

/-- The unrounded twist value at `K`. -/
def twv (p h D : ℕ) (t : Fin (D + 1) → ℕ) (K : (i : Fin (D + 1)) → ZMod (t i)) : ℂ :=
  vP (dlV p h (D + 1)) ((twz p h D).getD (K (Fin.last D)).val 0)

/-! ## Norms and bit sizes -/

end Box

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZFull
/-!
# KFZ Full: the whole chirp procedure body

`body j d18 = setupP; T1; …; T7; wP setC (PP j d18); popsL regL`.

**`body_runs`**: from `encK s f` (stack `0` = `gflat xs ++ l`, the chirp descriptor on `cS`, all
other managed stacks empty) the body ends in `encK s (update f 0 (gflat ys ++ l))` with
`ys = Q(Q(H θ̃⁻) ã)`, given the run of the Prop 18 call on its state.
-/

set_option autoImplicit false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue

/-- The body of the chirp procedure. -/
def body {E : ℕ} (j : Fin (7 + E)) (d18 : Fin HS) : IStmt HS (7 + E) :=
  .seq setupP.code (.seq (T1 E) (.seq (T2 E) (.seq (T3 E) (.seq (T4 E) (.seq (T5 E) (.seq (T6 E)
    (.seq (T7 E) (.seq (NLogN.KFZ.wP NLogN.KFZ.setC (NLogN.KFZ.PP j d18)) (NLogN.KFZ.popsL regL).code))))))))

/-- The empty-elsewhere condition of the input state. -/
def EmptyEx (f : Fin HS → List ℤ) : Prop := ∀ j, j ≠ 0 → j ≠ NLogN.KFZ.cS → f j = []

/-- The register values of the final pops. -/
def vsR (p h d σ K : ℕ) : List (Fin HS × ℤ) :=
  [(rgP, (p : ℤ)), (rgH, (h : ℤ)), (rgD, (d : ℤ)), (rgS, (σ : ℤ)), (rgK, (K : ℤ))]

/-! ## The product-phase state -/

/-- The parameters of the product phase. -/
def vsC (p h d σ : ℕ) : List (Fin HS × ℤ) :=
  [(NLogN.KFM.pR, 2 ^ p), (nX, 2 ^ (d * h + σ)), (nX2, ((2 * 2 ^ (d * h + σ) : ℕ) : ℤ)), (KFM.rN, 1),
    (rN2, ((2 * 2 ^ (d * h + σ) : ℕ) : ℤ))]

/-- The state entering the product phase (parameters pushed). -/
def FC (f : Fin HS → List ℤ) (p h D σ K : ℕ) (sv t : Fin (D + 1) → ℕ) (z : List ℤ) :
    Fin HS → List ℤ :=
  pushF (St f p h (D + 1) σ K [] z [] (gflat (bLz p h D sv t)) (gflat (aLz p h D sv t))
    (gflat (tLz p h D t)) (gflat (iLz p h D t))) (vsC p h (D + 1) σ)

/-- `F = Q(b̃ θ̃)`. -/
def zFz (p h D : ℕ) (sv t : Fin (D + 1) → ℕ) : List GI :=
  List.zipWith (gmulT p) (bLz p h D sv t) (tLz p h D t)

/-- `G = Q(Q(x ã) θ̃)`. -/
def zGz (p h D : ℕ) (sv t : Fin (D + 1) → ℕ) (xs : List GI) : List GI :=
  List.zipWith (gmulT p) (List.zipWith (gmulT p) xs (aLz p h D sv t)) (tLz p h D t)

/-- `y = Q(Q(H θ̃⁻) ã)`. -/
def yz (p h D : ℕ) (sv t : Fin (D + 1) → ℕ) (zh : List GI) : List GI :=
  List.zipWith (gmulT p) (List.zipWith (gmulT p) zh (iLz p h D t)) (aLz p h D sv t)

/-- Cost of the product phase without the Prop 18 call. -/
def cPP (w p N B Bd b len : ℕ) : ℕ :=
  cCmul w p N + (cCmul w p N + (cRepl 1 (2 * N) w + (cCmul w p N +
    (cM B Bd b len + (cCmul w p N + (cCmul w p N + 4 * cMvX (2 * N) w))))))

/-- Cost of the body without the Prop 18 call. -/
def cBody (B Bd b len p h D σ : ℕ) (sv t : Fin (D + 1) → ℕ) : ℕ :=
  10 ^ 10 * (D + 1 + 1) * (B + 2) ^ 2 +
  ((cPush 3 B + (cChirpTabU (dlV p h (D + 1)) (KFM.fmB (2 * (dlV p h (D + 1) + 2))) (NLogN.KFM.axL sv t)
    (∏ i, t i) + cPop 3 B)) +
  ((cPush 4 B + ((2 * cMvX (2 * 2 ^ ((D + 1) * h + σ)) (dlV p h (D + 1) + 3) +
    ((cItem (2 * 2 ^ ((D + 1) * h + σ)) (3000 * (dlV p h (D + 1) + 3 + NLogN.KFZ.gV p h (D + 1) + 3))
      (2 * (dlV p h (D + 1) + 3) + 1) + 2 * cCnt (2 * 2 ^ ((D + 1) * h + σ))) +
    2 * cMvX (2 * 2 ^ ((D + 1) * h + σ)) (p + 2))) + cPop 4 B)) +
  ((cPush 2 B + (cCnj (2 ^ ((D + 1) * h + σ)) (p + 2) + cPop 2 B)) +
  ((cPush 5 B + (cTwGU (dlV p h (D + 1)) (2 ^ (h + 2)) 1 (2 ^ (h + 1))
    (KFM.fmB (2 * (dlV p h (D + 1) + 2))) + cPop 5 B)) +
  ((cPush 4 B + ((2 * cMvX (2 * 2 ^ (h + 1)) (dlV p h (D + 1) + 3) +
    (cItem (2 * 2 ^ (h + 1)) (3000 * (dlV p h (D + 1) + 3 + NLogN.KFZ.gV p h (D + 1) + 3))
      (2 * (dlV p h (D + 1) + 3) + 1) + 2 * cCnt (2 * 2 ^ (h + 1)))) + cPop 4 B)) +
  ((cPush 2 B + ((cRepl (∏ i : Fin D, t i.castSucc) (2 * 2 ^ (h + 1)) (p + 2) +
    cMvX (2 * 2 ^ (h + 1)) (p + 2)) + cPop 2 B)) +
  ((cPush 2 B + (cCnj (2 ^ ((D + 1) * h + σ)) (p + 2) + cPop 2 B)) +
  ((cPush 5 B + (cPP (p + 2) p (2 ^ ((D + 1) * h + σ)) B Bd b len + cPop 5 B)) +
    (5 * (2 * B + 2) + 1)))))))))

section Body

variable {E : ℕ} (ext : Fin E → IStmt HS (7 + E)) {I : List Sym} (s : IState HS)

end Body

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZCost
/-!
# KFZ Cost: a closed-form bound for the body cost

**`cBody_le`**: with a master width `P` (bounding `Δ + 3`, `B`, `Bd`, `b`, the descriptor
length and `2h + 6`), `N = ∏ t_i`, `r = 2^{h+1}`, `d r ≤ N`:

`cBody ≤ 10^30 (N + 1)(fmB(8P) + P + 1) + 10^30 (P + 31)^6`.
-/

set_option autoImplicit false
set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option maxRecDepth 100000

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM

/-! ## Single pieces -/

/-! ## The pieces of `cBody` -/

section Pieces

variable {p h D σ B Bd b len P N : ℕ} {sv t : Fin (D + 1) → ℕ}

/-- `fmB` at the master width. -/
abbrev FM (P : ℕ) : ℕ := KFM.fmB (8 * P)

end Pieces

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZParam
/-!
# KFZ Param: the chirp procedure parameters for `n`

For `n ≥ n0` and `dP n = D + 1`: `h = ℓ - 1`, `σ = D + 1 - dpN n`, bits `bs`, sources `ss`.
`tBD n D i = 2^{h + bit i}`, `(D+1) h + σ = L`, `∏ t = 2^L`, the setup descriptor is
`descChirpN n`, the chirp-table descriptor is `descT (sBD n D) (tBD n D)`, and the register width
`B` satisfies `RB` and `B ≤ 5 p`.
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue IntMul.OpenAI.RP

noncomputable section

/-- The bit of axis `j`. -/
def bitN (n j : ℕ) : ℕ := if dpN n ≤ j then 1 else 0

/-- The bit list. -/
def bsN (n : ℕ) : List ℕ := (List.range (IntMul.OpenAI.Glue.dP n)).map (bitN n)

/-- The source list. -/
def ssN (n : ℕ) : List ℕ := (List.range (IntMul.OpenAI.Glue.dP n)).map (srcN n)

/-- `h = ℓ - 1`. -/
def hN (n : ℕ) : ℕ := ℓP n - 1

/-- `σ = d - dpN`. -/
def σN (n : ℕ) : ℕ := IntMul.OpenAI.Glue.dP n - dpN n

/-- The register width. -/
def BN (n : ℕ) : ℕ := dlV (pP n) (NLogN.KFZ.hN n) (IntMul.OpenAI.Glue.dP n) + (IntMul.OpenAI.Glue.dP n * NLogN.KFZ.hN n + σN n) + 8 + KP n + αP n + 1

section Facts

variable {n D : ℕ} (hn : n0 ≤ n) (hD : dP n = D + 1)
include hn hD

theorem length_bsN : (NLogN.KFZ.bsN n).length = D + 1 := by simp [NLogN.KFZ.bsN, hD]

theorem length_ssN : (ssN n).length = D + 1 := by simp [ssN, hD]

end Facts

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZCostR
/-!
# KFZ CostR: the chirp procedure cost against `cLineU`

**`cost_final`**: the Prop 18 call within its `P18IL` bound plus `cBody + 2` is at most
`(2 C₁₈ + Kc) · cLineU n (A₁₈ + 2)`.
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

attribute [local irreducible] fmA Cfm

/-- The constant of the body. -/
def Kc : ℝ := 10 ^ 30 * (2 * (32500 * (Cfm : ℝ) + 6) + 6 ^ 12) + 2

section Facts

variable {n D : ℕ} (hn : n0 ≤ n) (hD : dP n = D + 1)
include hn hD

end Facts

end

end NLogN.KFZ

end

section
-- Source module: Solutions.KFZMain
/-!
# KFZ Main: `P18IL → ChirpNeedR`

The chirp procedure `body` (`KFZFull`) is kmain's own procedure `0` in the merged table
`mExt ext₁₈ ∅ ∅ [body]`; the Prop 18 procedure is reached through `gI` (`mTable_g1`,
`IRuns.relabel`). The descriptor stack is `cS = 36` (the chirp descriptor stays there and is
restored by `setupP`).

Per call, `M̃r = mrtCall t F G H` with `F = Q(b̃ θ̃)`, `G = Q(θ̃ Q(ã x))` and `H` the Prop 18 result
at exactly these operands; `rec_out` gives `‖H − mrX F G‖ ≤ ErC hp`.
-/

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

namespace NLogN.KFZ

open IntMul NLogN.KIL KFB NLogN.KFM IntMul.OpenAI.Glue IntMul.OpenAI.RCost

noncomputable section

/-- The merged extension list: Prop 18, nothing, nothing, the chirp body. -/
def extZ {E18 : ℕ} (ext18 : Fin E18 → IStmt HS (7 + E18)) (entry18 : Fin (7 + E18))
    (d18 : Fin HS) : Fin (E18 + 0 + 0 + 1) → IStmt HS (7 + (E18 + 0 + 0 + 1)) :=
  mExt ext18 (Fin.elim0 : Fin 0 → IStmt HS (7 + 0)) (Fin.elim0 : Fin 0 → IStmt HS (7 + 0))
    (fun _ => body (NLogN.KFM.gI E18 (E18 + 0 + 0 + 1) 0 (by omega) entry18) d18)

end

end NLogN.KFZ

end