#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///
"""
Build Semantics/PIST/Matrices278.lean from rrc_pist_predictions_278_v1.json.

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

Generated Lean file is consumed by Semantics.PIST.Classify (classifyProxy/
classifyExact) when labels are computed in v2+.

Usage:
    python3 4-Infrastructure/shim/build_pist_matrices_278.py

Dependencies:
    shared-data/rrc_pist_predictions_278_v1.json  (pist_matrix_builder.py output)
"""
from __future__ import annotations
import json, sys
from pathlib import Path

ROOT = Path("/home/allaun/Research Stack")
PREDICTIONS_JSON = ROOT / "shared-data/rrc_pist_predictions_278_v1.json"
OUT_LEAN = ROOT / "0-Core-Formalism/lean/Semantics/Semantics/PIST/Matrices278.lean"


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


def lean_array_int(xs: list[int]) -> str:
    """Format a list of ints as a Lean Array literal."""
    inner = ", ".join(str(x) for x in xs)
    return f"#[{inner}]"


def lean_matrix_rows(rows: list[list[int]]) -> str:
    """Format an 8×8 matrix as a Lean Array (Array Int) literal."""
    inner = ",\n    ".join(lean_array_int(r) for r in rows)
    return f"#[\n    {inner}\n  ]"


def main() -> int:
    preds = json.loads(PREDICTIONS_JSON.read_text())
    raw = preds.get("predictions", [])
    print(f"Loaded {len(raw)} predictions from {PREDICTIONS_JSON}", file=sys.stderr)

    lines: list[str] = []
    lines.append("-- Semantics.PIST.Matrices278 — AUTO-GENERATED by build_pist_matrices_278.py")
    lines.append("-- DO NOT EDIT BY HAND.  Regenerate with:")
    lines.append("--   python3 4-Infrastructure/shim/build_pist_matrices_278.py")
    lines.append("--")
    lines.append("-- Source: shared-data/rrc_pist_predictions_278_v1.json")
    lines.append("-- (generated by 4-Infrastructure/shim/pist_matrix_builder.py)")
    lines.append("--")
    lines.append("-- This file is consumed by Semantics.PIST.Classify (classifyProxy/")
    lines.append("-- classifyExact) at compile time to produce pistProxyLabel/")
    lines.append("-- pistExactLabel for each invariant equation_id.")
    lines.append("")
    lines.append("namespace Semantics.PIST.Matrices278")
    lines.append("")
    lines.append("/-- 8×8 braid adjacency matrices keyed by invariant equation_id, stored as")
    lines.append("    an association list (key → matrix).  250 entries; linear lookup is fine.")
    lines.append("    Generated from rrc_pist_predictions_278_v1.json. -/")
    lines.append(f"def pistMatrices278 : List (String × Array (Array Int)) :=")
    lines.append("  [")

    entry_strs: list[str] = []
    for p in raw:
        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)
        entry_strs.append(f"    ({lean_str(eid)}, {matrix_lean})")

    lines.append(",\n".join(entry_strs))
    lines.append("  ]")
    lines.append("")
    lines.append("/-- Look up an 8×8 matrix by invariant equation_id. -/")
    lines.append("def findMatrix (eqId : String) : Option (Array (Array Int)) :=")
    lines.append("  pistMatrices278.find? (fun (k, _) => k = eqId) |>.map (fun (_, v) => v)")
    lines.append("")
    lines.append("end Semantics.PIST.Matrices278")
    lines.append("")

    OUT_LEAN.write_text("\n".join(lines))
    print(f"Wrote {OUT_LEAN} ({len(entry_strs)} entries)", file=sys.stderr)
    return 0


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