import RoleFlow.Flow.Models
import RoleFlow.Flow.Swap

/-! # The JAX printer (trusted)

A typed program `P` of the flow signature, with its roles `ρ`, is printed as two JAX functions
that also return the log-weight: the forward function from the lines of `P`, and the inverse
from the lines of the inverse program `P†` (`Prog.inverse`, the schedule of the recipes of the
derivation). The inverse program is typed too, with its own roles `ρ†` (inferred), and both are
printed by the same code (`emitFwd`): the log-weight is `L_θ`, the sum of `ℓ_o` over the spine
steps, accumulated in `ld` with one `ld += ℓ_o` on each spine step's row, the `;` aligned. For
`P†` it is the log-weight of `Prog.density` for `P†`; that it agrees a.e. with `-L_θ` at the
values `P†` recomputes (`Prog.inverse_density`) is not proved (`TODO.md`), but checked by
`scripts/compare.py`. Every variable is a Python expression; structural operations (`split`,
the changes of sort) emit slices or nothing, the others bind a fresh variable; a spine step that
prints nothing gets a row `pass ; ld += …`.

Each stage is a definition (`Linked.add`) printed as `{name}_fwd`, `{name}_inv` over its own
weights. A model is printed as **one program** `flow_fwd` (and its inverse program `flow_inv`),
so the stage pipeline is a typed program rather than a Python composition: its stages inlined
(`Flow.modelProg`), or, with `calls`, a call of each stage (`Flow.callProg`), a row
`y, l = {name}_fwd(w[a:b], x) ; ld += l`. A call is a spine step like any other, and the inverse
program calls `{name}_inv`. The roles are inferred (`Prog.infer`) on each stage and on the model.
`writeModule` refuses a model or a stage that is not typed, or whose inverse program is not.

The leaf table (one JAX expression per operation and per nonzero log-weight) is trusted: it is
written to agree with `Flow.semOp` and `Flow.ℓOp`. Functions act on one example; the caller
batches with `jax.vmap`. -/

namespace RoleFlow.Flow.Print

open FS

/-! ## The printer's state -/

structure Gen where
  /-- The rows so far: the code, and the log-weight statement after its `;` (`""`: none). -/
  lines : Array (String × String) := #[]
  next : ℕ := 0
  /-- Each permutation's expression (on `x`, of length `n`) with its index list, for
  `<module>_perms.py`. -/
  permChecks : Array (String × ℕ × List ℕ) := #[]

abbrev M := StateM Gen

def emitLine (l : String) : M Unit := modify fun g => { g with lines := g.lines.push (l, "") }

/-- Bind an expression to a fresh variable. -/
def bindE (e : String) : M String :=
  modifyGet fun g =>
    let v := s!"v{g.next}"
    (v, { g with lines := g.lines.push (s!"{v} = {e}", ""), next := g.next + 1 })

/-! ## Permutations

A permutation is printed as its steps (`Flow.PermS`), one numpy call each, applied in order; its
inverse as the steps of `PermS.inv`. `writeModule` also writes each permutation's expression with
its index list (`PermS.idxList`) to `<module>_perms.py`, which checks that numpy agrees. -/

/-- The comma-separated numbers. -/
def numsE (l : List ℕ) : String := ", ".intercalate (l.map toString)

/-- A step applied to the expression `x`; the result keeps the step's shape (`reshape` reads
row-major, so the next step can reshape it directly). -/
def stepE (x : String) : PermStep → String
  | .tr dims axes => s!"{x}.reshape({numsE dims}).transpose({numsE axes})"
  | .flip dims k => s!"jnp.flip({x}.reshape({numsE dims}), {k})"
  | .roll dims k s => s!"jnp.roll({x}.reshape({numsE dims}), {s}, axis={k})"

/-- The steps applied to the flat vector `x`, in order, flattened at the end; steps that do
nothing (`PermStep.isId`) are left out. -/
def permStepsE (x : String) (ps : PermS) : String :=
  match ps.filter (!·.isId) with
  | [] => x
  | qs => qs.foldl stepE x ++ ".reshape(-1)"

/-- The empty vector is printed as `bool`, which `jnp.concatenate` promotes to the dtype of what
it is joined with. -/
def concatE : List String → String
  | [] => "jnp.zeros(0, bool)"
  | es => s!"jnp.concatenate([{", ".intercalate es}])"

/-- Several values as one Python value: a tuple, or the value itself. -/
def tupE : List String → String
  | [e] => e
  | es => s!"({", ".intercalate es})"

