#!/usr/bin/env python3
"""Regenerate rrc_pist_predictions_250_v1.json from the classifier receipt.

Reads all 278 compiled equations, builds 8×8 strand adjacency matrices,
deduplicates by equation_id per the AGENTS.md merge rules, and writes
the predictions artifact.

Usage:
    python3 python/generate_predictions.py
    python3 python/generate_predictions.py \
        --receipt /path/to/rrc_equation_classifier_receipt.json \
        --output /path/to/rrc_pist_predictions_250_v1.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from pathlib import Path
from datetime import datetime, timezone
from collections import OrderedDict


def tokenize(text: str) -> list[str]:
    text = text.lower()
    return re.findall(r"[a-z0-9]+|[+\-*/=^(){}\[\]]", text)


def build_matrix(tokens: list[str]) -> list[list[int]]:
    vocab = sorted(set(tokens))
    matrix = [[0 for _ in range(8)] for _ in range(8)]
    for t_i, t_j in zip(tokens, tokens[1:]):
        matrix[vocab.index(t_i) % 8][vocab.index(t_j) % 8] += 1
    return matrix


def canonical_json_matrix(matrix: list[list[int]]) -> str:
    return json.dumps(matrix, separators=(",", ":"))


def build_record(text: str, equation_id: str, object_id: str, name: str) -> dict:
    tokens = tokenize(text)
    vocab = sorted(set(tokens))
    matrix = build_matrix(tokens)
    matrix_json = canonical_json_matrix(matrix)
    return {
        "equation_id": equation_id,
        "object_id": object_id,
        "name": name,
        "equation": text,
        "schema": "rrc_pist_predictions_250_v1",
        "claim_boundary": "matrix-only;no-classifier;no-lean-spectral",
        "proxy_pred": None,
        "exact_pred": None,
        "matrix_schema": "token_strand_adjacency_8x8_v1",
        "matrix_hash": hashlib.sha256(matrix_json.encode("utf-8")).hexdigest(),
        "matrix_8x8": matrix,
        "global_vocab_hash": hashlib.sha256(
            json.dumps(vocab, separators=(",", ":")).encode("utf-8")
        ).hexdigest(),
        "source_records": [],
    }


def main():
    parser = argparse.ArgumentParser(description="Regenerate PIST predictions")
    parser.add_argument("--receipt", type=Path,
                        default=Path("/home/allaun/Research Stack/archive/experimental-shim-probes/rrc_equation_classifier_receipt.json"))
    parser.add_argument("--output", type=Path,
                        default=Path("/home/allaun/Research Stack/shared-data/rrc_pist_predictions_250_v1.json"))
    args = parser.parse_args()

    receipt = json.loads(args.receipt.read_text())
    equations = receipt["compiled_equations"]

    print(f"Loaded {len(equations)} source equations", file=sys.stderr)

    # Dedup by object_id (invariant receipt) which matches the manifold's lookup key.
    merged: dict[str, dict] = OrderedDict()
    for eq in equations:
        rec = eq.get("equation_record", eq)
        ir = eq.get("invariant_receipt", {})
        eid = ir.get("object_id", rec["equation_id"])  # use object_id for lookup match
        name = rec.get("name", eid)
        text = rec.get("equation", "")

        if eid not in merged:
            oid = ir.get("object_id", rec["equation_id"])
            merged[eid] = build_record(text, rec["equation_id"], oid, name)
            merged[eid]["source_records"] = []

        merged[eid]["source_records"].append({
            "equation_record_id": rec["equation_id"],
            "object_id": ir.get("object_id", rec["equation_id"]),
            "name": name,
        })

    predictions = list(merged.values())

    # Build global vocab
    all_tokens = set()
    for p in predictions:
        all_tokens.update(tokenize(p.get("equation", "")))
    global_vocab = sorted(all_tokens)

    output = {
        "schema": "rrc_pist_predictions_250_v1",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "claim_boundary": "matrix-only;no-classifier;no-lean-spectral",
        "global_vocab_hash": hashlib.sha256(
            json.dumps(global_vocab, separators=(",", ":")).encode("utf-8")
        ).hexdigest(),
        "global_vocab": global_vocab,
        "matrix_schema": "token_strand_adjacency_8x8_v1",
        "summary": {
            "total_source_records": len(equations),
            "unique_equation_ids": len(predictions),
        },
        "predictions": predictions,
    }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    print(f"Wrote {args.output} ({len(predictions)} unique predictions from {len(equations)} source records)", file=sys.stderr)


if __name__ == "__main__":
    main()