import RoleFlow.Flow.Basic
import RoleFlow.Link
import RoleFlow.Infer
import Mathlib.Tactic.Ring

/-! # The trained models, as programs of the flow signature

The models of SemiInv2, with the same stages, weights and inits. A stage is written with a builder (`B`): an operation applied to
positions returns the positions of its results, so a stage is a function from the positions
of its weights and inputs to the positions it returns. Sorts are compared by value when the
program is checked (`Prog.WF`), so the builder needs no typing of its own.

A model is a list of stages over one weight context, the stages' weights side by side
(`Stage.at`). Its program (`modelProg`) runs the stages' bodies one after another in one
builder, each on what the previous one passes on, and is typed by inference as a whole; or
(`callProg`) it calls each stage, defined on its own weights in a tower of definitions
(`Linked`). A stage may factor out its first `fac` outputs: they are returned by the model, and
in `stageProgs` every later stage takes them as extra (unused) weights. Nothing here is trusted:
the printer (`RoleFlow.Flow.Print`) checks what it prints. -/

namespace RoleFlow.Flow

open FS

/-! ## The builder -/

/-- The builder's state: the next free position and the lines so far. -/
structure BSt (Sg : Sig FS) where
  n : ℕ
  lines : Array (Line Sg) := #[]

/-- A builder of lines of `Sg`. -/
abbrev BG (Sg : Sig FS) := StateM (BSt Sg)

/-- A builder of lines of the flow signature. -/
abbrev B := BG flowSig

namespace B

/-- `ys = o(xs)`, returning the positions `ys`. -/
def line {Sg : Sig FS} {as bs : List FS} (o : Sg.Op as bs) (xs : List ℕ) : BG Sg (List ℕ) :=
  modifyGet fun s => (List.range' s.n bs.length,
    { n := s.n + bs.length, lines := s.lines.push ⟨o, xs⟩ })

/-- `ys = o(xs)` for an operation of the flow signature. -/
def op {as bs : List FS} (o : FOp as bs) (xs : List ℕ) : B (List ℕ) := line (Sg := flowSig) o xs

def o0 {b : FS} (o : FOp [] [b]) : B ℕ := do return (← op o []).headD 0
def o1 {a b : FS} (o : FOp [a] [b]) (x : ℕ) : B ℕ := do return (← op o [x]).headD 0
def o2 {a₁ a₂ b : FS} (o : FOp [a₁, a₂] [b]) (x y : ℕ) : B ℕ := do return (← op o [x, y]).headD 0
def o3 {a₁ a₂ a₃ b : FS} (o : FOp [a₁, a₂, a₃] [b]) (x y z : ℕ) : B ℕ := do
  return (← op o [x, y, z]).headD 0
def o12 {a b₁ b₂ : FS} (o : FOp [a] [b₁, b₂]) (x : ℕ) : B (ℕ × ℕ) := do
  let ys ← op o [x]; return (ys.getD 0 0, ys.getD 1 0)

/-- An actnorm on `vec n`: `x ↦ x ⊙ exp s + b`. -/
def actnorm (n : ℕ) (x s b : ℕ) : B ℕ := do o2 (.add n) (← o2 (.scaleexp n) x s) b

/-- The affine coupling step on `vec n`: `x ↦ x ⊙ exp (tanh s) + t`. -/
def affine (n : ℕ) (x s t : ℕ) : B ℕ := do
  o2 (.add n) (← o2 (.scaleexp n) x (← o1 (.tanh n) s)) t

end B

/-- How a weight starts: `he:n` is He-normal with fan-in `n`, `zero:n` is zero, `small:n` is
normal with standard deviation `0.05` (`n` is the fan-in used when tests randomise it; `0`
means a bias). -/
def he (n : ℕ) : String := s!"he:{n}"
def small (n : ℕ) : String := s!"small:{n}"
def zero (n : ℕ) : String := s!"zero:{n}"
def bias : String := "zero:0"

/-- A stage: its weights with their inits, inputs and outputs, and its body from the positions
of the weights and the inputs to the returned positions. The first `fac` outputs are factored
out. -/
structure Stage where
  name : String
  wts : List (FS × String)
  ins : List FS
  outs : List FS
  fac : ℕ := 0
  body : (w : List ℕ) → (x : List ℕ) → B (List ℕ)

