"""Export the uniform-lobster rate certificate to a Lean data file.

Usage (from the repository root):
    python3 formal/scripts/export_lobster_cells.py > formal/MaxLeafFormal/LobsterData.lean

Reproduces the cover of [0, 3/4] that ancillary/lobster_rate_certificate.py
builds for B_lo = 1.71927 (same initial grid, same bisection order, same float Perron vectors
rationalised to denominator 10^15) and writes, for every cell [a, b], the vector w, the number
lam = max_j (w^T M_b)_j / w_j and c = max_i f_i / w_i. It also writes the right Perron vector v and
lam' = min_i (M v)_i / v_i at y* = 0.3453 for B_hi = 1.71928.

The generated file is data, not proof: Lean re-checks every cell condition in the kernel
(`MaxLeafFormal.Lobster.checkCells`, `checkUpper`), so a wrong export can only make a check fail.
"""
import hashlib
import importlib.util
import os
import sys
from fractions import Fraction as Fr

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SRC = os.path.join(ROOT, "ancillary", "lobster_rate_certificate.py")


def load():
    spec = importlib.util.spec_from_file_location("lobster_rate_certificate", SRC)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def q(x):
    x = Fr(x)
    assert x >= 0
    return f"q {x.numerator} {x.denominator}"


def main():
    sys.stdout.reconfigure(encoding="utf-8", newline="\n")
    sha = hashlib.sha256(open(SRC, "rb").read()).hexdigest()
    m = load()
    B = Fr("1.71927")
    B7 = B ** 7
    ymax, init, maxdepth = Fr(3, 4), 64, 40
    # the bisection of lobster_rate_certificate.lower, recording each accepted cell
    cells = [(ymax * i / init, ymax * (i + 1) / init, 0) for i in range(init)]
    out_cells = []
    while cells:
        a, b, d = cells.pop()
        if m.cell_check(a, b, B7) is None:
            assert d < maxdepth
            mid = (a + b) / 2
            cells += [(a, mid, d + 1), (mid, b, d + 1)]
            continue
        u, M = m.uM(b)
        Mf = [[float(x) for x in r] for r in M]
        w = [m.rat(x) for x in m.perron(Mf, left=True)]
        wM = [sum(w[i] * M[i][j] for i in range(3)) for j in range(3)]
        lam = max(wM[j] / w[j] for j in range(3))
        c = max(m.F[i] / w[i] for i in range(3))
        out_cells.append((a, b, w, lam, c))
    out_cells.sort(key=lambda t: t[0])
    ystar = Fr("0.3453")
    u, M = m.uM(ystar)
    Mf = [[float(x) for x in r] for r in M]
    v = [m.rat(x) for x in m.perron(Mf, left=False)]
    Mv = [sum(M[i][j] * v[j] for j in range(3)) for i in range(3)]
    lamp = min(Mv[i] / v[i] for i in range(3))
    out = [
        "import Mathlib",
        "",
        "/-!",
        "# Uniform-lobster certificate data (generated)",
        "",
        "Generated by `formal/scripts/export_lobster_cells.py` from the cover built by",
        f"`ancillary/lobster_rate_certificate.py`, SHA-256 `{sha}`.",
        f"{len(out_cells)} cells. Do not edit by hand. This file is data, not proof: every condition is",
        "re-checked in `MaxLeafFormal.Lobster`.",
        "-/",
        "",
        "set_option maxRecDepth 100000",
        "",
        "namespace MaxLeafFormal.Lobster",
        "",
        "/-- The rational `a / b`. -/",
        "def q (a b : ℕ) : ℚ := (a : ℚ) / b",
        "",
    ]
    CH = 500
    chunks = [out_cells[i:i + CH] for i in range(0, len(out_cells), CH)]
    for k, ch in enumerate(chunks):
        out += [
            f"/-- Cells `(a, b, w₀, w₁, w₂, λ, c)`, sorted by `a`: chunk {k}. -/",
            f"def cells{k} : List (ℚ × ℚ × ℚ × ℚ × ℚ × ℚ × ℚ) := [",
            ",\n".join(f"  ({q(a)}, {q(b)}, {q(w[0])}, {q(w[1])}, {q(w[2])}, {q(lam)}, {q(c)})"
                       for a, b, w, lam, c in ch),
            "]",
            "",
        ]
    out += [
        "/-- The chunks, in order. -/",
        "def chunks : List (List (ℚ × ℚ × ℚ × ℚ × ℚ × ℚ × ℚ)) := ["
        + ", ".join(f"cells{k}" for k in range(len(chunks))) + "]",
        "",
        f"/-- The right vector `v` at `y* = 0.3453`. -/",
        f"def vUp : ℚ × ℚ × ℚ := ({q(v[0])}, {q(v[1])}, {q(v[2])})",
        f"/-- `λ' = min_i (M v)_i / v_i`. -/",
        f"def lamUp : ℚ := {q(lamp)}",
        "",
        "end MaxLeafFormal.Lobster",
    ]
    print("\n".join(out))


if __name__ == "__main__":
    main()