/-! ## Leaves -/

/-- The function computing `Ref.logpdf`, where it is not `0`. -/
def logpdfE : Ref → Option String
  | .leb => none
  | .normal => some "logphi"
  | .logistic => some "loglogistic"

/-- An operation's results from its arguments, one expression per result; agrees with
`Flow.semOp`. -/
def emitOp {as bs : List FS} : FOp as bs → List String → M (List String)
  | .nil, _ => do return [← bindE "jnp.zeros(0, bool)"]
  | .split m _, xs => pure [s!"{xs.headD ""}[:{m}]", s!"{xs.headD ""}[{m}:]"]
  | .concat _ _, xs => do return [← bindE (concatE xs)]
  | .perm n ps inv, xs => do
    let qs := if inv then ps.inv else ps
    let e := permStepsE "x" qs
    modify fun g => if g.permChecks.any fun (e', n', _) => e' == e && n' == n then g
      else { g with permChecks := g.permChecks.push (e, n, qs.idxList n) }
    let x := xs.headD ""
    let y := permStepsE x qs
    if y == x then pure [x] else return [← bindE y]
  | .rebase .., xs => pure xs
  | .relu _, xs => do return [← bindE s!"jax.nn.relu({xs.headD ""})"]
  | .tanh _, xs => do return [← bindE s!"jnp.tanh({xs.headD ""})"]
  | .layernorm _, xs => do return [← bindE s!"layernorm({xs.headD ""})"]
  | .gelu _, xs => do return [← bindE s!"jax.nn.gelu({xs.headD ""}, approximate=True)"]
  | .oddpow _ k inv, xs => do
    return [← bindE (if inv then s!"oddroot({2 * k + 1}, {xs.headD ""})"
      else s!"{xs.headD ""} ** {2 * k + 1}")]
  | .rows2 m n, xs => do
    let e ← bindE (xs.headD "")
    return [s!"{e}[:{m * n}]", s!"{e}[{m * n}:]"]
  | .add _ inv, xs => do
    return [← bindE s!"{xs.getD 0 ""} {if inv then "-" else "+"} {xs.getD 1 ""}"]
  | .scaleexp _ inv, xs => do
    return [← bindE s!"{xs.getD 0 ""} * jnp.exp({if inv then "-" else ""}{xs.getD 1 ""})"]
  | .triu c n inv, xs => do
    return [← bindE s!"triu_{if inv then "inv" else "fwd"}({c}, {n}, {xs.getD 0 ""}, {xs.getD 1 ""})"]
  | .mvmul m n, xs => do return [← bindE s!"{xs.getD 0 ""}.reshape({m}, {n}) @ {xs.getD 1 ""}"]
  | .mul _, xs => do return [← bindE s!"{xs.getD 0 ""} * {xs.getD 1 ""}"]
  | .addc c h w, xs => do
    return [← bindE s!"({xs.getD 0 ""}.reshape({c}, {h * w}) + {xs.getD 1 ""}[:, None]).reshape(-1)"]
  | .conv1x1 cin cout h w, xs => do
    return [← bindE
      s!"({xs.getD 0 ""}.reshape({cout}, {cin}) @ {xs.getD 1 ""}.reshape({cin}, {h * w})).reshape(-1)"]
  | .conv3x3 cin cout h w, xs => do
    return [← bindE
      s!"conv3x3({xs.getD 0 ""}.reshape({cout}, {cin}, 3, 3), {xs.getD 1 ""}.reshape({cin}, {h}, {w})).reshape(-1)"]
  | .attend m h, xs => do
    return [← bindE s!"attend({m}, {h}, {xs.getD 0 ""}, {xs.getD 1 ""}, {xs.getD 2 ""})"]

