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

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

import Definitions.Def_IntMul_Kappa_K_Defs16

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.SSPow
/-! # Powers of two: `PW : m :: L ↦ 2^m :: L` (recursive on the bits of `m`) -/

namespace NLogN

open IntMul

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

def pow2Base : IStmt NS P := .seq (.op (.drop PW)) (.op (.const PW 1))

def pow2A : IStmt NS P :=
  .seq (.op (.const T7 2)) (.seq (.op (.split PW T7 T6)) (.seq (.op (.drop T7)) (.op (.mv T6 PB))))

def pow2D : IStmt NS P :=
  .seq (.op (.dup PW T6)) (.seq (.op (.shl PW T6)) (.seq (.op (.drop T6)) (.op (.mv PB T6))))

def pow2E : IStmt NS P :=
  .ite T6 (fun c => c == .pos) (.seq (.op (.const T7 2)) (.seq (.op (.shl PW T7)) (.op (.drop T7))))
    .skip

def pow2Body (pP : Fin P) : IStmt NS P :=
  .ite PW (fun c => c == .zero) pow2Base
    (.seq pow2A (.seq (.call pP) (.seq pow2D (.seq pow2E (.op (.drop T6))))))

def pow2CostS (m : ℕ) : ℕ := 20000 * (m + 1)

end NLogN

end

section
-- Source module: NLogN.SSProg
/-!
# The Schönhage–Strassen procedure `FM`: program text

`FM` multiplies modulo `2^N+1`: `ARG : a :: b :: L ↦ (a*b mod 2^N+1) :: L`, with `N` on top of
`PRM` and the ring rulers `2^N`, `2^N+1` on top of `RN`, `RF`.  Small `N` use schoolbook
multiplication; otherwise the inputs are cut into `K = 2^k` pieces of `M` bits, twisted and
transformed in the ring `ℤ/(2^N'+1)`, multiplied pointwise by recursive calls, transformed back,
untwisted, lifted to signed coefficients, carried into `M`-bit digits and recombined.
-/

namespace NLogN

open IntMul

variable {P : ℕ}

/-- `ARG : a :: b :: L ↦ a'' :: b'' :: L` with `x'' = (x mod 2^N) - ⌊x/2^N⌋`. -/
def fmPre : IStmt NS P :=
  .seq (.op (.split ARG RN T1)) (.seq (.op (.neg ARG)) (.seq (.op (.add T1 ARG ARG))
    (.seq (.op (.mv ARG T1)) (.seq (.op (.split ARG RN T2)) (.seq (.op (.neg ARG))
      (.seq (.op (.add T2 ARG ARG)) (.op (.mv T1 ARG))))))))

/-- Move `K` integers (`K` on top of `DL`) from `X` to `Y`. -/
def mvKS (X Y : Fin NS) : IStmt NS P := .seq (.op (.dup DL CC)) (.seq (NLogN.mvN X Y CC) (.op (.drop CC)))

/-- One piece: cut `M` bits off `XA`, twist by `W = 2^(iQ)`, reduce, push to `QS`. -/
def prepStep : IStmt NS P :=
  .seq (.op (.split XA RMs T4)) (.seq (.op (.shl T4 W)) (.seq (fred T4) (.seq (.op (.mv T4 QS))
    (.seq (.op (.shl W RQ)) (decC CC)))))

/-- `ARG : x :: L ↦ L`, `DX : R ↦ (forward transform of the twisted pieces of x) ++ R`. -/
def fmPrep (pD : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.mv ARG XA)) (.seq (.op (.const W 1)) (.op (.dup DL CC))))
    (.seq (.loop CC (fun c => c != .zero) prepStep)
      (.seq (.seq (.op (.drop CC)) (.seq (.op (.drop W)) (.op (.drop XA))))
        (.seq (mvKS QS DX) (.call pD))))

/-- One pointwise product by a recursive call. -/
def pwStep (pF : Fin P) : IStmt NS P :=
  .seq (.op (.mv DX ARG)) (.seq (.op (.mv NLogN.V ARG)) (.seq (.call pF) (.seq (.op (.mv ARG NLogN.P1))
    (decC CC))))

def fmPw (pF : Fin P) : IStmt NS P :=
  .seq (.op (.dup DL CC)) (.seq (.loop CC (fun c => c != .zero) (pwStep pF)) (.op (.drop CC)))

/-- Fix a negative remainder of the signed split. -/
def fixS : IStmt NS P :=
  .seq (.op (.dup RMs T6)) (.seq (.op (.add T6 T5 T5)) (.seq (.op (.const T6 (-1)))
    (.op (.add T6 T4 T4))))

/-- Untwist, lift to the symmetric range, add the carry and split off an `M`-bit digit. -/
def postStep : IStmt NS P :=
  .seq (.seq (.op (.mv DX T4)) (.op (.shl T4 W))) (.seq (fred T4)
    (.seq (.seq (.op (.shl T4 W2)) (.seq (.op (.dup HB T5)) (.op (.add T5 T4 T4)))) (.seq (fred T4)
      (.seq (.seq (.op (.dup HB T5)) (.seq (.op (.neg T5)) (.seq (.op (.add T5 T4 T4))
        (.seq (.op (.mv CYS T5)) (.seq (.op (.add T5 T4 T4)) (.op (.split T4 RMs T5)))))))
      (.seq (.ite T5 (fun c => c == .neg) fixS .skip)
        (.seq (.op (.mv T5 NLogN.DG)) (.seq (.op (.mv T4 CYS)) (.seq (.op (.split W RQ T5))
          (.seq (.op (.drop T5)) (decC CC))))))))))

def fmPost : IStmt NS P :=
  .seq (.seq (.op (.dup RN W)) (.seq (.op (.const CYS 0)) (.op (.dup DL CC))))
    (.seq (.loop CC (fun c => c != .zero) postStep) (.seq (.op (.drop CC)) (.op (.drop W))))

def hornerStep : IStmt NS P := .seq (.op (.cat ZZ RMs NLogN.DG)) (decC CC)

def fmHorner : IStmt NS P :=
  .seq (.seq (.op (.const ZZ 1)) (.op (.dup DL CC)))
    (.seq (.loop CC (fun c => c != .zero) hornerStep) (.op (.drop CC)))

/-- `ZZ := (ZZ - 2^N - CYS) mod (2^N+1)`, pushed to `ARG`. -/
def fmFinal : IStmt NS P :=
  .seq (.seq (.op (.dup RN T4)) (.seq (.op (.neg T4)) (.seq (.op (.add T4 ZZ ZZ))
    (.seq (.op (.mv CYS T4)) (.seq (.op (.neg T4)) (.op (.add T4 ZZ ZZ)))))))
    (.seq (fred ZZ) (.op (.mv ZZ ARG)))

/-! ## Parameters -/

/-- `ZZ := ℓ = blen N`, `T1 := k = ℓ/2 - blen ℓ`, `DL := K = 2^k` (the latter via `pow2`). -/
def parA (pP : Fin P) : IStmt NS P :=
  .seq (.seq (blenS PRM ZZ) (.seq (blenS ZZ CYS) (.seq (.op (.const T4 2)) (.seq (.op (.dup ZZ T1))
    (.seq (.op (.split T1 T4 T2)) (.seq (.op (.drop T2)) (.seq (.op (.drop T4))
      (.seq (.op (.mv CYS T2)) (.seq (.op (.neg T2)) (.seq (.op (.add T2 T1 T1))
        (.op (.dup T1 PW))))))))))))
    (.seq (.call pP) (.op (.mv PW DL)))

/-- `T2 := M = N/K`, `RMs := 2^M`. -/
def parB (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup PRM T2)) (.seq (.op (.split T2 DL NLogN.T3)) (.seq (.op (.drop NLogN.T3))
    (.op (.dup T2 PW)))))
    (.seq (.call pP) (.op (.mv PW RMs)))

/-- `T2 := Q = (2M+2k+3)/K + 1`, `T3 := N' = Q*K`, `RQ := 2^Q`, `DE := 2^(2Q)`. -/
def parC (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup T2 NLogN.T3)) (.seq (.op (.add NLogN.T3 T2 T2)) (.seq (.op (.dup T1 NLogN.T3))
    (.seq (.op (.dup T1 T4)) (.seq (.op (.add NLogN.T3 T4 T4)) (.seq (.op (.add T4 T2 T2))
      (.seq (.op (.const NLogN.T3 3)) (.seq (.op (.add NLogN.T3 T2 T2)) (.seq (.op (.split T2 DL NLogN.T3))
        (.seq (.op (.drop NLogN.T3)) (.seq (.op (.const NLogN.T3 1)) (.seq (.op (.add NLogN.T3 T2 T2))
          (.seq (.op (.dup T2 NLogN.T3)) (.seq (.op (.shl NLogN.T3 DL)) (.op (.dup T2 PW))))))))))))))))
    (.seq (.call pP) (.seq (.op (.mv PW RQ)) (.seq (.op (.dup RQ T4)) (.seq (.op (.shl T4 RQ))
      (.op (.mv T4 NLogN.DE))))))

