#!/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).

Python's role: read the predictions artifact, serialize each 8×8 matrix as a
Lean Array (Array Int) literal, emit an association list keyed by rrc_eq_<hex>.

The generated Lean file is consumed by SilverSight.PIST.Classify
(classifyProxy/classifyExact) to produce pistProxyLabel/pistExactLabel.

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, sys
from pathlib import Path


def lean_str(s: str) -> str:
    s = s.replace("\\", "\\\\").replace('"', '\\"')
    return f'"{s}"'


def lean_array_int(xs: list[int]) -> str:
    inner = ", ".join(str(x) for x in xs)
    return f"#[{inner}]"


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


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate SilverSight PIST Matrices250.lean"
    )
    parser.add_argument(
        "--predictions",
        type=Path,
        default=Path("/home/allaun/Research Stack/shared-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)

    # Deterministic order by equation_id for reproducible builds.
    raw_sorted = sorted(raw, key=lambda p: p.get("equation_id", ""))

    entries: list[str] = []
    for p in raw_sorted:
        eid = p.get("equation_id", "")
        mat = p.get("matrix_8x8", [])
        if not eid or len(mat) != 8:
            print(f"WARNING: skipping malformed entry {eid}", file=sys.stderr)
            continue
        matrix_lean = lean_matrix_rows(mat)
        entries.append(f"    ({lean_str(eid)}, {matrix_lean})")

    # Include a content hash so the generated file can be validated.
    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.PIST.Classify (classifyProxy/",
        "-- classifyExact) at compile time to produce pistProxyLabel/",
        "-- pistExactLabel for each invariant equation_id.",
        "",
        "namespace SilverSight.PIST.Matrices250",
        "",
        "/-- All matrices in this list have dimension 8×8 (the canonical strand count).",
        "    Stored as association list (key → matrix). Generated from",
        "    rrc_pist_predictions_250_v1.json.",
        f"    Entries: {len(entries)}. -/",
        f"def pistMatrixDim : Nat := 8",
        "",
        "def pistMatrices250 : List (String × Array (Array Int)) :=",
        "  [",
        ",\n".join(entries),
        "  ]",
        "",
        "/-- Look up an 8×8 matrix by invariant equation_id. -/",
        "def findMatrix (eqId : String) : Option (Array (Array Int)) :=",
        "  pistMatrices250.find? (fun (k, _) => k = eqId) |>.map (fun (_, v) => v)",
        "",
        "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(entries)} entries)", file=sys.stderr)
    return 0


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