/-- The log-weight `ℓ_o` on the arguments, where it is not `0`; agrees with `Flow.ℓOp`
(`rebase r r'` has `log ρ_{r'} - log ρ_r`, a term for each density that is not `1`;
`oddpow_logdet p x` is `Σ log p x^(p-1)`, which `opowLogdet` sets to `0` on the null set of the
axes;
`triu_logdet c n U = n Σ_a U_aa` is `triuLogdet c n U`). -/
def ldOp {as bs : List FS} : FOp as bs → Option (List String → String)
  | .rebase r r' _ => match logpdfE r, logpdfE r' with
    | none, none => none
    | none, some f' => some fun xs => s!"{f'}({xs.getD 0 ""})"
    | some f, none => some fun xs => s!"-{f}({xs.getD 0 ""})"
    | some f, some f' => some fun xs => s!"{f'}({xs.getD 0 ""}) - {f}({xs.getD 0 ""})"
  | .scaleexp _ inv => some fun xs => s!"{if inv then "-" else ""}jnp.sum({xs.getD 1 ""})"
  | .triu c n inv => some fun xs => s!"{if inv then "-" else ""}triu_logdet({c}, {n}, {xs.getD 1 ""})"
  | .oddpow _ k inv => some fun xs =>
    let p := 2 * k + 1
    if inv then s!"-oddpow_logdet({p}, oddroot({p}, {xs.getD 0 ""}))"
    else s!"oddpow_logdet({p}, {xs.getD 0 ""})"
  | _ => none

/-- Put `ld += t` after the row the current line printed, if it printed one (the rows were `n`
before it), or else on a row `pass` of its own. -/
def addLd (n : ℕ) (t : String) : M Unit := modify fun g =>
  if n < g.lines.size then
    { g with lines := g.lines.modify (g.lines.size - 1) fun (c, _) => (c, s!"ld += {t}") }
  else { g with lines := g.lines.push ("pass", s!"ld += {t}") }

/-! ## Calls -/

/-- The weights at the positions `ps`, as a list: a slice `w[a:b]` if they are the function's
own weights `a, …, b - 1`. -/
def weightsE (env : Array String) (ps : List ℕ) : String :=
  match ps with
  | a :: _ =>
    if ps == List.range' a ps.length && ps.all fun j => env.getD j "" == s!"w[{j}]" then
      s!"w[{a}:{a + ps.length}]"
    else s!"[{", ".intercalate (ps.map (env.getD · "MISSING"))}]"
  | [] => "[]"