/-- The stage as a program with the weights `Θ ++ C`, its own at `off` in `Θ`. -/
def Stage.at (s : Stage) (Θ C : List FS) (off : ℕ) : Prog flowSig :=
  let nw := Θ.length + C.length
  let (ret, st) := (s.body (List.range' off s.wts.length) (List.range' nw s.ins.length)).run
    { n := nw + s.ins.length }
  { wts := Θ ++ C, ins := s.ins, outs := s.outs, lines := st.lines.toList, ret }

/-- The stage as a program on its own weights. -/
def Stage.prog (s : Stage) : Prog flowSig := s.at (s.wts.map (·.1)) [] 0

/-- A model's weights: the stages' weights side by side. -/
def modelWts (stages : List Stage) : List FS := (stages.map (·.wts.map (·.1))).flatten

/-- The stages of a model as programs over the model's weights: each with the weights
`Θ ++ C`, `C` what the earlier stages factored out, and with that `C`. -/
def stageProgs (stages : List Stage) : List (Stage × Prog flowSig × List FS) :=
  let Θ := modelWts stages
  (stages.foldl (fun (acc, C, off) t =>
    (acc ++ [(t, t.at Θ C off, C)], C ++ t.outs.take t.fac, off + t.wts.length)) ([], [], 0)).1