/-- The ring `2^N'+1` on `RN`, `RF`. -/
def parD (pP : Fin P) : IStmt NS P :=
  .seq (.op (.dup NLogN.T3 PW)) (.seq (.call pP) (.seq (.op (.dup PW T4)) (.seq (.op (.mv PW RN))
    (.seq (.op (.const T5 1)) (.seq (.op (.add T5 T4 T4)) (.op (.mv T4 RF)))))))

/-- `W2 := 2^(N'-k)`, `HB := 2^(N'-1)`, `PRM := N'`, and the temporaries are dropped. -/
def parE (pP : Fin P) : IStmt NS P :=
  .seq (.seq (.op (.dup NLogN.T3 PW)) (.seq (.op (.dup T1 T4)) (.seq (.op (.neg T4))
    (.op (.add T4 PW PW)))))
    (.seq (.call pP) (.seq (.seq (.op (.mv PW W2)) (.seq (.op (.dup NLogN.T3 PW))
      (.seq (.op (.const T4 (-1))) (.op (.add T4 PW PW)))))
    (.seq (.call pP) (.seq (.op (.mv PW HB)) (.seq (.op (.mv NLogN.T3 PRM)) (.seq (.op (.drop T2))
      (.seq (.op (.drop T1)) (.op (.drop ZZ)))))))))

def fmParams (pP : Fin P) : IStmt NS P :=
  .seq (parA pP) (.seq (parB pP) (.seq (NLogN.parC pP) (.seq (parD pP) (parE pP))))

/-- Drop the inner ring and the transform parameters (keeping `RMs`, `DL`). -/
def fmUnpar : IStmt NS P :=
  .seq (.op (.drop RN)) (.seq (.op (.drop RF)) (.seq (.op (.drop PRM)) (.seq (.op (.drop W2))
    (.seq (.op (.drop HB)) (.seq (.op (.drop NLogN.DE)) (.op (.drop RQ)))))))

def fmMain (pF pD pT pP : Fin P) : IStmt NS P :=
  .seq fmPre (.seq (fmParams pP) (.seq (fmPrep pD) (.seq (mvKS DX U) (.seq (mvKS U NLogN.V)
    (.seq (fmPrep pD) (.seq (fmPw pF) (.seq (mvKS NLogN.P1 DX) (.seq (.call pT) (.seq fmPost
      (.seq fmUnpar (.seq fmHorner (.seq fmFinal (.seq (.op (.drop RMs)) (.op (.drop DL)))))))))))))))

/-- Schoolbook multiplication and reduction for small `N`. -/
def fmBase : IStmt NS P := .seq (smul ARG XA QS W CC U NLogN.V) (fred ARG)

/-- Threshold on the bit length of `N`. -/
def fmThr : ℕ := 2048

def fmBody (pF pD pT pP : Fin P) : IStmt NS P :=
  .seq (.seq (blenS PRM ZZ) (.seq (.op (.const T4 (-(fmThr : ℤ)))) (.op (.add T4 ZZ ZZ))))
    (.ite ZZ (fun c => c == .neg) (.seq (.op (.drop ZZ)) fmBase)
      (.seq (.op (.drop ZZ)) (fmMain pF pD pT pP)))

end NLogN

end

section
-- Source module: NLogN.SSFM1
/-!
# `FM`: correctness of the preprocessing segments
-/

namespace NLogN

open IntMul

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

/-! ## Pieces -/

/-- The quotient after `i` cuts of `M` bits. -/
def qit (M i : ℕ) (x : ℤ) : ℤ := (fun y : ℤ => sg y (y.natAbs / 2 ^ M))^[i] x

def prepCost (N n k : ℕ) : ℕ := difCost n k + 40000 * 2 ^ k * (n + 3) + 1000 * (N + n + 10)

end NLogN

end

section
-- Source module: NLogN.SSFM2
/-!
# `FM`: pointwise products, untwisting with carries, and recombination
-/

namespace NLogN

open IntMul

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

/-- The specification of a call of `FM` in the ring `2^n+1`, with cost at most `C`. -/
def FMSpec (pF : Fin P) (iprog : Fin P → IStmt NS P) (I : List Sym) (n C : ℕ) : Prop :=
  ∀ (s : IState NS) (a b : ℤ) (LA LP LN LF : List Item), 0 ≤ a ∧ a ≤ 2 ^ n → 0 ≤ b ∧ b ≤ 2 ^ n →
    s.G ARG = .int a :: .int b :: LA → s.G PRM = .int n :: LP → s.G RN = .int (2 ^ n) :: LN →
    s.G RF = .int (2 ^ n + 1) :: LF →
    IRuns I iprog (.call pF) s (s.upd ARG (.int ((a * b) % (2 ^ n + 1)) :: LA)) C

end NLogN

end

section
-- Source module: NLogN.SSFM3
/-!
# `FM`: untwisting, lifting and carrying
-/

namespace NLogN

open IntMul

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

/-- Untwist by `2^e`, unweight by `2^(n-k)` and lift to the symmetric range. -/
def liftE (n k e : ℕ) (d : ℤ) : ℤ :=
  (((d * 2 ^ e) % (2 ^ n + 1)) * 2 ^ (n - k) + 2 ^ (n - 1)) % (2 ^ n + 1) - 2 ^ (n - 1)

/-- The lifted coefficients of the inverse transform `D` (with `K` entries). -/
def liftL (n k Q K : ℕ) (D : List ℤ) : List ℤ :=
  (List.range K).map fun i => liftE n k ((K - i) * Q) (D.getD i 0)

end NLogN

end

section
-- Source module: NLogN.SSFM4
/-!
# `FM`: parameter computation
-/

namespace NLogN

open IntMul

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

/-- The parameters of `FM` at size `N`. -/
def fmL (N : ℕ) : ℕ := zsz (N : ℤ)
def fmB (N : ℕ) : ℕ := zsz (fmL N : ℤ)
def fmK (N : ℕ) : ℕ := fmL N / 2 - fmB N
def fmM (N : ℕ) : ℕ := N / 2 ^ fmK N
def fmQ (N : ℕ) : ℕ := (2 * fmM N + 2 * fmK N + 3) / 2 ^ fmK N + 1
def fmN' (N : ℕ) : ℕ := fmQ N * 2 ^ fmK N

end NLogN

end

section
-- Source module: NLogN.SSCorrect
/-!
# Correctness of one Schönhage–Strassen level
-/

namespace NLogN

open IntMul

/-- The folded input `x mod 2^N - ⌊x / 2^N⌋`. -/
def fold (N : ℕ) (x : ℤ) : ℤ := x % 2 ^ N + -(x / 2 ^ N)

end NLogN

end

section
-- Source module: NLogN.SSFM5
/-!
# One level of `FM`
-/

namespace NLogN

open IntMul

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

