"""Export the two RS-theta antinorm certificates to a Lean data file.

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

For each certificate ancillary/cert_final_{0,1}.json this writes the point set Z, the
parameters y, e1, e2, m0 and, for every condition, the exact convex coefficients that the exact
verifier ancillary/verify_cert.py uses: each stored float is rationalised with
Fraction.limit_denominator(10**15) and the vector is renormalised to sum 1 (`ratmu`).

The generated file is data, not proof: Lean re-checks every condition in the kernel, so a wrong
export can only make a check fail.
"""
import hashlib
import json
import os
import sys
from fractions import Fraction as Fr

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def ratmu(lst):
    # identical to verify_cert.ratmu
    mu = [(k, Fr(v).limit_denominator(10 ** 15)) for k, v in lst]
    s = sum(v for _, v in mu)
    return [(k, v / s) for k, v in mu]


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


def mu_lean(mu):
    return "[" + ", ".join(f"({k}, {q(v)})" for k, v in mu) + "]"


def export(i, path):
    raw = open(path, "rb").read()
    sha = hashlib.sha256(raw).hexdigest()
    cert = json.loads(raw)
    Z = [tuple(Fr(*v) for v in z) for z in cert["Z"]]
    n = len(Z)
    m0 = cert["m0"]
    out = []
    out.append(f"/-! Certificate {i}: `{os.path.relpath(path, ROOT)}`, SHA-256 `{sha}`, |Z| = {n}, m0 = {m0}. -/")
    out.append("")
    out.append(f"def y{i} : ℚ := {q(Fr(*cert['y']))}")
    out.append(f"def e1_{i} : ℚ := {q(Fr(*cert['e1']))}")
    out.append(f"def e2_{i} : ℚ := {q(Fr(*cert['e2']))}")
    out.append(f"def m0_{i} : ℕ := {m0}")
    out.append("")
    out.append(f"def Z{i} : List (ℚ × ℚ × ℚ) := [")
    out.append(",\n".join(f"  ({q(z[0])}, {q(z[1])}, {q(z[2])})" for z in Z))
    out.append("]")
    out.append("")
    out.append(f"/-- Convex coefficients for the pair conditions, row-major: entry `i * {n} + j` is for `(i, j)`. -/")
    out.append(f"def pairs{i} : List (List (ℕ × ℚ)) := [")
    rows = []
    for a in range(n):
        for b in range(n):
            rows.append("  " + mu_lean(ratmu(cert["combos"][f"{a},{b}"])))
    out.append(",\n".join(rows))
    out.append("]")
    out.append("")
    out.append(f"/-- Convex coefficients for `Init(j)`, `j = 0, …, {m0}`. -/")
    out.append(f"def inits{i} : List (List (ℕ × ℚ)) := [")
    out.append(",\n".join("  " + mu_lean(ratmu(cert["inits"][str(m)])) for m in range(m0 + 1)))
    out.append("]")
    out.append("")
    out.append(f"/-- Convex coefficients for the tail vector `(y/2, 0, 0)` at scale `e1 * e2 ^ {m0 + 1}`. -/")
    out.append(f"def tail{i} : List (ℕ × ℚ) := {mu_lean(ratmu(cert['inits']['inf']))}")
    out.append("")
    return "\n".join(out)


def main():
    header = [
        "import Mathlib",
        "",
        "/-!",
        "# RS-θ certificate data (generated)",
        "",
        "Generated by `formal/scripts/export_rstheta_cert.py`; do not edit by hand.",
        "This file is data, not proof: every condition is re-checked in `MaxLeafFormal.Antinorm3`.",
        "-/",
        "",
        "set_option maxRecDepth 100000",
        "",
        "namespace MaxLeafFormal.RSThetaCert",
        "",
        "/-- The rational `a / b`. -/",
        "def q (a b : ℕ) : ℚ := (a : ℚ) / b",
        "",
    ]
    body = [export(i, os.path.join(ROOT, "ancillary", f"cert_final_{i}.json")) for i in (0, 1)]
    print("\n".join(header) + "\n".join(body) + "\nend MaxLeafFormal.RSThetaCert")


if __name__ == "__main__":
    main()