/-- A model as one program, each stage run by `run` (from the positions of its weights and
inputs) in one builder, on the outputs the previous stage does not factor out. It returns the
factored-out outputs, then the rest of the last stage's. -/
def modelProgOf {Sg : Sig FS} (run : Stage → List ℕ → List ℕ → BG Sg (List ℕ))
    (stages : List Stage) : Option (Prog Sg) := do
  let s ← stages.head?
  let t ← stages.getLast?
  let Θ := modelWts stages
  let ((cs, y, _), st) := (stages.foldlM (fun (cs, x, off) (u : Stage) => do
      let y ← run u (List.range' off u.wts.length) x
      return (cs ++ y.take u.fac, y.drop u.fac, off + u.wts.length))
    (([] : List ℕ), List.range' Θ.length s.ins.length, 0)).run { n := Θ.length + s.ins.length }
  return { wts := Θ, ins := s.ins
           outs := (stages.map fun u => u.outs.take u.fac).flatten ++ t.outs.drop t.fac
           lines := st.lines.toList, ret := cs ++ y }

/-- A model as one program: the stages' bodies one after another. -/
def modelProg (stages : List Stage) : Option (Prog flowSig) := modelProgOf (·.body) stages

/-- A model as one program of calls: each stage a call of the definition of its name in `L`
(`Linked.find`), on its inputs, then its weights. -/
def callProg {Sg : Sig FS} (L : Linked flowSig Sg) (stages : List Stage) : Option (Prog Sg) :=
  modelProgOf (fun u w x => match L.find u.name with
    | some ⟨_, _, o⟩ => B.line o (x ++ w)
    | none => pure []) stages

/-! ## Stages shared by the models -/

/-- `score n`: the last stage, so that the forward log-weight is the log-density. -/
def scoreStage (n : ℕ) : Stage :=
  { name := "base", wts := [], ins := [vec n], outs := [gauss n]
    body := fun _ x => do return [← B.o1 (.score n) (x.headD 0)] }

/-- Data in `(0, 1)ⁿ`, given in its logit chart, as a vector. -/
def logitStage (n : ℕ) : Stage :=
  { name := "logit", wts := [], ins := [unit n], outs := [vec n]
    body := fun _ x => do return [← B.o1 (.logit n) (x.headD 0)] }

/-! ## Two dimensions -/

/-- An affine coupling on `vec (m + m)`, after an actnorm: `x₂ ↦ x₂ ⊙ exp (tanh s) + t` with
`(s, t)` an MLP of `x₁` with two hidden layers. The output is `(x₂', x₁)`. -/
def couplingN (name : String) (m hid : ℕ) : Stage where
  name := name
  wts := [(vec (m + m), bias), (vec (m + m), bias), (vec (hid * m), he m), (vec hid, bias),
    (vec (hid * hid), he hid), (vec hid, bias), (vec ((m + m) * hid), zero hid), (vec (m + m), bias)]
  ins := [vec (m + m)]
  outs := [vec (m + m)]
  body w x := do
    let w (i : ℕ) := w.getD i 0
    let x ← B.actnorm (m + m) (x.headD 0) (w 0) (w 1)
    let (x₁, x₂) ← B.o12 (.split m m) x
    let a ← B.o2 (.mvmul hid m) (w 2) x₁
    let a ← B.o2 (.add hid) a (w 3)
    let a ← B.o1 (.relu hid) a
    let a ← B.o2 (.mvmul hid hid) (w 4) a
    let a ← B.o2 (.add hid) a (w 5)
    let a ← B.o1 (.relu hid) a
    let a ← B.o2 (.mvmul (m + m) hid) (w 6) a
    let a ← B.o2 (.add (m + m)) a (w 7)
    let (s, t) ← B.o12 (.split m m) a
    let y ← B.affine m x₂ s t
    let y ← B.o2 (.concat m m) y x₁                        -- (x₂', x₁)
    return [y]

/-- A learned invertible linear map on `vec (m + m)` in block-LU form: a diagonal scale
`exp logS`, then the shears `a ↦ a + U b` and `b ↦ b + L a`. -/
def linmix (name : String) (m : ℕ) : Stage where
  name := name
  wts := [(vec (m + m), bias), (vec (m * m), zero m), (vec (m * m), zero m)]
  ins := [vec (m + m)]
  outs := [vec (m + m)]
  body w x := do
    let w (i : ℕ) := w.getD i 0
    let x ← B.o2 (.scaleexp (m + m)) (x.headD 0) (w 0)
    let (a, b) ← B.o12 (.split m m) x
    let u ← B.o2 (.mvmul m m) (w 1) b
    let a ← B.o2 (.add m) a u
    let l ← B.o2 (.mvmul m m) (w 2) a
    let b ← B.o2 (.add m) b l
    return [← B.o2 (.concat m m) a b]

/-- Augmentation: the data `x : vec n` and a hidden standard-normal input `u : gauss k`, which
`sample` turns into data and which is appended to `x`. -/
def augment (n k : ℕ) : Stage where
  name := "aug"
  wts := []
  ins := [vec n, gauss k]
  outs := [vec (n + k)]
  body _ x := do
    let e ← B.o1 (.sample k) (x.getD 1 0)
    return [← B.o2 (.concat n k) (x.headD 0) e]

/-- A flow on `ℝⁿ` (`n` even): `blocks` couplings of width `hid`, each after a `linmix` if
`mix`, then `score`. -/
def toyN (n blocks hid : ℕ) (mix : Bool := false) : List Stage :=
  ((List.range blocks).map fun k =>
    (if mix then [linmix s!"mix{k}" (n / 2)] else []) ++ [couplingN s!"c{k}" (n / 2) hid]).flatten ++
  [scoreStage (n / 2 + n / 2)]

/-- A flow on the plane: `toyN` on the data augmented by `2 * augPairs` hidden dimensions. -/
def toy2d (blocks hid : ℕ) (mix : Bool) (augPairs : ℕ) : List Stage :=
  (if augPairs = 0 then [] else [augment 2 (2 * augPairs)]) ++
  toyN (2 * (1 + augPairs)) blocks hid mix

/-! ## Permutations of images

Permutations as steps (`PermS`): `σ.trans τ` (first `σ`) is `τ ++ σ`, since a gather reads
`τ`'s index through `σ`'s. -/

/-- Squeezing `cin` channels of `(s f) × (s f)` by `f`, then a rotation of the flat vector by
`shift`. Output channel `chan c'`, with `c' = ci * f² + f dy + dx`, at `(r, q)` reads channel
`ci` at `(f r + dy, f q + dx)`. -/
def squeezeC (cin f s : ℕ) (chan : PermS) (shift : ℕ) : PermS :=
  let N := cin * (f * f) * (s * s)
  (if shift % N = 0 then [] else [.roll [N] 0 (N - shift % N)]) ++
  [.tr [cin, s, f, s, f] [0, 2, 4, 1, 3]] ++ chan.map (PermStep.embed 1 (s * s))

/-- `squeezeC` with one input channel. -/
def squeeze (f s : ℕ) (chan : PermS) (shift : ℕ) : PermS := squeezeC 1 f s chan shift

/-- Reordering `a × b` channels as `b × a`. -/
def chanSwap (a b : ℕ) : PermS := [.tr [b, a] [1, 0]]

/-- At `f = 2`: rows, columns, the checkerboard (`swap 1 2`, `swap 1 3` of the four positions). -/
def pairings2 : List PermS := [[], [.tr [2, 2] [1, 0]], [.roll [4] 0 3, .flip [4] 0]]

/-- A channel permutation of the channels of `n` pixels each: output channel `a` reads input
channel `π a`. -/
def chanPerm (n : ℕ) (π : PermS) : PermS := π.map (PermStep.embed 1 n)

/-- The rotation of `k` channels by `s`: channel `a` reads `a + s`. -/
def chanRot (k s : ℕ) : PermS := [.roll [k] 0 ((k - s % k) % k)]

/-! ## MNIST -/

/-- The coupling network of a Glow coupling on `c` channels of `h × w` (`hid` wide), reading
`a` and the weights at `w i₀, …`: a 3 × 3 conv, relu, a 1 × 1 conv, relu, a 3 × 3 conv to
`(s, t)`. -/
def glowNet (c h w hid : ℕ) (W : ℕ → ℕ) (a : ℕ) : B (ℕ × ℕ) := do
  let u ← B.o2 (.conv3x3 c hid h w) (W 0) a
  let u ← B.o2 (.addc hid h w) u (W 1)
  let u ← B.o1 (.relu (hid * (h * w))) u
  let u ← B.o2 (.conv1x1 hid hid h w) (W 2) u
  let u ← B.o2 (.addc hid h w) u (W 3)
  let u ← B.o1 (.relu (hid * (h * w))) u
  let u ← B.o2 (.conv3x3 hid (c + c) h w) (W 4) u
  let u ← B.o2 (.addc (c + c) h w) u (W 5)
  B.o12 (.rows2 c (h * w)) u

/-- The deeper coupling network of `glowResAt`: `u₁ = relu (conv₃ₓ₃ a)`, two residual layers
`uᵢ₊₁ = uᵢ + conv₃ₓ₃ (relu (conv₁ₓ₁ (relu uᵢ)))`, then `(s, t) = conv₃ₓ₃ (relu u₃)`. -/
def resNet (c h w hid : ℕ) (W : ℕ → ℕ) (a : ℕ) : B (ℕ × ℕ) := do
  let u ← B.o2 (.conv3x3 c hid h w) (W 0) a
  let u ← B.o2 (.addc hid h w) u (W 1)
  let mut u ← B.o1 (.relu (hid * (h * w))) u
  for j in [0:2] do
    let g ← B.o1 (.relu (hid * (h * w))) u
    let g ← B.o2 (.conv1x1 hid hid h w) (W (2 + 4 * j)) g
    let g ← B.o2 (.addc hid h w) g (W (3 + 4 * j))
    let g ← B.o1 (.relu (hid * (h * w))) g
    let g ← B.o2 (.conv3x3 hid hid h w) (W (4 + 4 * j)) g
    let g ← B.o2 (.addc hid h w) g (W (5 + 4 * j))
    u ← B.o2 (.add (hid * (h * w))) u g
  let g ← B.o1 (.relu (hid * (h * w))) u
  let g ← B.o2 (.conv3x3 hid (c + c) h w) (W 10) g
  let g ← B.o2 (.addc (c + c) h w) g (W 11)
  B.o12 (.rows2 c (h * w)) g

/-- The weights of `glowNet`. -/
def glowNetW (c hid : ℕ) : List (FS × String) :=
  [(vec (hid * (c * 9)), he (c * 9)), (vec hid, bias), (vec (hid * hid), he hid), (vec hid, bias),
   (vec ((c + c) * (hid * 9)), zero (hid * 9)), (vec (c + c), bias)]

/-- The weights of `resNet`. -/
def resNetW (c hid : ℕ) : List (FS × String) :=
  [(vec (hid * (c * 9)), he (c * 9)), (vec hid, bias)] ++
  (List.replicate 2 [(vec (hid * hid), he hid), (vec hid, bias),
    (vec (hid * (hid * 9)), zero (hid * 9)), (vec hid, bias)]).flatten ++
  [(vec ((c + c) * (hid * 9)), zero (hid * 9)), (vec (c + c), bias)]

/-- A Glow block at any resolution, after an actnorm: `c + c` channels of `h × w`, permuted by
`sq`; the halves are coupled in turn (`b' = b ⊙ exp sA(a) + tA(a)`, then
`a' = a ⊙ exp sB(b') + tB(b')`), and `sq⁻¹` restores the layout. `net` is the coupling network
and `netW` its weights. -/
def coupleBlock (name : String) (c h w : ℕ) (sq : PermS)
    (netW : List (FS × String)) (net : (ℕ → ℕ) → ℕ → B (ℕ × ℕ)) : Stage where
  name := name
  wts := [(vec (c * (h * w) + c * (h * w)), bias), (vec (c * (h * w) + c * (h * w)), bias)] ++
    netW ++ netW
  ins := [vec (c * (h * w) + c * (h * w))]
  outs := [vec (c * (h * w) + c * (h * w))]
  body W x := do
    let W (i : ℕ) := W.getD i 0
    let k := netW.length
    let x ← B.actnorm (c * (h * w) + c * (h * w)) (x.headD 0) (W 0) (W 1)
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) sq) x
    let (a, b) ← B.o12 (.split (c * (h * w)) (c * (h * w))) x
    let (sA, tA) ← net (fun i => W (2 + i)) a
    let b' ← B.affine (c * (h * w)) b sA tA
    let (sB, tB) ← net (fun i => W (2 + k + i)) b'
    let a' ← B.affine (c * (h * w)) a sB tB
    let y ← B.o2 (.concat (c * (h * w)) (c * (h * w))) a' b'
    return [← B.o1 (.perm (c * (h * w) + c * (h * w)) sq true) y]

/-- Factoring out half of the dimensions, with Glow's learned prior: after `sq`, the first half
`x₁` passes on, and the second is scored under `N (μ, σ²)` with `(μ, log σ)` a conv net of
`x₁`. The scored half is factored out, so it comes first: the stage returns `(z, x₁)`. -/
def splitPrior (c h w hid : ℕ) (sq : PermS) : Stage where
  name := "split"
  wts := [(vec (hid * (c * 9)), he (c * 9)), (vec hid, bias),
    (vec ((c + c) * (hid * 9)), zero (hid * 9)), (vec (c + c), bias)]
  ins := [vec (c * (h * w) + c * (h * w))]
  outs := [gauss (c * (h * w)), vec (c * (h * w))]
  fac := 1
  body W x := do
    let W (i : ℕ) := W.getD i 0
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) sq) (x.headD 0)
    let (x₁, x₂) ← B.o12 (.split (c * (h * w)) (c * (h * w))) x
    let u ← B.o2 (.conv3x3 c hid h w) (W 0) x₁
    let u ← B.o2 (.addc hid h w) u (W 1)
    let u ← B.o1 (.relu (hid * (h * w))) u
    let u ← B.o2 (.conv3x3 hid (c + c) h w) (W 2) u
    let u ← B.o2 (.addc (c + c) h w) u (W 3)
    let (μ, logσ) ← B.o12 (.rows2 c (h * w)) u
    let z ← B.o2 (.sub (c * (h * w))) x₂ μ
    let z ← B.o2 (.scalenexp (c * (h * w))) z logσ
    let z ← B.o1 (.score (c * (h * w))) z
    return [z, x₁]