/-- A row calling `{name}_fwd` (or `_inv`) on the weights `ws` and the values `x`, binding its
`nb` results and its log-weight; returns their expressions. -/
def emitCall (name : String) (inv : Bool) (ws x : String) (nb : ℕ) : M (List String × String) :=
  modifyGet fun g =>
    let vs := (List.range' g.next nb).map fun i => s!"v{i}"
    let l := s!"l{g.next + nb}"
    let row := s!"{tupE vs}, {l} = {name}_{if inv then "inv" else "fwd"}({ws}, {x})"
    ((vs, l), { g with lines := g.lines.push (row, ""), next := g.next + nb + 1 })

/-! ## Programs -/

/-- An operation of the signature the printed program is over: of the flow signature, or a call
(`Linked.view`). -/
abbrev View (Sg : Sig FS) := ∀ {as bs : List FS}, Sg.Op as bs → Callee flowSig.Op as bs

/-- Print the lines from the expressions of the context they start in, each spine step with
its `ℓ_o` (for a call, the log-weight it returns); returns the expressions of all variables, by
position. -/
def emitLines {Sg : Sig FS} (view : View Sg) (env : Array String) (ls : List (Line Sg))
    (ρ : List Role) : M (Array String) :=
  (ls.zip ρ).foldlM (fun env (l, r) => do
    let xs := l.args.map fun a => env.getD a "MISSING"
    let n := (← get).lines.size
    let (vs, ld) ← match view l.op with
      | .base o => (·, (ldOp o).map (· xs)) <$> emitOp o xs
      | .call name inv k =>
        (fun (vs, t) => (vs, some t)) <$>
          emitCall name inv (weightsE env (l.args.drop k)) (tupE (xs.take k)) l.bs.length
    if r = .spine then
      if let some t := ld then addLd n t
    return env ++ vs.toArray) env

/-- The column of the `;`: after the longest row that has a log-weight statement, ignoring rows
longer than this. -/
def ldCol : ℕ := 96

/-- `def name(w, x)` returning `(value, ld)`, from the body's rows: `ld` starts at `0`, and the
log-weight statements are aligned at one column. -/
def emitDef (name : String) (body : M String) : M String := do
  let saved := (← get).lines
  modify fun g => { g with lines := #[("pass", "ld  = jnp.zeros(())")] }
  let y ← body
  emitLine s!"return {y}, ld"
  let rows := (← get).lines
  modify fun g => { g with lines := saved }
  let col := rows.foldl (fun m (c, t) => if t.isEmpty || c.length > ldCol then m else max m c.length) 0
  return s!"def {name}(w, x):\n" ++ rows.foldl (fun s (c, t) =>
    s ++ "    " ++ (if t.isEmpty then c else s!"{c.pushn ' ' (col - c.length)} ; {t}") ++ "\n") ""

/-- The inputs of a function taking `n` values as `x`. -/
def argsE (n : ℕ) : List String :=
  if n = 1 then ["x"] else (List.range n).map fun j => s!"x[{j}]"

/-- The function of `P` from the expressions of its weights and inputs: the returned expressions,
after the rows that compute them and accumulate the log-weight. -/
def emitFwd {Sg : Sig FS} (view : View Sg) (P : Prog Sg) (ρ : List Role) (env₀ : List String) :
    M (List String) := do
  let env ← emitLines view env₀.toArray P.lines ρ
  return P.ret.map fun k => env.getD k "MISSING"

/-- `{name}_fwd` from `P` with its roles `ρ`, and `{name}_inv` from `P† = P.inverse ρ` with its
roles `ρi`, both with the log-weight. The weights are the list `w`. -/
def emitPair {Sg : Sig FS} (view : View Sg) (name : String) (P : Prog Sg) (ρ ρi : List Role) :
    M String := do
  let ws := (List.range P.nw).map fun i => s!"w[{i}]"
  let fwd ← emitDef s!"{name}_fwd" do return tupE (← emitFwd view P ρ (ws ++ argsE P.ins.length))
  let inv ← emitDef s!"{name}_inv" do
    return tupE (← emitFwd view (P.inverse ρ) ρi (ws ++ argsE P.outs.length))
  return fwd ++ "\n" ++ inv

/-! ## Modules -/

/-- Helpers the leaf table refers to, matching `Flow.FlowOps`. -/
def preamble : String := "# Generated by roleflow/RoleFlow/Flow/Print.lean (scripts/Jax.lean). Do not edit.
import math
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax


def logphi(x):
    \"\"\"log stdNormalPDF.\"\"\"
    return -0.5 * jnp.sum(x * x) - 0.5 * x.size * math.log(2 * math.pi)


def loglogistic(t):
    \"\"\"log logisticPDF: Σ log σ'(t), with σ'(t) = σ(t) σ(-t).\"\"\"
    return -jnp.sum(jax.nn.softplus(t) + jax.nn.softplus(-t))


def conv3x3(K, v):
    \"\"\"fconv3x3: SAME-zero-padded cross-correlation, K : (cout, cin, 3, 3), v : (cin, h, w).\"\"\"
    return lax.conv_general_dilated(v[None], K, (1, 1), 'SAME')[0]


def layernorm(v):
    \"\"\"flayernorm, eps = 1e-5.\"\"\"
    mu = jnp.mean(v)
    return (v - mu) / jnp.sqrt(jnp.mean((v - mu) ** 2) + 1e-5)


def attend(m, h, q, K, V):
    \"\"\"fattend: rows j * h + . of K and V.\"\"\"
    s = K.reshape(m, h) @ q / math.sqrt(h)
    return jax.nn.softmax(s) @ V.reshape(m, h)


def triu_mat(c, U):
    \"\"\"triuMat's block: upper triangular, diagonal exp U_aa, U row-major c x c.\"\"\"
    U = U.reshape(c, c)
    return jnp.triu(U, 1) + jnp.diag(jnp.exp(jnp.diag(U)))


def triu_fwd(c, n, x, U):
    \"\"\"ftriu (with m = c * n): the block applied at each of the n pixels.\"\"\"
    return (triu_mat(c, U) @ x.reshape(c, n)).reshape(-1)


def triu_inv(c, n, x, U):
    \"\"\"ftriuInv: a triangular solve at each pixel.\"\"\"
    return jax.scipy.linalg.solve_triangular(triu_mat(c, U), x.reshape(c, n)).reshape(-1)


def triu_logdet(c, n, U):
    \"\"\"triuLogdet: n * sum_a U_aa.\"\"\"
    return n * jnp.sum(jnp.diag(U.reshape(c, c)))


def oddroot(p, y):
    \"\"\"oroot: the real p-th root of y, for odd p.\"\"\"
    return jnp.cbrt(y) if p == 3 else jnp.sign(y) * jnp.abs(y) ** (1.0 / p)


def oddpow_logdet(p, x):
    \"\"\"opowLogdet: Σ log (p x_i^(p-1)), the log-weight of x ↦ x^p for odd p.\"\"\"
    return x.size * jnp.log(p) + (p - 1) * jnp.sum(jnp.log(jnp.abs(x)))

"

/-- The postamble: the model's `forward`/`inverse` (the model as one program) and
`init_weights`. -/
def postamble (nOut : ℕ) : String := s!"

def _flat(ws):
    return [a for w in ws for a in w]


def forward(ws, x):
    \"\"\"The model, one program: ws[k] is stage k's weights as a list. Returns (zs, logdet), zs
    the factored-out latents followed by the last stage's outputs.\"\"\"
    z, ld = flow_fwd(_flat(ws), x)
    return {if nOut = 1 then "[z]" else "list(z)"}, ld


def inverse(ws, zs):
    \"\"\"The inverse program of the model, from latents shaped like forward's. Returns
    (x, logdet of the inverse).\"\"\"
    return flow_inv(_flat(ws), {if nOut = 1 then "zs[0]" else "tuple(zs)"})


def init_weights(key, test=False):
    \"\"\"Weights per stage, from each stage's WINIT. For training, `he:n` is He-normal with
    fan-in n, `small:n` is normal with standard deviation 0.05, and `zero:n` is zero, so every
    coupling starts as the identity. `test=True` randomises every weight, so that no block is
    trivial.\"\"\"
    ws = []
    for _, _, dims, inits, _ in STAGES:
        w = []
        for d, s in zip(dims, inits):
            key, k = jax.random.split(key)
            kind, fan = s.split(':')
            fan = int(fan)
            if test:
                std = math.sqrt(1.0 / fan) if fan else 0.1
            else:
                std = (math.sqrt(2.0 / fan) if kind == 'he' else
                       0.05 if kind == 'small' else 0.0)
            w.append(std * jax.random.normal(k, (d,)))
        ws.append(w)
    return ws
"

/-- A stage's two functions (over its own weights), and its weights' dimensions and inits. -/
def emitStage {Sg : Sig FS} (view : View Sg) (s : Stage) (Q : Prog Sg) (ρ ρi : List Role) :
    M String := do
  let code ← emitPair view s.name Q ρ ρi
  return code ++ s!"
{s.name}_WDIMS = {s.wts.map (·.1.dim)}
{s.name}_WINIT = {s.wts.map fun (_, i) => s!"'{i}'"}
"

/-- The definitions of a tower, in order, each with the dimensions and inits of the weights of
the stage of its name. -/
def emitDefs (stages : List Stage) : {Sg : Sig FS} → Linked flowSig Sg → M (List String)
  | _, .base => pure []
  | _, .add L name Q ρ ρi .. => do
    let earlier ← emitDefs stages L
    let code ← match stages.find? (·.name == name) with
      | some s => emitStage L.view s Q ρ ρi
      | none => emitPair L.view name Q ρ ρi
    return earlier ++ [code]

/-- A certified model: its stages, as the definitions of a tower `L` (each over its own
weights), and the model as one program `P` at the top of the tower with its roles `ρ` and the
roles `ρi` of `P.inverse ρ`: the stages inlined, or a program of calls of them. -/
structure Cert where
  stages : List Stage
  Sg : Sig FS
  L : Linked flowSig Sg
  P : Prog Sg
  ρ : List Role
  ρi : List Role

/-- The module: the stages, the model (`flow_fwd` from `P`, `flow_inv` from `P†`), and the
Python interface; with each permutation's expression and index list. -/
def emitModuleChecks (c : Cert) : String × Array (String × ℕ × List ℕ) :=
  let (defs, g) := (show M (List String) from do
    return (← emitDefs c.stages c.L) ++ [← emitPair c.L.view "flow" c.P c.ρ c.ρi]).run {}
  let names := ", ".intercalate (c.stages.map fun s =>
    s!"({s.name}_fwd, {s.name}_inv, {s.name}_WDIMS, {s.name}_WINIT, {s.fac})")
  (preamble ++ "\n\n" ++ "\n\n".intercalate defs ++ s!"

# Each stage on its own: (forward, inverse, weight dimensions, inits, outputs it factors out).
# A stage's functions take its own weights.
STAGES = [{names}]

# The dimensions of the latents `forward` returns: the factored-out ones, then the last outputs.
ZDIMS = {c.P.outs.map FS.dim}

# The dimensions of the input's parts (more than one: the data, then hidden inputs).
IN_DIMS = {c.P.ins.map FS.dim}
" ++ postamble c.P.outs.length, g.permChecks)

/-- The module (`emitModuleChecks`). -/
def emitModule (c : Cert) : String := (emitModuleChecks c).1

/-- A script that checks each permutation's numpy expression, applied to `0, …, n - 1`, against
its index list (as Lean computes it). -/
def permCheckScript (module : String) (checks : Array (String × ℕ × List ℕ)) : String :=
  let rows := checks.foldl (fun s (e, n, a) => s ++ s!"    ({e.quote}, {n}, {a}),\n") ""
  s!"# Generated with {module} by roleflow/RoleFlow/Flow/Print.lean. Do not edit.
# Each permutation's numpy expression, applied to 0, …, n - 1, against its index list (as Lean
# computes it).
import numpy as np

CHECKS = [
{rows}]

for e, n, a in CHECKS:
    assert np.array_equal(eval(e, \{\"np\": np, \"jnp\": np, \"x\": np.arange(n)}), np.array(a)), e
print(f\"{module}: \{len(CHECKS)} permutations ok\")
"

/-- The roles of `P` and of its inverse program `P.inverse ρ`, by inference, with the typings;
`what` names `P` in the error. -/
def inferBoth {Sg : Sig FS} (what : String) (P : Prog Sg) :
    Except String ((ρ : List Role) × (ρi : List Role) ×' P.Typed ρ ∧ (P.inverse ρ).Typed ρi) :=
  match h : P.infer with
  | none => throw s!"{what}: role inference fails"
  | some ρ => match h' : (P.inverse ρ).infer with
    | none => throw s!"{what}: role inference fails on the inverse program"
    | some ρi => pure ⟨ρ, ρi, (P.infer_iff ρ).1 h, ((P.inverse ρ).infer_iff ρi).1 h'⟩

/-- The operation, if it is a permutation that is not `PermS.valid` (its index map is then not a
permutation, and the model does not read it as the steps). -/
def badPerm? : {as bs : List FS} → FOp as bs → Option String
  | _, _, .perm n ps _ => if ps.valid n then none else some s!"{ps.text "x"} on {n} entries"
  | _, _, _ => none

/-- Certify a model: add each stage, after the front end (`swapArgs`), as a definition of a
tower (`Linked.add`, with its roles and its inverse program's, inferred), then infer the roles
of the model as one program and of its inverse program: the stages inlined (`modelProg`, after
the front end), or, with `calls`, a call of each (`callProg`). Fails if something is not typed,
if two stages have one name, or if a permutation is not valid. -/
def certify (stages : List Stage) (calls : Bool := false) : Except String Cert := do
  if let some e := stages.findSome? fun s => s.prog.lines.findSome? fun l => badPerm? l.op then
    throw s!"not a permutation: {e}"
  let names := stages.map (·.name)
  if names.eraseDups.length != names.length then throw "two stages have the same name"
  let ⟨Sg, L⟩ ← stages.foldlM (fun (⟨_, L⟩ : Sigma (Linked flowSig)) s =>
      let Q := (swapArgs s.prog).mapOp fun o => L.lift o
      match inferBoth s!"stage {s.name}" Q with
      | .error e => throw e
      | .ok ⟨ρ, ρi, hT, hTi⟩ =>
        if hi : Q.ins = [] then throw s!"stage {s.name} has no inputs"
        else if ho : Q.outs = [] then throw s!"stage {s.name} has no outputs"
        else pure ⟨_, L.add s.name Q ρ ρi hT hTi hi ho⟩) ⟨flowSig, .base⟩
  let P? := if calls then callProg L stages
    else (modelProg stages).map fun P => (swapArgs P).mapOp fun o => L.lift o
  let some P := P? | throw "the model has no stages"
  match inferBoth "the model" P with
  | .error e => throw e
  | .ok ⟨ρ, ρi, _⟩ => return ⟨stages, Sg, L, P, ρ, ρi⟩

/-- Certify a model (`certify`) and write its module, and `<module>_perms.py`
(`permCheckScript`). -/
def writeModule (path : System.FilePath) (stages : List Stage) (calls : Bool := false) :
    IO Unit := do
  match certify stages calls with
  | .error e => throw (IO.userError e)
  | .ok c =>
    if let some d := path.parent then IO.FS.createDirAll d
    let (code, checks) := emitModuleChecks c
    IO.FS.writeFile path code
    let stem := path.fileStem.getD "module"
    IO.FS.writeFile (path.withFileName s!"{stem}_perms.py")
      (permCheckScript (path.fileName.getD "module") checks)

end RoleFlow.Flow.Print