"""Export the three Directed Steiner Tree certificates to a Lean data file.

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

Reads ancillary/dst_certificates.json and writes, for each certificate
j = 0, 1, 2, the parameters y, e1, e2, the point list Z and, for every membership condition of
(C1)-(C3), the witness that dst_certificate_verify.member finds: a point z_k of Z (coefficients
[(k, 1)]) or two points z_a, z_b and the rational mu of `mu_interval` (coefficients
[(a, mu), (b, 1 - mu)]).  It also writes the weights lambda_j.

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

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SRC = os.path.join(ROOT, "ancillary", "dst_certificates.json")


def comp(P, Q):
    return (P[0] * (Q[0] + Q[1]), P[1] * (Q[0] / 2 + Q[1]))


def mu_interval(z, w, P):
    # identical to dst_certificate_verify.mu_interval
    lo, hi = F(0), F(1)
    for k in range(2):
        d, r = z[k] - w[k], P[k] - w[k]
        if d > 0:
            hi = min(hi, r / d)
        elif d < 0:
            lo = max(lo, r / d)
        elif r < 0:
            return None
    return lo if lo <= hi else None


def witness(P, Z):
    """The witness of dst_certificate_verify.member, as index/coefficient pairs."""
    for k, z in enumerate(Z):
        if z[0] <= P[0] and z[1] <= P[1]:
            return [(k, F(1))]
    order = sorted(range(len(Z)), key=lambda k: Z[k])
    for a, b in zip(order, order[1:]):
        mu = mu_interval(Z[a], Z[b], P)
        if mu is not None:
            return [(a, mu), (b, 1 - mu)]
    for a in range(len(Z)):
        for b in range(a + 1, len(Z)):
            mu = mu_interval(Z[a], Z[b], P)
            if mu is not None:
                return [(a, mu), (b, 1 - mu)]
    raise SystemExit(f"no witness for {P}")


def q(x):
    x = F(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 main():
    sys.stdout.reconfigure(encoding="utf-8", newline="\n")
    raw = open(SRC, "rb").read()
    # Hash LF-normalised bytes, so that the header does not depend on the checkout's line endings.
    sha = hashlib.sha256(raw.replace(b"\r\n", b"\n")).hexdigest()
    data = json.loads(raw)
    certs = data["certificates"]
    lam = [F(s) for s in data["lambda"]]
    assert len(certs) == 3 and len(lam) == 3
    out = [
        "import Mathlib",
        "",
        "/-!",
        "# Directed Steiner Tree certificate data (generated)",
        "",
        "Generated by `formal/scripts/export_dst_cert.py` from",
        f"`ancillary/dst_certificates.json`, SHA-256 `{sha}`.",
        "Do not edit by hand. This file is data, not proof: every condition is re-checked in",
        "`MaxLeafFormal.DSTCert`.",
        "-/",
        "",
        "set_option maxRecDepth 100000",
        "",
        "namespace MaxLeafFormal.DSTCert",
        "",
        "/-- The rational `a / b`. -/",
        "def q (a b : ℕ) : ℚ := (a : ℚ) / b",
        "",
    ]
    ys, e1s, e2s = [], [], []
    for j, c in enumerate(certs):
        y, e1, e2 = F(c["y"]), F(c["e1"]), F(c["e2"])
        Z = [(F(a), F(b)) for a, b in c["Z"]]
        n = len(Z)
        ys.append(y), e1s.append(e1), e2s.append(e2)
        out.append(f"/-! Certificate {j}: |Z| = {n}. -/")
        out.append("")
        out.append(f"def Z{j} : List (ℚ × ℚ) := [")
        out.append(",\n".join(f"  ({q(z[0])}, {q(z[1])})" for z in Z))
        out.append("]")
        out.append("")
        out.append(f"/-- (C1) witnesses, row-major: entry `a * {n} + b` is for `z_a ∘ z_b`. -/")
        out.append(f"def pairs{j} : List (List (ℕ × ℚ)) := [")
        out.append(",\n".join("  " + mu_lean(witness(comp(z, w), Z)) for z in Z for w in Z))
        out.append("]")
        out.append("")
        out.append(f"/-- (C2) witness for `(1, 0) / e2`. -/")
        out.append(f"def c2w{j} : List (ℕ × ℚ) := {mu_lean(witness((1 / e2, F(0)), Z))}")
        out.append("")
        out.append(f"/-- (C3) witnesses: entry `a` is for `((y, 1) ∘ z_a) / e1`. -/")
        out.append(f"def c3w{j} : List (List (ℕ × ℚ)) := [")
        out.append(",\n".join(
            "  " + mu_lean(witness(tuple(x / e1 for x in comp((y, F(1)), z)), Z)) for z in Z))
        out.append("]")
        out.append("")
    vec = lambda xs: "![" + ", ".join(q(x) for x in xs) + "]"
    out.append(f"def y : Fin 3 → ℚ := {vec(ys)}")
    out.append(f"def e₁ : Fin 3 → ℚ := {vec(e1s)}")
    out.append(f"def e₂ : Fin 3 → ℚ := {vec(e2s)}")
    out.append(f"def lam : Fin 3 → ℚ := {vec(lam)}")
    out.append("def Zq : Fin 3 → List (ℚ × ℚ) := ![Z0, Z1, Z2]")
    out.append("def pairs : Fin 3 → List (List (ℕ × ℚ)) := ![pairs0, pairs1, pairs2]")
    out.append("def c2w : Fin 3 → List (ℕ × ℚ) := ![c2w0, c2w1, c2w2]")
    out.append("def c3w : Fin 3 → List (List (ℕ × ℚ)) := ![c3w0, c3w1, c3w2]")
    out.append("")
    out.append("end MaxLeafFormal.DSTCert")
    print("\n".join(out))


if __name__ == "__main__":
    main()