/-- The multiscale coupling flow for MNIST with fixed permutations (`mnistCoupling` of
`semiInv`): `logit 784`, `fine` blocks on `4 × 14 × 14`, `splitPrior`, `coarse` residual
blocks on `8 × 7 × 7`, `score 392`. -/
def mnistCoupling (fine coarse hidF hidC : ℕ) : List Stage :=
  [logitStage 784] ++
  (List.range fine).map (fun k =>
    coupleBlock s!"fine{k}" 2 14 14
      (squeeze 2 14 (pairings2.getD (k % 3) []) (if k % 2 = 0 then 0 else 29))
      (glowNetW 2 hidF) (glowNet 2 14 14 hidF)) ++
  [splitPrior 2 14 14 hidF (squeeze 2 14 (pairings2.getD 2 []) 0)] ++
  (List.range coarse).map (fun k =>
    coupleBlock s!"coarse{k}" 4 7 7
      (squeezeC 2 2 7 (if k % 2 = 0 then [] else chanSwap 2 4) (15 * (k % 4)))
      (resNetW 4 hidC) (resNet 4 7 7 hidC)) ++
  [scoreStage 392]

/-! ## Glow -/

/-- One step of Glow on `c + c` channels of `h × w`, in the image layout (squeezed by `sq` and
back): an actnorm, the invertible 1 × 1 convolution `W = P L U` (`U`, `L` by `triu`, `L`
between two channel reversals, `P` the channel permutation `π`), and one affine coupling. -/
def glowStepAt (name : String) (c h w hid : ℕ) (sq : PermS)
    (π : PermS) : Stage where
  name := name
  wts := [(vec (c * (h * w) + c * (h * w)), bias), (vec (c * (h * w) + c * (h * w)), bias),
    (vec ((c + c) * (c + c)), small 64), (vec ((c + c) * (c + c)), small 64)] ++ glowNetW c hid
  ins := [vec (c * (h * w) + c * (h * w))]
  outs := [vec (c * (h * w) + c * (h * w))]
  body W x := do
    let W (i : ℕ) := W.getD i 0
    let rev := chanPerm (h * w) [.flip [c + c] 0]
    let x ← B.actnorm (c * (h * w) + c * (h * w)) (x.headD 0) (W 0) (W 1)
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) sq) x                           -- squeeze
    let x ← B.o2 (.triu (c + c) (h * w)) x (W 2)                                    -- W = P L U
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) rev) x
    let x ← B.o2 (.triu (c + c) (h * w)) x (W 3)
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) rev) x
    let x ← B.o1 (.perm (c * (h * w) + c * (h * w)) (chanPerm (h * w) π)) x
    let (a, b) ← B.o12 (.split (c * (h * w)) (c * (h * w))) x   -- the coupling
    let (s, t) ← glowNet c h w hid (fun i => W (4 + i)) a
    let b' ← B.affine (c * (h * w)) b s t
    let y ← B.o2 (.concat (c * (h * w)) (c * (h * w))) a b'
    return [← B.o1 (.perm (c * (h * w) + c * (h * w)) sq true) y]                   -- unsqueeze