/-- The cost of one level of `FM` beyond the recursive calls. -/
def fmLevelCost (N : ℕ) : ℕ :=
  2 * difCost (fmN' N) (fmK N) + ditCost (fmN' N) (fmK N) +
    200000 * 2 ^ fmK N * (fmN' N + 3) + 2000 * (fmL N + 1) ^ 2 + 400000 * (N + 3)

end NLogN

end

section
-- Source module: NLogN.SSParam
/-!
# Parameter inequalities for `FM`

For `ℓ = fmL N ≥ 2048` and `2^k ∣ N` (`k = fmK N`), the recursive size `fmN' N` is smaller,
keeps the divisibility invariant, and the potential `fmPhi` absorbs one level.
-/

namespace NLogN

/-- `λ = blen (ℓ^4)`. -/
def fmLam (N : ℕ) : ℕ := zsz (((fmL N) ^ 4 : ℕ) : ℤ)

/-- The constant of the potential. -/
def fmA : ℕ := 2 ^ 2100

/-- The potential `A · N · ℓ · λ`. -/
def fmPhi (N : ℕ) : ℕ := fmA * N * fmL N * fmLam N

/-! ## Arithmetic helpers -/

end NLogN

end

section
-- Source module: NLogN.SSMain
/-!
# The top-level Schönhage–Strassen program
-/

namespace NLogN

open IntMul

section blen

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

end blen

/-- The FM size used for `n`-bit inputs. -/
def mainN0 (n : ℕ) : ℕ := 2 ^ (zsz ((n + 1 : ℕ) : ℤ) + 1)

def mainACost (n : ℕ) : ℕ :=
  2000 * (n + 2) * (zsz ((n + 1 : ℕ) : ℤ) + 2) + 2000 * (zsz ((n + 1 : ℕ) : ℤ) + 1) *
    (zsz (zsz ((n + 1 : ℕ) : ℤ) : ℤ) + 2) + 60000 * (mainN0 n + 3)

def mainCost (n : ℕ) : ℕ :=
  2 * readRCost n + mainACost n + fmPhi (mainN0 n) + shlCost n + outRCost (2 * n)

end NLogN

end

section
-- Source module: Solutions.KFMMul
/-!
# KFM 6: the Schönhage–Strassen FM as a `MulOK` multiplication procedure

nlogn's sorry-free Fermat-ring multiplier `fm_spec'` (SSFM6) multiplies `0 ≤ a, b ≤ 2^N` modulo
`2^N + 1`.  This file wraps it into an integer multiplier of arbitrary signed operands:

* **`mulBody`** (32 stacks `NS`, procedure `0` of the 5-procedure table `mulProg`, whose
  procedures `1–4` are `fmBody`, `difBody`, `ditBody`, `pow2Body`): signs are stripped (sign
  kept on `T3`), `n = zsz a + zsz b` is computed with `blenS`, the FM parameters
  `N = mainN0 n = 2^(zsz(n+1)+1)`, `2^N`, `2^N+1` are pushed as in SS `mainBody`, FM is called
  (the product `< 2^n < 2^N` so no reduction happens), parameters are dropped and the sign is
  restored.  **`mulBody_runs`**: `ARG = a :: b :: L ↦ ARG = a*b :: L`, all other stacks
  unchanged, cost `mulCost (zsz a + zsz b)`.
* **`mulOK_of_embed`**: for any `HS`-table `iprog'` containing the renamed `mulProg`
  (`iprog' (g j) = (mulProg j).ren ssE g`) and a wrapper procedure `pM` with body `mulWrap g`,
  `MulOK I iprog' pM ⟨fun a b => mulCostH (zsz a + zsz b), fun _ _ => True⟩`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false
set_option linter.unusedSectionVars false
set_option linter.style.longLine false

namespace NLogN.KFM

open IntMul NLogN

section NSLevel

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

/-- Take the absolute value of the top of `ARG`, flipping the sign token on `T3`. -/
def absC : IStmt NS P :=
  .ite ARG (fun c => c == .neg) (.seq (.op (.neg ARG)) (.op (.neg NLogN.T3))) .skip

/-- The sign flip of `σ` by the sign of `x`. -/
def fl (x σ : ℤ) : ℤ := if x < 0 then -σ else σ

/-- Strip both signs; the product sign goes to `T3`. -/
def signC : IStmt NS P :=
  .seq (.op (.const NLogN.T3 1)) (.seq absC (.seq (.op (.mv ARG T2)) (.seq absC (.op (.mv T2 ARG)))))

/-- Bit lengths: push `zsz x + zsz y + 1` onto `ZZ`. -/
def lenC : IStmt NS P :=
  .seq (blenS ARG ZZ) (.seq (.op (.mv ARG T2)) (.seq (blenS ARG ZZ) (.seq (.op (.mv T2 ARG))
    (.seq (.op (.mv ZZ T4)) (.seq (.op (.add T4 ZZ ZZ)) (.seq (.op (.const T4 1))
      (.op (.add T4 ZZ ZZ))))))))

/-- The FM parameters from `ZZ = n+1` (the tail of SS `mainBody.mainA`, with the
`pow2` procedure `pP`). -/
def parC (pP : Fin P) : IStmt NS P :=
  .seq (blenS ZZ PW) (.seq (.seq (.op (.const T4 1)) (.op (.add T4 PW PW)))
    (.seq (.call pP) (.seq (.seq (.op (.mv PW PRM)) (.op (.dup PRM PW))) (.seq (.call pP)
      (.seq (.op (.dup PW T4)) (.seq (.op (.mv PW RN)) (.seq (.op (.const T5 1))
        (.seq (.op (.add T5 T4 T4)) (.op (.mv T4 RF))))))))))

/-- Drop the FM parameters and `ZZ`, apply the sign from `T3`, drop it. -/
def cleanC : IStmt NS P :=
  .seq (.op (.drop PRM)) (.seq (.op (.drop RN)) (.seq (.op (.drop RF)) (.seq (.op (.drop ZZ))
    (.seq (.ite NLogN.T3 (fun c => c == .neg) (.op (.neg ARG)) .skip) (.op (.drop NLogN.T3))))))

/-- The signed multiplier body: `ARG = a :: b :: L ↦ a*b :: L`. -/
def mulBody (pF pP : Fin P) : IStmt NS P :=
  .seq signC (.seq lenC (.seq (parC pP) (.seq (.call pF) cleanC)))

/-- The cost of `mulBody` on operands of total bit length `n`. -/
def mulCost (n : ℕ) : ℕ :=
  300 * (n + 2) + 5000 * (n + 3) * (zsz ((n + 1 : ℕ) : ℤ) + 2) + mainACost n +
    fmPhi (mainN0 n) + 10 * (mainN0 n + n + 20)

end NSLevel

/-! ## The 5-procedure NS table -/

/-- Procedure `0`: the signed multiplier; `1–4`: nlogn's FM, `dif`, `dit`, `pow2`. -/
def mulProg : Fin 5 → IStmt NS 5 :=
  ![mulBody 1 4, fmBody 1 2 3 4, difBody 2, ditBody 3, pow2Body 4]

/-! ## Embedding into the `HS` stack machine -/

/-- The `NS` stacks are the first 32 `HS` stacks. -/
def ssE : Fin NS → Fin HS := Fin.castLE (by decide)

/-- The `HS` wrapper: move `a`, `b` from `ARGM` to `ssE ARG`, call the renamed multiplier
(procedure `g 0`), move the product back. -/
def mulWrap {P' : ℕ} (g : Fin 5 → Fin P') : IStmt HS P' :=
  .seq (.op (.mv ARGM (ssE ARG))) (.seq (.op (.mv ARGM (ssE ARG))) (.seq (.call (g 0))
    (.op (.mv (ssE ARG) ARGM))))

/-- The cost of one `MulOK` call on operands of total bit length `n`. -/
def mulCostH (n : ℕ) : ℕ := mulCost n + 400 * (n + 3)

end NLogN.KFM

end

section
-- Source module: Solutions.KFSMain
/-!
# KFS §2: Lemma 12 on arrays of `w`-bit integers and of Gaussian records

Route (b) of `MACHINE-PLAN.md` §3 item 2(iii): convert the items to bits (`KFL.toBitsI`, sign
plus `w` magnitude bits per integer), so that the item axis becomes an innermost bit axis of
size `w + 1`, interchange the two `2^u` fields with the fine Lemma 12 wrapper `KFL.swI`, and
convert back (`KFL.fromBitsI`).

* `swW_spec`: integer arrays on `binShape Pre G Suf (2^u)`, cost `≤ K (w+1) (V+1) u^τ`.
* `swAx_spec`: the same for any two equal `2^u` axes `p < q` of an arbitrary positive shape.
* `swAxG_spec`: Gaussian arrays (`gflat`, `GB w`) on any shape, in the `encK` form with a tail
  `l` after the array on stack 0 (records of `r` entries: use the shape `sh ++ [r]`).

All codes are `IStmt HS 1` programs run under `iprog = KFL.l12prog`; they depend on the
Lemma 12 prime `q`, on `w` and on the shape.  Any tail `R` is allowed on stack 0.
-/

set_option linter.unusedSimpArgs false
set_option linter.unusedVariables false

namespace NLogN.KFS

open IntMul KIL KFB KFL KFE AModel AModelL12

/-- `τ = 1 - 2^-50`. -/
noncomputable abbrev τS : ℝ := 1 - (2 : ℝ)⁻¹ ^ 50

/-! ## Helpers -/

/-! ## Integer arrays on `binShape` -/

/-- **The code**: to bits, Lemma 12 on the bit array `binShape Pre G (Suf ++ [w+1]) (2^u)`,
back to integers. -/
def swW (q w : ℕ) (Pre G Suf : Shape) (u : ℕ) : IStmt HS 1 :=
  .seq (toBitsI (vol (binShape Pre G Suf (2 ^ u))) w)
    (.seq (swI q (vol Pre) (vol G) (vol Suf * (w + 1)) u)
      (fromBitsI (vol (binShape Pre G Suf (2 ^ u))) w))

/-! ## Any two equal `2^u` axes of an arbitrary shape -/

/-- The prefix before axis `p`. -/
def preS (sh : Shape) (p : ℕ) : Shape := sh.take p
/-- The axes strictly between `p` and `q`. -/
def midS (sh : Shape) (p q : ℕ) : Shape := (sh.drop (p + 1)).take (q - p - 1)
/-- The suffix after axis `q`. -/
def sufS (sh : Shape) (q : ℕ) : Shape := sh.drop (q + 1)

/-- **The code** for the axes `p < q` of `sh`. -/
def swAx (q' w : ℕ) (sh : Shape) (p q u : ℕ) : IStmt HS 1 :=
  swW q' w (preS sh p) (midS sh p q) (sufS sh q) u

/-! ## Gaussian arrays (`gflat`), `encK` form -/

/-- **The code** for Gaussian arrays: the integer code on `sh ++ [2]`. -/
def swAxG (q' w : ℕ) (sh : Shape) (p q u : ℕ) : IStmt HS 1 := swAx q' w (sh ++ [2]) p q u

end NLogN.KFS

end

section
-- Source module: Solutions.KFTBase
/-!
# KFT §0: a counted repetition with a literal count

`rpt A C D n b`: push the literal `n` on `A`, turn it into `n` marks on `C` (KFC `pro`, linear
`cUL n`), run `b` once per mark (`uloop`, `64` per iteration), clean up (`epi`) and drop `n`.
`rpt_runs`: if `b` takes `T k` to `T (k+1)` whatever `A, C, D` hold, then `rpt` takes `T 0` to
`T n` at cost `n (c + 64) + cRpt n`, `cRpt n = O(n + 1)`.
-/

set_option linter.unusedSimpArgs 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. -/
def rpt (A C D : Fin HS) (n : ℕ) (b : IStmt HS Pp) : IStmt HS Pp :=
  .seq (.op (.const A n)) (.seq (pro A C D) (.seq (uloop C D b) (.seq (epi C D) (.op (.drop A)))))

/-- Overwrite the three counter stacks. -/
def cf3 (T : IState HS) (A C D : Fin HS) (LA LC LD : List Item) : IState HS :=
  ((T.upd A LA).upd C LC).upd D LD

/-- Its fixed overhead. -/
def cRpt (n : ℕ) : ℕ := 3 * zsz (n : ℤ) + cUL n + n * 31 + 46

end NLogN.KFT

end

section
-- Source module: Solutions.KFTDig
/-!
# KFT §1: digits in

`digI b e M S`: with `x ≥ 0`, `x < 2^{bM}` on stack `tX`, it cuts the `M` radix-`2^b` digits
`a_j = ⌊x / 2^{bj}⌋ mod 2^b` off `x` (one `split` by the ruler `2^b` each, cost `O(b)`
independently of the size of the rest), and leaves on stack `0` the Gaussian array
`digG b e M S x = [(a_0 2^e, 0), …, (a_{M-1} 2^e, 0), (0,0), …]` of `S` entries (`gflat` form),
i.e. the scaled input coefficients `a_j 2^{p-b-2}` of §9.2.3 padded to `S` entries (`e = p-b-2`).
Cost `O(M (b + e + 1) + S + 1)`, all loop counters unary.
-/

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

namespace NLogN.KFT

open IntMul KFC KFB

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

abbrev tX : Fin HS := 56
abbrev tR : Fin HS := 57
abbrev tE : Fin HS := 58
abbrev tT : Fin HS := 59
abbrev tA : Fin HS := 60
abbrev tC : Fin HS := 61
abbrev tD : Fin HS := 62

/-! ## The denotation -/

/-- The first `M` radix-`2^b` digits of `x`. -/
def chk (b M x : ℕ) : List ℤ := (List.range M).map fun j => ((x / 2 ^ (b * j) % 2 ^ b : ℕ) : ℤ)

/-- **The input array**: scaled digits, padded with zeros to `S` entries. -/
def digG (b e M S x : ℕ) : List GI :=
  (chk b M x).map (fun a => (a * 2 ^ e, 0)) ++ List.replicate (S - M) (0, 0)

/-! ## Loop 1: cut the digits -/

def cutB : IStmt HS Pp := .op (.split tX tR tT)

/-! ## Loop 2: zeros -/

def zB : IStmt HS Pp := .op (.const 0 0)

/-! ## Loop 3: scaled digits onto stack `0` -/

def putB : IStmt HS Pp := .seq (.op (.const 0 0)) (.seq (.op (.mv tT 0)) (.op (.shl 0 tE)))

def fG (e : ℕ) (a : ℤ) : GI := (a * 2 ^ e, 0)

/-! ## The program -/

/-- **Digits in.** -/
def digI (b e M S : ℕ) : IStmt HS Pp :=
  .seq (.op (.const tR (2 ^ b))) (.seq (rpt tA tC tD M cutB) (.seq (.op (.drop tX))
    (.seq (rpt tA tC tD (2 * (S - M)) zB) (.seq (.op (.const tE (2 ^ e)))
    (.seq (rpt tA tC tD M putB) (.seq (.op (.drop tE)) (.op (.drop tR))))))))

/-- Its cost. -/
def cDig (b e M S : ℕ) : ℕ :=
  (b + 3) + (M * (splitCost b + 17 + 64) + 60 + cRpt M) + 2 +
    ((2 * (S - M)) * (2 + 64) + 60 + cRpt (2 * (S - M))) + (e + 3) +
    (M * (2 + (76 * b + 101) + shlCost e + 64) + 60 + cRpt M) + (2 * (e + 1) + 2) + (2 * (b + 1) + 2)

end NLogN.KFT

end

section
-- Source module: Solutions.KFTCarry
/-!
# KFT §2: carries out

`carI b N`: stack `0` holds the Gaussian array `gflat zs` (`N` entries, real parts `c_j ≥ 0`,
imaginary parts discarded) on top of `R0`.  Phase 1 runs the radix-`2^b` carry scan
`t := y + c_j`, digit `t mod 2^b`, carry `t / 2^b` (one `add` and one `split` by the ruler `2^b`
per entry; the carry stays below `2^B`).  Phase 2 assembles the digits from the top with `cat`
(`O(b)` each) starting from the sentinel `carry + 1`, while `2^{bN}` is built by `shl`; one final
`add` removes the sentinel.  Result on stack `0`: the integer `Σ_j c_j 2^{bj}` (`valN`), i.e.
`∑_{h<N} c_h Q^h` of `RF.carry_digits` with `Q = 2^b` (`valN_eq_sum`).
Cost `O((N + 1)(B + b + 1))` for `GB B zs`.
-/

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

namespace NLogN.KFT

open IntMul KFC KFB

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

abbrev tY : Fin HS := 56
abbrev tP : Fin HS := 58
abbrev tG : Fin HS := 59

/-! ## Pure carrying -/

/-- The carry scan: digits and final carry. -/
def carN (b : ℕ) : List ℕ → ℕ → List ℕ × ℕ
  | [], y => ([], y)
  | c :: cs, y => ((y + c) % 2 ^ b :: (carN b cs ((y + c) / 2 ^ b)).1,
      (carN b cs ((y + c) / 2 ^ b)).2)

/-- Radix-`2^b` value, least significant first. -/
def valN (b : ℕ) : List ℕ → ℕ
  | [] => 0
  | c :: cs => c + 2 ^ b * valN b cs

/-- Horner from the most significant digit. -/
def hor (b a : ℕ) (ds : List ℕ) : ℕ := ds.foldl (fun a d => a * 2 ^ b + d) a

/-! ## The denotation -/

/-- Real parts as naturals. -/
def reN (zs : List GI) : List ℕ := zs.map fun z => z.1.toNat

/-- **The output**: `Σ_j Re(z_j) 2^{bj}`. -/
def carOut (b : ℕ) (zs : List GI) : ℕ := valN b (reN zs)

/-! ## Phase 1 -/

def c1B : IStmt HS Pp := .seq (.op (.add 0 tY tY)) (.seq (.op (.split tY tR tG)) (.op (.drop 0)))

def cB1 (b B : ℕ) : ℕ := 400 * (2 * B + 1) + (splitCost b + 17) + (2 * B + 2)

/-! ## Phase 2 -/

def c2B : IStmt HS Pp := .seq (.op (.cat tY tR tG)) (.op (.shl NLogN.KFT.tP tR))

/-! ## The program -/

/-- **Carries out.** -/
def carI (b N : ℕ) : IStmt HS Pp :=
  .seq (.op (.const tR (2 ^ b))) (.seq (.op (.const tY 0)) (.seq (rpt tA tC tD N c1B)
    (.seq (.op (.const NLogN.KFT.tP 1)) (.seq (.op (.add tY NLogN.KFT.tP tY)) (.seq (.op (.const NLogN.KFT.tP 1))
    (.seq (rpt tA tC tD N c2B) (.seq (.op (.neg NLogN.KFT.tP)) (.seq (.op (.add tY NLogN.KFT.tP 0))
    (.op (.drop tR))))))))))

/-- Its cost. -/
def cCar (b B N : ℕ) : ℕ :=
  (b + 3) + 2 + (N * (NLogN.KFT.cB1 b B + 64) + 60 + cRpt N) + 3 + 400 * (B + 2) + 3 +
    (N * (catCost b + shlCost b + 64) + 60 + cRpt N) + 20 + 400 * (2 * (b * N) + B + 4) +
    (2 * (b + 1) + 2)

end NLogN.KFT

end

section
-- Source module: Solutions.KFTOut
/-!
# KFT §3: linear cost bounds, the `bin_{2n}` write, and carries-then-write

* `cCar_le`: `cCar b B N ≤ 10^5 (N + 1)(B + b + 1)`.
* `writeI_spec`: `write A W` appends `bin m v` to the output for `v < 2^m` (ruler `2^m` on `W`),
  at cost `outRCost m ≤ 200 (m + 1)`.
* `carW_spec`: `carI` followed by the write: given the coefficient array `zs` on stack `0` and
  the ruler `2^m` (`m = 2n`) on stack `W`, the output gains `bin m (carOut b zs)` (when
  `carOut b zs < 2^m`), and stack `0` and `W` lose the array and the ruler.
* `carOut_eq_sum`: `carOut b zs = ∑_{h < N} c_h (2^b)^h` (`c_h = Re z_h`), the sum of
  `RF.carry_digits`.
-/

set_option linter.unusedSimpArgs false
set_option maxRecDepth 20000

namespace NLogN.KFT

open IntMul KFC KFB

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

/-- The output step. -/
def writeI (A W : Fin HS) : IStmt HS Pp := .op (.write A W)

/-- The input digits as naturals. -/
def chkN (b M x : ℕ) : List ℕ := (List.range M).map fun j => x / 2 ^ (b * j) % 2 ^ b

/-- Carries, then the write. -/
def carW (b N : ℕ) (W : Fin HS) : IStmt HS Pp := .seq (carI b N) (writeI 0 W)

end NLogN.KFT

end

section
-- Source module: Solutions.KFURevBase
/-!
# KFU §3: labelled layouts for the axis reversal

Addresses of the working array are tracked through *layouts*: lists of atom labels (`2j` is the
main part of axis `j`, of size `mz j`; `2j+1` is its extra bit, of size `2`, present when
`bg j`).  An axis of size `tz j = mz j · (2 if bg j)` is the pair `[mz j, 2]` in row-major order.

* `Dl ax0 lay xs`: the array `xs` (shape `ax0.map tz`) viewed in the layout `lay`.
* `POp g σ n`: an address map `g = permIdx σ` on `n` digits, `σ` onto `[0, n)`.
* `Dl_step`: a program realizing `pull (g lay').map asz → lay'.map asz` by `g` turns
  `Dl ax0 (g lay')` into `Dl ax0 lay'`.  `swapPos` and `KIL.swapMid` are `POp`s.
* `Dl_layA`: for the standard layout `layA ax` of an axis order `ax`, `Dl` is an ordinary pull
  onto `ax.map tz`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN.KFU

open AModel AModelL12

/-! ## Position permutations -/

/-- `g` is the digit permutation `permIdx σ` on `n`-digit lists, `σ` maps `[0,n)` onto itself. -/
structure POp (g : List ℕ → List ℕ) (σ : ℕ → ℕ) (n : ℕ) : Prop where
  eq : ∀ b : List ℕ, b.length = n → g b = permIdx σ b
  lt : ∀ i < n, σ i < n
  surj : ∀ j < n, ∃ i < n, σ i = j

/-! ### `swapPos` -/

/-! ### `swapMid` -/

/-- Index map of `KIL.swapMid nA nF nG` (input layout `A G F C`, output `A F G C`). -/
def mvIdx (nA nF nG i : ℕ) : ℕ :=
  if i < nA then i else if i < nA + nF then i + nG else if i < nA + nF + nG then i - nF else i

/-! ## Layouts -/

section Layout

variable (mz : ℕ → ℕ) (bg : ℕ → Bool)

/-- Size of an atom. -/
def asz (l : ℕ) : ℕ := if l % 2 = 0 then mz (l / 2) else 2

/-- Size of an axis. -/
def tz (j : ℕ) : ℕ := if bg j then mz j * 2 else mz j

/-- Merged digit of axis `j` read from the address `b` in the layout `lay`. -/
def dig (lay b : List ℕ) (j : ℕ) : ℕ :=
  if bg j then b.getD (lay.idxOf (2 * j)) 0 * 2 + b.getD (lay.idxOf (2 * j + 1)) 0
  else b.getD (lay.idxOf (2 * j)) 0

/-- The source address (in the original axis order `ax0`). -/
def srcF (ax0 lay b : List ℕ) : List ℕ := ax0.map (dig bg lay b)

/-- Every atom of every axis occurs in the layout. -/
def Full (ax0 lay : List ℕ) : Prop := ∀ j ∈ ax0, 2 * j ∈ lay ∧ (bg j = true → 2 * j + 1 ∈ lay)

/-- The array `xs` (shape `ax0.map tz`) in the layout `lay`. -/
def Dl {α : Type*} (ax0 lay : List ℕ) (xs : List α) (d : α) : List α :=
  AModel.pull (ax0.map (tz mz bg)) (lay.map (asz mz)) (srcF bg ax0 lay) d xs

/-- The atoms of one axis. -/
def atm (j : ℕ) : List ℕ := 2 * j :: if bg j then [2 * j + 1] else []

/-- The standard layout of an axis order. -/
def layA (L : List ℕ) : List ℕ := L.flatMap (atm bg)

end Layout

/-! ## The standard layout -/

section Std

variable (mz : ℕ → ℕ) (bg : ℕ → Bool)

/-- The reversal of the first `n` digits. -/
def revH (n : ℕ) (m : List ℕ) : List ℕ := (m.take n).reverse ++ m.drop n

end Std

end NLogN.KFU

end

section
-- Source module: Solutions.KFUPad
/-!
# KFU §1: padding and deletion scans on integer arrays (`w`-bit entries)

* `padI X B Z`: stack `0` holds `X` blocks of `B` entries (`xs`, `|xs| = X·B`); afterwards it
  holds `padPre B Z xs X` (every block followed by `Z` zeros).  Per block: `B` moves to the
  buffer `zO`, `Z` zero pushes; finally the whole buffer goes back (`mvW`, the two reversals
  cancel).  `padPre (F·Y) ((F'-F)·Y) xs X = KIL.padL X F F' Y 0 xs` by definition.
* `unpadI X B C`: the deletion scan (blocks of stride `C`, keep the first `B`).
  `unpPre (F·Y) (F'·Y) ys X = KIL.unpadL X F F' Y ys`.
* `padAllC S T n` / `unpadAllC S T n`: one scan per field; they compute `embed S T 0` and
  `extract S T 0` (`padAll_spec`, `unpadAll_spec`).

Costs `≤ 10^4 (w+1)(V + X + 1)` per scan (`V` the padded volume); all counters unary.
-/

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

namespace NLogN.KFU

open IntMul KFC KFT AModel AModelL12

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

/-- Inner loop counters and the buffer. -/
abbrev zA : Fin HS := 57
abbrev zC : Fin HS := 58
abbrev zD : Fin HS := 59
abbrev zO : Fin HS := 63
/-- Outer loop counters. -/
abbrev oA : Fin HS := 60
abbrev oC : Fin HS := 61
abbrev oD : Fin HS := 62

/-! ## Counted inner loops -/

/-! ## The padding scan -/

/-- The padded prefix after `k` blocks. -/
def padPre (B Z : ℕ) (xs : List ℤ) (k : ℕ) : List ℤ :=
  (List.range k).flatMap fun x => (xs.drop (x * B)).take B ++ List.replicate Z 0

def padBlk (B Z : ℕ) : IStmt HS Pp :=
  .seq (rpt NLogN.KFU.zA zC zD B (.op (.mv 0 zO))) (rpt NLogN.KFU.zA zC zD Z (.op (.const zO 0)))

def cBlk (B Z w : ℕ) : ℕ :=
  (B * (76 * w + 101 + 64) + 60 + cRpt B) + (Z * (2 + 64) + 60 + cRpt Z)

/-- **The padding scan.** -/
def padI (X B Z : ℕ) : IStmt HS Pp :=
  .seq (.op (.mark zO)) (.seq (rpt oA oC oD X (padBlk B Z)) (.seq (mvW zO 0) (.op (.unmark zO))))

def cPadI (X B Z w : ℕ) : ℕ :=
  1 + ((X * (NLogN.KFU.cBlk B Z w + 64) + 60 + cRpt X) + (X * (B + Z) * kap w + 30 + 1))

/-! ## The deletion scan -/

/-- The kept prefix after `k` blocks of stride `C`. -/
def unpPre (B C : ℕ) (ys : List ℤ) (k : ℕ) : List ℤ :=
  (List.range k).flatMap fun x => (ys.drop (x * C)).take B

def unpBlk (B C : ℕ) : IStmt HS Pp :=
  .seq (rpt NLogN.KFU.zA zC zD B (.op (.mv 0 zO))) (rpt NLogN.KFU.zA zC zD (C - B) (.op (.drop 0)))

def cUBlk (B C w : ℕ) : ℕ :=
  (B * (76 * w + 101 + 64) + 60 + cRpt B) + ((C - B) * (2 * w + 2 + 64) + 60 + cRpt (C - B))

/-- **The deletion scan.** -/
def unpadI (X B C : ℕ) : IStmt HS Pp :=
  .seq (.op (.mark zO)) (.seq (rpt oA oC oD X (unpBlk B C)) (.seq (mvW zO 0) (.op (.unmark zO))))

def cUnpadI (X B C w : ℕ) : ℕ :=
  1 + ((X * (cUBlk B C w + 64) + 60 + cRpt X) + (X * B * kap w + 30 + 1))

end NLogN.KFU

end

section
-- Source module: Solutions.KFUPadAll
/-!
# KFU §2: padding / deletion of every field

`padAllC S T n` runs one padding scan per field `k < n` (`S ≤ T` pointwise); after all
`|S|` scans stack `0` holds `embed S T 0 xs` (`padAll_spec`).  `unpadAllC S T n` runs the deletion
scans and computes `extract S T 0 ys` (`unpadAll_spec`).  Cost `≤ |S| · 10^4 (w+1)(2 vol T + 1)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC KFT AModel AModelL12

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

/-! ## Shape facts -/

/-! ## All padding scans -/

/-- The `k`-th padding scan of `S ↦ T`. -/
def padStep (S T : Shape) (k : ℕ) : IStmt HS Pp :=
  padI (AModel.vol (T.take k)) (S.getD k 0 * AModel.vol (S.drop (k + 1)))
    ((T.getD k 0 - S.getD k 0) * AModel.vol (S.drop (k + 1)))

def padAllC (S T : Shape) : ℕ → IStmt HS Pp
  | 0 => .skip
  | k + 1 => .seq (padAllC S T k) (padStep S T k)

/-- The per-scan cost unit. -/
def cScan (VT w : ℕ) : ℕ := 10 ^ 4 * (w + 1) * (2 * VT + 1) + 1

/-! ## All deletion scans -/

def unpadStep (S T : Shape) (k : ℕ) : IStmt HS Pp :=
  unpadI (AModel.vol (S.take k)) (S.getD k 0 * AModel.vol (T.drop (k + 1))) (T.getD k 0 * AModel.vol (T.drop (k + 1)))

def unpadAllC (S T : Shape) : ℕ → IStmt HS Pp
  | 0 => .skip
  | k + 1 => .seq (unpadAllC S T k) (unpadStep S T k)

end NLogN.KFU

end

section
-- Source module: Solutions.KFURev
/-!
# KFU §4: one-bit moves on layouts, and the two bit phases of the axis reversal

* `mvFC a f g c` / `mvBC a g f c`: push the four block volumes on stacks `1 … 4`, run KFC
  `fwdF` / `bwdF` (Lemma 3: a block `F` past `G`, cost linear in `vol F` resp. `vol G`), pop.
* `mvF_lay` / `mvB_lay`: the same on `Dl` layouts: `LA ++ LF ++ LG ++ LC ↦ LA ++ LG ++ LF ++ LC`.
* Phase 1 `p1C pre L Z`: every extra bit of the axes `L` moves (one `fwdF` with `vol F = 2`)
  behind the main atoms: `pre ++ layA L ++ Z ↦ pre ++ mains L ++ bitsR L ++ Z`.
* Phase 3 `p3C pre L Z`: the inverse placement with `bwdF` (`vol G = 2`):
  `pre ++ mains L ++ bitsF L ++ Z ↦ pre ++ layA L ++ Z`.
Each move costs `≤ cMv V w = 100 · kap w · (V + 1)`.
-/

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

namespace NLogN.KFU

open IntMul KFC AModel AModelL12
open List (Perm)

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

/-! ## The two move programs -/

def mvFC (a f g c : ℕ) : IStmt HS Pp :=
  .seq (.op (.const 1 a)) (.seq (.op (.const 2 f)) (.seq (.op (.const 3 g)) (.seq (.op (.const 4 c))
    (.seq fwdF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.seq (.op (.drop 3)) (.op (.drop 4)))))))))

def mvBC (a g f c : ℕ) : IStmt HS Pp :=
  .seq (.op (.const 1 a)) (.seq (.op (.const 2 g)) (.seq (.op (.const 3 f)) (.seq (.op (.const 4 c))
    (.seq bwdF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.seq (.op (.drop 3)) (.op (.drop 4)))))))))

/-- Parameter pushes and pops. -/
def cK4 (a b c d : ℕ) : ℕ :=
  3 * (zsz (a : ℤ) + zsz (b : ℤ) + zsz (c : ℤ) + zsz (d : ℤ)) + 16

/-- Cost of one bit move. -/
def cMv (V w : ℕ) : ℕ := 100 * kap w * (V + 1)

/-! ## Layout facts -/

section Lay

variable (mz : ℕ → ℕ) (bg : ℕ → Bool)

def mains (L : List ℕ) : List ℕ := L.map (2 * ·)
def bitsF (L : List ℕ) : List ℕ := (L.filter bg).map (2 * · + 1)
def bitsR (L : List ℕ) : List ℕ := bitsF bg L.reverse

end Lay

/-! ## Moves on layouts -/

end NLogN.KFU

end

section
-- Source module: Solutions.KFURevPh
/-!
# KFU §5: the bit phases of the reversal

Fixed: a reference layout `lay0` (`Nodup`, `Full`), every layout below is a permutation of it,
`V = vol (lay0.map asz)`.

* `p1_spec`: `p1C pre L Z` takes `pre ++ layA L ++ Z` to `pre ++ mains L ++ bitsR L ++ Z`.
* `p3_spec`: `p3C pre L Z` takes `pre ++ mains L ++ bitsF L ++ Z` to `pre ++ layA L ++ Z`.
Cost `≤ (|L| + 1)(cMv V w + 1)` each.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

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

section Ph

variable (mz : ℕ → ℕ) (bg : ℕ → Bool)

/-- Phase 1: the extra bits of `L` move behind the main atoms (innermost axis first). -/
def p1C : List ℕ → List ℕ → List ℕ → IStmt HS Pp
  | _, [], _ => .skip
  | pre, j :: L, Z => .seq (p1C (pre ++ atm bg j) L Z)
      (if bg j = true then mvFC (AModel.vol ((pre ++ [2 * j]).map (asz mz))) (AModel.vol ([2 * j + 1].map (asz mz)))
        (AModel.vol ((mains L ++ bitsR bg L).map (asz mz))) (AModel.vol (Z.map (asz mz))) else .skip)

/-- Phase 3: the extra bits of `L` move back in front, after their main atoms. -/
def p3C : List ℕ → List ℕ → List ℕ → IStmt HS Pp
  | _, [], _ => .skip
  | pre, j :: L, Z => .seq
      (if bg j = true then mvBC (AModel.vol ((pre ++ [2 * j]).map (asz mz))) (AModel.vol ([2 * j + 1].map (asz mz)))
        (AModel.vol ((mains L).map (asz mz))) (AModel.vol ((bitsF bg L ++ Z).map (asz mz))) else .skip)
      (p3C (pre ++ atm bg j) L Z)

end Ph

end NLogN.KFU

end

section
-- Source module: Solutions.KFURevTop
/-!
# KFU §6: the axis reversal

`revC q' w u mz bg L cid` reverses the axis order `L` of an integer array of shape
`(L ++ [cid]).map tz` (the last axis `cid`, e.g. the Gaussian component, stays):

1. phase 1 (`p1C`): every extra bit of a big axis (`tz j = 2 · 2^u`) moves behind the main atoms
   (KFC `fwdF`, `O(V)` each);
2. phase 2 (`swC`): the `|L|` main atoms (all of size `2^u`) are reversed by `⌊|L|/2⌋` swaps
   `swapPos k (|L|-1-k)` (KFS `swAx`, Lemma 12, `K (w+1)(V+1) u^τ` each);
3. phase 3 (`p3C`): the extra bits move back behind their main atoms (KFC `bwdF`).

`rev_spec`: the result is `pull X Y (revH |L|) 0 xs`, `X = (L ++ [cid]).map tz`,
`Y = (L.reverse ++ [cid]).map tz`, at cost `≤ 3 (|L|+1)(cMv V w + 1 + K (w+1)(V+1) u^τ)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

/-! ## Reversal of a prefix by swaps -/

def rIdx (r k i : ℕ) : ℕ := if i < k ∨ (r - k ≤ i ∧ i < r) then r - 1 - i else i

/-- The list after the first `k` swaps `(i, r-1-i)`. -/
def rv (r k : ℕ) (l : List ℕ) : List ℕ := permIdx (rIdx r k) l

/-! ## Phase 2 -/

/-- What `KFS.swAx_spec` provides, for a fixed prime `q'` and constant `K`. -/
def SwOK (q' : ℕ) (K : ℝ) : Prop :=
  ∀ (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 →
    ∀ (I : List Sym) (t : IState HS) (xs : List ℤ) (R : List Item),
    t.G 0 = ints xs ++ R → xs.length = vol sh → (∀ x ∈ xs, zsz x ≤ w) →
    ∃ c : ℕ, IRuns I KFL.l12prog (KFS.swAx q' w sh p q u) t
      (t.upd 0 (ints (AModel.pull sh sh (swapPos p q) 0 xs) ++ R)) c ∧
      (c : ℝ) ≤ K * ((w : ℝ) + 1) * ((vol sh : ℝ) + 1) * (u : ℝ) ^ KFS.τS

def swC (q' w u : ℕ) (sh : Shape) (r : ℕ) : ℕ → IStmt HS 1
  | 0 => .skip
  | k + 1 => .seq (swC q' w u sh r k) (KFS.swAx q' w sh k (r - 1 - k) u)

/-! ## The reversal -/

def revC (q' w u : ℕ) (mz : ℕ → ℕ) (bg : ℕ → Bool) (L : List ℕ) (cid : ℕ) : IStmt HS 1 :=
  .seq (p1C mz bg [] L [2 * cid])
    (.seq (swC q' w u ((mains L ++ bitsR bg L ++ [2 * cid]).map (asz mz)) L.length (L.length / 2))
      (p3C mz bg [] L.reverse [2 * cid]))

end NLogN.KFU

end

section
-- Source module: Solutions.KFURot
/-!
# KFU §7: one prefix-controlled cyclic shift with a linear offset table

`rotStep Q P B ac`: generate the offset table `tbl Q ac` (KIL2 `genP`, register code), push
`P, Q, B`, run KFC `rotF`, pop.  For the shape `A ++ Q :: C` and coefficients `cs` (one per
field of `A`), `rotStep_spec` computes the full-shape pull by

  `rotG |A| Q f b = b.take |A| ++ ((b_{|A|} + Q - f (b.take |A|)) mod Q) :: b.drop (|A|+1)`,
  `f a = (Σ_k cs_k a_k) mod Q`,

i.e. the entry at address `a` moves to the address whose `|A|`-th digit is
`(a_{|A|} + Σ_k cs_k a_k) mod Q` (the controls `a_k`, `k < |A|`, are the outer fields).
Cost `cRotF (vol A) Q (vol C) w + 4·10^5 (|A|+1) vol A (Bb+2) + O(Bb)`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

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

/-- The full-shape address map of a controlled shift. -/
def rotG (n Q : ℕ) (f : List ℕ → ℕ) (b : List ℕ) : List ℕ :=
  b.take n ++ ((b.getD n 0 + Q - f (b.take n)) % Q) :: b.drop (n + 1)

/-- The controlled shift: table, parameters, `rotF`, pops. -/
def rotStep (Q P B : ℕ) (ac : List (ℕ × ℕ)) : IStmt HS Pp :=
  .seq (KIL.genP Q ac).code (.seq (.op (.const 1 P)) (.seq (.op (.const 2 Q)) (.seq (.op (.const 4 B))
    (.seq rotF (.seq (.op (.drop 1)) (.seq (.op (.drop 2)) (.op (.drop 4))))))))

/-- Its cost. -/
def cRot (P Q B w Bb n : ℕ) : ℕ :=
  400000 * (n + 1) * P * (Bb + 2) + NLogN.KFC.cRotF P Q B w + 3 * (zsz (P : ℤ) + zsz (Q : ℤ) + zsz (B : ℤ)) + 12

end NLogN.KFU

end

section
-- Source module: Solutions.ResampleCRT
/-!
# §9.3 An explicit cyclic-convolution layout

* `Phi`: `ℤ/S → ∏ ℤ/s_i`, `k ↦ (μ_i k)_i` (the coordinate map of `x ↦ ∏ x_i^{μ_i}`), an additive
  equivalence for pairwise coprime `s_i` and units `μ_i` (`PhiEquiv`); every additive
  equivalence carries cyclic convolution to cyclic convolution (`conv_addEquiv`), so
  `Φ(u ∗ v) = Φu ∗ Φv` (`Phi_conv`).
* Mixed radix: `k mod P_i = Σ_{j<i} a_j P_j` (`mod_prefix_eq_sum`) and the coordinate formula
  (42): `μ_i k ≡ a_i + μ_i Σ_{j<i} a_j P_j (mod s_i)` (`coord_formula`).
* Triangular realization: the prefix-controlled cyclic shifts, processed for `i = d-1, …, 0`,
  send every valid address `a` to `b` of (42) (`shifts_realize`); each shift is a bijection of
  the padded box preserving the valid set `V` and its complement (`upd_inv`, `upd_valid`,
  `upd_invalid`).
-/

open Finset

namespace IntMul.OpenAI.RCRT

/-! ### Convolution is preserved by additive equivalences -/

section conv
variable {G H : Type*} [AddCommGroup G] [AddCommGroup H] [Fintype G] [Fintype H]
  {R : Type*} [CommRing R]

end conv

/-! ### The map `Φ` -/

variable {d : ℕ}

/-- `Φ(k)_i = μ_i k mod s_i`. -/
def Phi (s : Fin d → ℕ) (μ : Fin d → ℕ) (k : ZMod (∏ i, s i)) : (i : Fin d) → ZMod (s i) :=
  fun i => (μ i : ZMod (s i)) * ZMod.castHom (Finset.dvd_prod_of_mem s (Finset.mem_univ i))
    (ZMod (s i)) k

/-- For pairwise coprime `s_i` and units `μ_i`, `Φ` is an additive equivalence. -/
noncomputable def PhiEquiv (s : Fin d → ℕ) [∀ i, NeZero (s i)]
    (hcop : Pairwise (Function.onFun Nat.Coprime s)) (μ : Fin d → ℕ)
    (hμ : ∀ i, Nat.Coprime (μ i) (s i)) : ZMod (∏ i, s i) ≃+ ((i : Fin d) → ZMod (s i)) :=
  (ZMod.prodEquivPi s hcop).toAddEquiv.trans
    (AddEquiv.piCongrRight (fun i =>
      { toFun := fun x => (μ i : ZMod (s i)) * x
        invFun := fun x => (μ i : ZMod (s i))⁻¹ * x
        left_inv := fun x => by
          simp only [← mul_assoc]
          rw [ZMod.inv_mul_of_unit _ ((ZMod.isUnit_iff_coprime _ _).2 (hμ i)), one_mul]
        right_inv := fun x => by
          simp only [← mul_assoc]
          rw [ZMod.coe_mul_inv_eq_one _ (hμ i), one_mul]
        map_add' := fun x y => mul_add _ _ _ }))

/-! ### Mixed radix and the coordinate formula (42) -/

/-- Prefix product `P_i = ∏_{j<i} s_j`. -/
def Ppre (s : ℕ → ℕ) (i : ℕ) : ℕ := ∏ j ∈ range i, s j

/-- Mixed-radix digit `a_j = ⌊k / P_j⌋ mod s_j`. -/
def digit (s : ℕ → ℕ) (k j : ℕ) : ℕ := k / Ppre s j % s j

/-! ### Realization by prefix-controlled cyclic shifts -/

section shifts
variable (s μ : ℕ → ℕ)

/-- Valid controls for target `i`: every coordinate `j ≤ i` is in `[0, s_j)`. -/
def validUpTo (x : ℕ → ℕ) (i : ℕ) : Prop := ∀ j ≤ i, x j < s j

instance (x : ℕ → ℕ) (i : ℕ) : Decidable (validUpTo s x i) :=
  inferInstanceAs (Decidable (∀ j ≤ i, x j < s j))

/-- Offset from the lower coordinates: `μ_i Σ_{j<i} x_j P_j`. -/
def offset (x : ℕ → ℕ) (i : ℕ) : ℕ := μ i * ∑ j ∈ range i, x j * Ppre s j

/-- The prefix-controlled cyclic shift of target `i`. -/
def upd (i : ℕ) (x : ℕ → ℕ) : ℕ → ℕ :=
  if validUpTo s x i then Function.update x i ((x i + IntMul.OpenAI.RCRT.offset s μ x i) % s i) else x

/-- Its inverse (subtract the same offset). -/
def updInv (i : ℕ) (x : ℕ → ℕ) : ℕ → ℕ :=
  if validUpTo s x i then
    Function.update x i ((x i + (s i - IntMul.OpenAI.RCRT.offset s μ x i % s i)) % s i) else x

end shifts

end IntMul.OpenAI.RCRT

end

section
-- Source module: Solutions.KFUPhi
/-!
# KFU §8: the programs `Φ` and `Φ⁻¹`

Parameters: `d` axes, sources `s_j = sF j` (`1 ≤ s_j ≤ t_j`), padded sizes
`t_j = tP d h bL j = 2^h · (2 if bL j)`, multipliers `μ j`.  Arrays are integer lists with a
trailing component axis of size `2` (Gaussian records in `gflat` form); the component axis has
index `d`.

* `shK f = ((range d).reverse ++ [d]).map f` (k-order, `x_0` innermost: flat index `2k + c`),
* `shR f = (range d ++ [d]).map f` (reversed order, `x_0` outermost).

`phiC`: pad `S_k → T_k` (one scan per axis), reverse the axes (`revC`: Lemma 12 swaps of the
`2^h`-atoms + single-bit moves), delete the padding `T_r → S_r`, run the `d` prefix-controlled
shifts `i = d-1, …, 0` (`rotStep`: offset table `Σ_{j<i} (μ_i P_j mod s_i) x_j mod s_i`), pad
`S_r → T_r`.  `phi_spec`: the result is `PhiD xs`.

`phiInvC`: the inverse (delete padding, inverse shifts `i = 0, …, d-1`, pad, reverse back,
delete padding).  `phiInv_spec`: the result is `PhiInvD ys`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

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

/-! ## Generic shape lemmas -/

/-! ## Parameters and shapes -/

section Par

variable (d h : ℕ) (bL : ℕ → Bool) (sF : ℕ → ℕ)

def mzP (j : ℕ) : ℕ := if j < d then 2 ^ h else 2
def bgP (j : ℕ) : Bool := decide (j < d) && bL j
/-- The padded sizes `t_j` (`2` for the component axis `d`). -/
def tP (j : ℕ) : ℕ := tz (mzP d h) (bgP d bL) j
/-- The source sizes `s_j` (`2` for the component axis `d`). -/
def sP (j : ℕ) : ℕ := if j < d then sF j else 2
/-- Reversed order (`x_0` outermost) and k-order (`x_0` innermost). -/
def shR (f : ℕ → ℕ) : Shape := (List.range d ++ [d]).map f
def shK (f : ℕ → ℕ) : Shape := ((List.range d).reverse ++ [d]).map f

end Par

/-! ## The shifts -/

/-- Coefficients of the forward shift of target `i`: `μ_i P_j mod s_i`. -/
def csF (sF μ : ℕ → ℕ) (i : ℕ) : List ℕ :=
  (List.range i).map fun j => μ i * IntMul.OpenAI.RCRT.Ppre sF j % sF i

/-- Coefficients of the inverse shift: `-(μ_i P_j) mod s_i`. -/
def csB (sF μ : ℕ → ℕ) (i : ℕ) : List ℕ :=
  (List.range i).map fun j => (sF i - μ i * IntMul.OpenAI.RCRT.Ppre sF j % sF i) % sF i

/-- One shift program. -/
def rotI (d : ℕ) (sF : ℕ → ℕ) (cs : ℕ → List ℕ) (i : ℕ) : IStmt HS Pp :=
  NLogN.KFU.rotStep (sF i) (vol ((List.range i).map sF))
    (vol (((List.range d).drop (i + 1) ++ [d]).map (sP d sF))) (((List.range i).map sF).zip (cs i))

/-- Its address map. -/
def gRot (sF : ℕ → ℕ) (cs : ℕ → List ℕ) (i : ℕ) (b : List ℕ) : List ℕ :=
  NLogN.KFU.rotG i (sF i) (fun a => KIL.lf (cs i) a % sF i) b

/-- Targets `k-1, …, 0` (forward). -/
def rotsC (d : ℕ) (sF : ℕ → ℕ) (cs : ℕ → List ℕ) : ℕ → IStmt HS Pp
  | 0 => .skip
  | k + 1 => .seq (rotI d sF cs k) (rotsC d sF cs k)

/-- Targets `0, …, k-1` (backward). -/
def rotsBC (d : ℕ) (sF : ℕ → ℕ) (cs : ℕ → List ℕ) : ℕ → IStmt HS Pp
  | 0 => .skip
  | k + 1 => .seq (rotsBC d sF cs k) (rotI d sF cs k)

def Gf (sF : ℕ → ℕ) (cs : ℕ → List ℕ) : ℕ → List ℕ → List ℕ
  | 0 => id
  | k + 1 => gRot sF cs k ∘ Gf sF cs k

def Gb (sF : ℕ → ℕ) (cs : ℕ → List ℕ) : ℕ → List ℕ → List ℕ
  | 0 => id
  | k + 1 => Gb sF cs k ∘ gRot sF cs k

/-- Per-shift cost bound. -/
def cRb (d w Vs PM Bb : ℕ) : ℕ := 100 * kap w * Vs + 400000 * d * PM * (Bb + 2) + 1

end NLogN.KFU

end

section
-- Source module: Solutions.KFUPhiRun
/-!
# KFU §9: the runs of `Φ` and `Φ⁻¹`

`phiC` (stack `0`: `ints xs ++ R`, `xs` on `S_k = shK (sP)`):
1. pad `S_k → T_k` (`d+1` scans),
2. reverse the axes on `T_k → T_r` (`revC`, Lemma 12 swaps + bit moves),
3. delete the padding `T_r → S_r`,
4. the shifts `i = d-1, …, 0` (forward coefficients `csF`),
5. pad `S_r → T_r`.

Result `PhiD xs = embed S_r T_r (pull S_r S_r (Gf d) (pull S_k S_r (revH d) xs))` (steps 1–3
collapse by `extract_pull_embed'`).  `phiInvC` runs the inverse chain; result
`PhiInvD ys = pull S_r S_k (revH d) (pull S_r S_r (Gb d) (extract S_r T_r ys))`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

section Defs

variable (q' w h d : ℕ) (bL : ℕ → Bool) (sF μ : ℕ → ℕ)

/-- The program `Φ`. -/
def phiC : IStmt HS 1 :=
  .seq (padAllC (shK d (sP d sF)) (shK d (NLogN.KFU.tP d h bL)) (d + 1))
    (.seq (revC q' w h (mzP d h) (bgP d bL) (List.range d).reverse d)
      (.seq (unpadAllC (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) (d + 1))
        (.seq (rotsC d sF (csF sF μ) d) (padAllC (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) (d + 1)))))

/-- The program `Φ⁻¹`. -/
def phiInvC : IStmt HS 1 :=
  .seq (unpadAllC (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) (d + 1))
    (.seq (rotsBC d sF (csB sF μ) d)
      (.seq (padAllC (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) (d + 1))
        (.seq (revC q' w h (mzP d h) (bgP d bL) (List.range d) d)
          (unpadAllC (shK d (sP d sF)) (shK d (NLogN.KFU.tP d h bL)) (d + 1)))))

/-- Denotation of `Φ`. -/
noncomputable def PhiD (xs : List ℤ) : List ℤ :=
  AModelL12.embed (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) 0
    (AModel.pull (shR d (sP d sF)) (shR d (sP d sF)) (Gf sF (csF sF μ) d) 0
      (AModel.pull (shK d (sP d sF)) (shR d (sP d sF)) (revH d) 0 xs))

/-- Denotation of `Φ⁻¹`. -/
def PhiInvD (ys : List ℤ) : List ℤ :=
  AModel.pull (shR d (sP d sF)) (shK d (sP d sF)) (revH d) 0
    (AModel.pull (shR d (sP d sF)) (shR d (sP d sF)) (Gb sF (csB sF μ) d) 0
      (extract (shR d (sP d sF)) (shR d (NLogN.KFU.tP d h bL)) 0 ys))

end Defs

/-- The cost bound of `Φ` and `Φ⁻¹` (`VT` padded volume, `VS` source volume,
`PM = ∏_{j<d-1} s_j`). -/
noncomputable def cPhi (d w VT VS PM Bb : ℕ) (K : ℝ) (h : ℕ) : ℝ :=
  3 * (((d + 1) * cScan VT w + 1 : ℕ) : ℝ) +
    3 * ((d : ℝ) + 1) * (NLogN.KFU.cMv VT w + 1 + K * ((w : ℝ) + 1) * ((VT : ℝ) + 1) * (h : ℝ) ^ KFS.τS) +
    ((d * cRb d w VS PM Bb + 1 : ℕ) : ℝ)

end NLogN.KFU

end

section
-- Source module: Solutions.KFUPhiMath
/-!
# KFU §10: what `Φ` and `Φ⁻¹` compute (against `RCRT.upd`)

For a valid address `b` of `S_r` (`shR (sP)`), viewed as the function `fnL b j = b_j`:

* `fn_gRotF`: the forward step `gRot (csF) i` is `RCRT.updInv i` (subtract `μ_i Σ_{j<i} b_j P_j`),
* `fn_gRotB`: the inverse step `gRot (csB) i` is `RCRT.upd i`,

so `Gf d` undoes the `RCRT` shift chain `foldr upd (range d)` and `Gb d` is that chain.  With
`(42)` (`RCRT.coord_formula`) the shift chain sends the digit address of `k < P_d` to
`Φ(k) = (μ_j k mod s_j)_j`.  End results (`phiAddr k c = [μ_0 k mod s_0, …, μ_{d-1} k mod s_{d-1}, c]`):

* `PhiD_at`: `(PhiD xs)[enc_{T_r} (phiAddr k c)] = xs[2k + c]`,
* `PhiD_zero`: `PhiD xs` is `0` at every `T_r`-address that is not an `S_r`-address,
* `PhiInvD_at`: `(PhiInvD ys)[2k + c] = ys[enc_{T_r} (phiAddr k c)]`.

Hypotheses: `0 < s_j` and `μ_j P_j ≡ 1 (mod s_j)` for `j < d`.
-/

set_option linter.unusedSimpArgs false
set_option linter.unnecessarySeqFocus false

namespace NLogN.KFU

open IntMul KFC AModel AModelL12

/-- An address as a function. -/
def fnL (b : List ℕ) : ℕ → ℕ := fun j => b.getD j 0

section Shift

variable {d : ℕ} {sF μ : ℕ → ℕ}

/-- `updInv` applied for targets `0, …, k-1`. -/
def updInvs (sF μ : ℕ → ℕ) : ℕ → (ℕ → ℕ) → (ℕ → ℕ)
  | 0, x => x
  | k + 1, x => OpenAI.RCRT.updInv sF μ k (updInvs sF μ k x)

end Shift

/-! ## Digit addresses and `Φ`-addresses -/

section Addr

variable (d : ℕ) (sF μ : ℕ → ℕ)

/-- The digit address of `2k + c` in reversed order: `[a_0, …, a_{d-1}, c]`. -/
def digAddr (k c : ℕ) : List ℕ := (List.range d).map (OpenAI.RCRT.digit sF k) ++ [c]

/-- The `Φ`-address: `[μ_0 k mod s_0, …, μ_{d-1} k mod s_{d-1}, c]`. -/
def phiAddr (k c : ℕ) : List ℕ := (List.range d).map (fun j => μ j * k % sF j) ++ [c]

end Addr

/-! ## Encoding the k-order address -/

/-! ## Pointwise results -/

end NLogN.KFU

end