#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
"""
Build formal/SilverSight/PIST/Matrices250.lean from
shared-data/rrc_pist_predictions_250_v1.json (or a local copy).

Each 8×8 matrix is emitted as its own named `def` to avoid overwhelming
Lean's kernel with a single large term. `findMatrix` is a `match` on the
equation_id string that delegates to the individual def.

Usage:
    python3 python/build_pist_matrices_250.py
    python3 python/build_pist_matrices_250.py \
        --predictions /path/to/rrc_pist_predictions_250_v1.json \
        --out-lean formal/SilverSight/PIST/Matrices250.lean
"""
from __future__ import annotations
import argparse, hashlib, json, re, sys
from pathlib import Path


def lean_matrix_def_name(eid: str) -> str:
    """Sanitize equation_id to a valid Lean identifier."""
    name = re.sub(r"[^a-zA-Z0-9_]", "_", eid)
    if name[0].isdigit():
        name = "m_" + name
    return name


def lean_array_int(xs: list[int]) -> str:
    return "#[" + ", ".join(str(x) for x in xs) + "]"


def lean_matrix_rows(rows: list[list[int]]) -> str:
    inner = ",\n    ".join(lean_array_int(r) for r in rows)
    return "#[\n    " + inner + "\n  ]"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate SilverSight PIST Matrices250.lean"
    )
    parser.add_argument(
        "--predictions",
        type=Path,
        default=Path("data/rrc_pist_predictions_250_v1.json"),
        help="Path to rrc_pist_predictions_250_v1.json",
    )
    parser.add_argument(
        "--out-lean",
        type=Path,
        default=Path("formal/SilverSight/PIST/Matrices250.lean"),
        help="Output Lean module path",
    )
    args = parser.parse_args()

    preds = json.loads(args.predictions.read_text())
    raw = preds.get("predictions", [])
    print(f"Loaded {len(raw)} predictions from {args.predictions}", file=sys.stderr)

    raw_sorted = sorted(raw, key=lambda p: p.get("object_id", p.get("equation_id", "")))

    defs: list[str] = []
    find_cases: list[str] = []
    for p in raw_sorted:
        eid = p.get("object_id") or p.get("equation_id", "")
        mat = p.get("matrix_8x8", [])
        if not eid or len(mat) != 8:
            print(f"WARNING: skipping malformed entry {p.get('equation_id', '?')}", file=sys.stderr)
            continue
        name = lean_matrix_def_name(eid)
        defs.append(f"def {name} : Array (Array Int) :=\n  {lean_matrix_rows(mat)}")
        find_cases.append(f"  | \"{eid}\" => some {name}")

    content_blob = json.dumps(
        {p.get("equation_id"): p.get("matrix_8x8") for p in raw_sorted},
        sort_keys=True, separators=(",", ":"),
    ).encode("utf-8")
    content_hash = hashlib.sha256(content_blob).hexdigest()

    lines = [
        "-- SilverSight.PIST.Matrices250 — AUTO-GENERATED by python/build_pist_matrices_250.py",
        "-- DO NOT EDIT BY HAND. Regenerate with:",
        "--   python3 python/build_pist_matrices_250.py",
        "--",
        "-- Source: rrc_pist_predictions_250_v1.json",
        "-- Content hash (SHA-256 canonical matrices): " + content_hash,
        "--",
        "-- This file is consumed by SilverSight.RRC.Q16_16Manifold at compile time",
        "-- to produce pistProxyLabel/pistExactLabel for each equation_id.",
        "",
        "namespace SilverSight.PIST.Matrices250",
        "",
        f"def pistMatrixDim : Nat := 8",
        "",
        *defs,
        "",
        "/-- Look up an 8×8 matrix by invariant equation_id. -/",
        "def findMatrix (eqId : String) : Option (Array (Array Int)) :=",
        "  match eqId with",
        *find_cases,
        "  | _ => none",
        "",
        "end SilverSight.PIST.Matrices250",
        "",
    ]

    args.out_lean.parent.mkdir(parents=True, exist_ok=True)
    args.out_lean.write_text("\n".join(lines))
    print(f"Wrote {args.out_lean} ({len(defs)} entries)", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())