/-- Glow for MNIST, in two levels: `logit 784`, `fine` steps on `4 × 14 × 14`, `splitPrior`,
`coarse` steps on `8 × 7 × 7`, `score 392`. -/
def mnistGlow (fine coarse hidF hidC : ℕ) : List Stage :=
  [logitStage 784] ++
  (List.range fine).map (fun k =>
    glowStepAt s!"fine{k}" 2 14 14 hidF (squeeze 2 14 [] 0) (chanRot 4 (k + 1))) ++
  [splitPrior 2 14 14 hidF (squeeze 2 14 (pairings2.getD 2 []) 0)] ++
  (List.range coarse).map (fun k =>
    glowStepAt s!"coarse{k}" 4 7 7 hidC (squeezeC 2 2 7 [] 0) (chanRot 8 (k + 1))) ++
  [scoreStage 392]

/-! ## TARFlow on image patches -/

/-- The weights of one transformer layer: attention `W_q, W_k, W_v, W_o`, then the MLP
`W₁, b₁, W₂, b₂` of width `4 h`. -/
def tfLayerW (h : ℕ) : List (FS × String) :=
  [(vec (h * h), he h), (vec (h * h), he h), (vec (h * h), he h), (vec (h * h), small h),
   (vec (4 * h * h), he h), (vec (4 * h), bias), (vec (h * (4 * h)), small (4 * h)), (vec h, bias)]

/-- The weights of `tarflowBlock`: the token embedding, one position embedding per token, the
layers, the head to `(α, μ)`, and `(α₁, μ₁)` for the first token. -/
def tarflowW (T d h L : ℕ) : List (FS × String) :=
  [(vec (h * d), he d)] ++ List.replicate T (vec h, small h) ++
  (List.replicate L (tfLayerW h)).flatten ++
  [(vec ((d + d) * h), zero h), (vec (d + d), bias), (vec d, bias), (vec d, bias)]

/-- One TARFlow block on `T ≥ 1` tokens of dimension `d`, in the token layout given by `σ`:
`zₙ = (xₙ − μₙ) ⊙ exp (−tanh αₙ)`, with `(αₙ, μₙ)` read off a causal transformer (width `h`,
`L` pre-norm layers of single-head attention and a GELU MLP, a key-value cache per layer) at
token `n − 1`, and learned `(α₁, μ₁)`. `σ⁻¹` restores the layout. -/
def tarflowBlock (name : String) (T d h L : ℕ) (σ : PermS) : Stage where
  name := name
  wts := tarflowW T d h L
  ins := [vec (d * T)]
  outs := [vec (d * T)]
  body ws x := do
    let w (i : ℕ) := ws.getD i 0
    let W_in := w 0
    let pos (j : ℕ) := w (1 + j)
    let lw (ℓ i : ℕ) := w (1 + T + 8 * ℓ + i)
    let W_hd := w (1 + T + 8 * L); let b_hd := w (2 + T + 8 * L)
    let α₁ := w (3 + T + 8 * L); let μ₁ := w (4 + T + 8 * L)
    let x ← B.o1 (.perm (d * T) σ) (x.headD 0)
    -- the tokens, split off the front one by one
    let mut xs : Array ℕ := #[]
    let mut r := x
    for j in [0:T - 1] do
      let (t, r') ← B.o12 (.split d (d * (T - 1 - j))) r
      xs := xs.push t; r := r'
    xs := xs.push r
    let nil ← B.o0 .nil
    let mut K : Array ℕ := Array.replicate L nil
    let mut Vs : Array ℕ := Array.replicate L nil
    let mut αμ := (α₁, μ₁)
    let mut y := 0
    for j in [0:T] do
      let xj := xs.getD j 0
      -- the affine step of token j
      let a ← B.o1 (.tanh d) αμ.1
      let z ← B.o2 (.sub d) xj αμ.2
      let z ← B.o2 (.scalenexp d) z a
      -- joined onto the tokens before it
      y ← if j = 0 then pure z else B.o2 (.concat (d * j) d) y z
      if j + 1 < T then
        -- token j through the transformer, appending to the caches, predicts token j + 1
        let e ← B.o2 (.mvmul h d) W_in xj
        let mut e ← B.o2 (.add h) e (pos j)
        for ℓ in [0:L] do
          let l ← B.o1 (.layernorm h) e
          let k ← B.o2 (.mvmul h h) (lw ℓ 1) l
          let k ← B.o2 (.concat (h * j) h) (K.getD ℓ 0) k
          K := K.set! ℓ k
          let v ← B.o2 (.mvmul h h) (lw ℓ 2) l
          let v ← B.o2 (.concat (h * j) h) (Vs.getD ℓ 0) v
          Vs := Vs.set! ℓ v
          let q ← B.o2 (.mvmul h h) (lw ℓ 0) l
          let a ← B.o3 (.attend (j + 1) h) q k v
          let a ← B.o2 (.mvmul h h) (lw ℓ 3) a
          e ← B.o2 (.add h) e a
          let l ← B.o1 (.layernorm h) e
          let m ← B.o2 (.mvmul (4 * h) h) (lw ℓ 4) l
          let m ← B.o2 (.add (4 * h)) m (lw ℓ 5)
          let m ← B.o1 (.gelu (4 * h)) m
          let m ← B.o2 (.mvmul h (4 * h)) (lw ℓ 6) m
          let m ← B.o2 (.add h) m (lw ℓ 7)
          e ← B.o2 (.add h) e m
        let l ← B.o1 (.layernorm h) e
        let p ← B.o2 (.mvmul (d + d) h) W_hd l
        let p ← B.o2 (.add (d + d)) p b_hd
        let (α, μ) ← B.o12 (.split d d) p
        αμ := (α, μ)
    return [← B.o1 (.perm (d * T) σ true) y]

/-- The image `s f × s f` cut into `s × s` patches of `f × f`, patch-major. -/
def patchify (f s : ℕ) : PermS := squeeze f s [] 0 ++ chanSwap (s * s) (f * f)

/-- The token order reversed (`T` tokens of `d`). -/
def revTokens (T d : ℕ) : PermS := [.flip [T, d] 0]

/-- The image `n × n` rolled by `k` pixels down and right. -/
def roll2 (n k : ℕ) : PermS := [.roll [n, n] 0 ((n - k % n) % n), .roll [n, n] 1 ((n - k % n) % n)]

/-- TARFlow for Fashion-MNIST: `logit 784`, then `blocks` TARFlow blocks on the `(28 / p)²`
patches of `p × p` pixels, then `score 784`. Odd blocks reverse the token order; blocks `2, 3`
(mod 4) roll the image by `p / 2` pixels first. -/
def mnistTarflow (p blocks h L : ℕ) : Option (List Stage) := do
  let s := 28 / p
  guard (p * s = 28)
  let bs := (List.range blocks).map fun k =>
    let σ₁ := if k % 2 = 1 then patchify p s ++ revTokens (s * s) (p * p) else patchify p s
    let σ := if k % 4 ≥ 2 then roll2 (p * s) (p / 2) ++ σ₁ else σ₁
    tarflowBlock s!"block{k}" (s * s) (p * p) h L σ
  pure ([logitStage (p * p * (s * s))] ++ bs ++ [scoreStage (p * p * (s * s))])

end RoleFlow.Flow