#!/usr/bin/env python3
"""pist_matrix_builder — reproducible 8×8 braid adjacency matrix from RRC equations.

Builds token→strand adjacency matrices.  This is a feature-extraction shim
only — it produces no classifier output and no Lean spectral analysis.
proxy_pred/exact_pred are left null for the Lean surface.

Output schema: rrc_pist_predictions_250_v1  (claim_boundary: matrix-only)
"""

import hashlib
import json
import os
import re
import sys
from collections import defaultdict

N_STRANDS = 8

RECEIPT_JSON = os.path.join(
    os.path.dirname(__file__), "../..",
    "archive/experimental-shim-probes/rrc_equation_classifier_receipt.json",
)
OUTPUT_FILE = os.path.join(
    os.path.dirname(__file__), "../..",
    "shared-data/rrc_pist_predictions_250_v1.json",
)


# ═══════════════════════════════════════════════════════════════════════
# §1  PARSING
# ═══════════════════════════════════════════════════════════════════════

def load_equations() -> tuple[list[dict], int]:
    """Load equations grouped by invariant_receipt.object_id.

    Returns:
        (equations, total_source_records)
            equations — one per unique object_id with provenance fields
            total_source_records — count of all compiled_equations entries

    Representative selection (deterministic):
        Within each object_id group, the record with the lexicographically
        smallest ``equation_record.equation_id`` supplies the name and
        equation text for tokenization.  All records are captured in
        ``source_records`` for auditability.
    """
    with open(RECEIPT_JSON) as f:
        d = json.load(f)
    raw = d.get("compiled_equations", [])

    groups = defaultdict(list)
    for idx, eq in enumerate(raw):
        er = eq.get("equation_record", {})
        ir = eq.get("invariant_receipt", {})
        oid = ir.get("object_id", "")
        if not oid:
            continue
        groups[oid].append({
            "index": idx,
            "equation_id": er.get("equation_id", ""),
            "name": er.get("name", ""),
            "equation_text": er.get("equation", ""),
            "shape": ir.get("shape", "HoldForUnlawfulOrUnderspecifiedShape"),
            "status": ir.get("status", "HOLD"),
        })

    equations = []
    for oid, records in sorted(groups.items()):
        # Deterministic representative: lexicographically smallest equation_id
        rep = min(records, key=lambda r: r["equation_id"])
        equations.append({
            "equation_id": oid,
            "name": rep["name"],
            "equation_text": rep["equation_text"],
            "shape": rep["shape"],
            "status": rep["status"],
            "source_records": [
                {"equation_id": r["equation_id"], "name": r["name"]}
                for r in sorted(records, key=lambda r: r["equation_id"])
            ],
        })

    return equations, len(raw)


def tokenize(text: str) -> list[str]:
    """Split text into tokens.

    Normalization rules (versioned, deterministic):
      1. Replace '-' with '_'
      2. Split on '_' and ':'
      3. Drop empty strings

    Token source: ``equation_record.name`` (the human-readable equation title).
    The full equation text is preserved in ``equation_record.equation`` but is
    not the tokenization source.
    """
    normalized = text.replace("-", "_")
    return [p for p in re.split(r"[_:]", normalized) if p]


# ═══════════════════════════════════════════════════════════════════════
# §2  VOCABULARY
# ═══════════════════════════════════════════════════════════════════════

def build_global_vocabulary(equations: list[dict]) -> dict[str, int]:
    """Build global token→index mapping (sorted, deterministic)."""
    all_tokens = set()
    for eq in equations:
        for t in tokenize(eq["name"]):
            all_tokens.add(t)
    sorted_tokens = sorted(all_tokens)
    return {t: i for i, t in enumerate(sorted_tokens)}


def global_vocab_hash(vocab: dict[str, int]) -> str:
    """SHA256 of ``|``-joined sorted vocabulary tokens."""
    joined = "|".join(vocab.keys())
    return hashlib.sha256(joined.encode("utf-8")).hexdigest()


# ═══════════════════════════════════════════════════════════════════════
# §3  STRAND ADJACENCY MATRIX
# ═══════════════════════════════════════════════════════════════════════

def strand_for_token(token: str, vocab: dict[str, int]) -> int:
    """Map token to strand index via global vocabulary position (mod 8)."""
    return vocab[token] % N_STRANDS


def build_adjacency_matrix(tokens: list[str], vocab: dict[str, int]) -> list[list[int]]:
    """Build 8×8 adjacency matrix from token ordering.

    For each adjacent pair (t_i, t_{i+1}) in the original token sequence,
    increment ``M[strand(t_i)][strand(t_{i+1})]`` by 1.

    No symmetrization, no operator projection, no diagonal self-crossings.
    """
    M = [[0] * N_STRANDS for _ in range(N_STRANDS)]
    strands = [strand_for_token(t, vocab) for t in tokens]
    for i in range(len(strands) - 1):
        s1, s2 = strands[i], strands[i + 1]
        M[s1][s2] += 1
    return M


# ═══════════════════════════════════════════════════════════════════════
# §4  HERMITE POLYNOMIAL KERNEL (RRC sieve refinement)
#
# The generalized bilinear generating function (Mehler kernel) lifts the
# RRC similarity metric from linear (bigram frequency) to polynomial:
#
#     K_ij = H_{p,q}(M[i·], M[j·] | ρ)
#
# where M[i·] is the i-th row of the adjacency matrix, H_{p,q} is the
# generalized Hermite–Kampé de Fériet polynomial (Giani et al. 2025, Lemma 1),
# and ρ ∈ (0,1) is a coupling parameter.
#
# For p=q=0 (standard Mehler):
#     K_ij = exp(2·M[i]·M[j]·ρ - (|M[i]|²+|M[j]|²)·ρ²) / √(1-ρ²)
#
# This is the simplest non-trivial kernel and already captures algebraic
# variety structure beyond linear adjacency.
# ═══════════════════════════════════════════════════════════════════════

_HERMITE_RHO: float = 0.5  # coupling parameter (tunable, 0 < ρ < 1)


def _dot(a: list[int], b: list[int]) -> int:
    """Dot product of two integer vectors."""
    return sum(x * y for x, y in zip(a, b))


def _norm2(a: list[int]) -> int:
    """Squared Euclidean norm of integer vector."""
    return sum(x * x for x in a)


def mehler_kernel_row(i: int, M: list[list[int]], rho: float) -> list[float]:
    """Compute Mehler kernel row K[i·] for adjacency matrix M.

    K_ij = exp(2·M[i]·M[j]·ρ - (|M[i]|²+|M[j]|²)·ρ²) / √(1-ρ²)

    This is the generalized bilinear generating function G from
    Giani et al. (2025) eq. (5) with p=q=0, x=M[i], y=M[j].
    """
    import math
    n = len(M)
    row_i = M[i]
    norm_i = _norm2(row_i)
    denom = math.sqrt(1.0 - rho * rho)
    result = []
    for j in range(n):
        row_j = M[j]
        inner = 2.0 * _dot(row_i, row_j) * rho - (norm_i + _norm2(row_j)) * rho * rho
        result.append(math.exp(inner) / denom)
    return result


def build_hermite_kernel(M: list[list[int]]) -> list[list[float]]:
    """Build full 8×8 Hermite kernel matrix from adjacency M.

    Returns K[i][j] = mehler_kernel_row(i, M, _HERMITE_RHO)[j]
    """
    n = len(M)
    return [mehler_kernel_row(i, M, _HERMITE_RHO) for i in range(n)]


def canonical_hermite_json(K: list[list[float]]) -> str:
    """Row‑major JSON, fixed 8×8 nesting, scientific notation for floats."""
    return json.dumps(K, separators=(",", ":"))


def hermite_hash(K: list[list[float]]) -> str:
    return hashlib.sha256(canonical_hermite_json(K).encode("utf-8")).hexdigest()


# ═══════════════════════════════════════════════════════════════════════
# §5  HASHING
# ═══════════════════════════════════════════════════════════════════════

def canonical_matrix_json(M: list[list[int]]) -> str:
    """Row‑major JSON, integers only, no whitespace, fixed 8×8 nesting."""
    return json.dumps(M, separators=(",", ":"))


def matrix_hash(M: list[list[int]]) -> str:
    return hashlib.sha256(canonical_matrix_json(M).encode("utf-8")).hexdigest()


# ═══════════════════════════════════════════════════════════════════════
# §5  MAIN
# ═══════════════════════════════════════════════════════════════════════

def main() -> int:
    equations, total_source_records = load_equations()
    print(f"Loaded {len(equations)} unique equations "
          f"(from {total_source_records} source records)", flush=True)

    vocab = build_global_vocabulary(equations)
    gvh = global_vocab_hash(vocab)
    print(f"Global vocabulary: {len(vocab)} unique tokens  hash={gvh}", flush=True)

    predictions = []
    hash_counts = {}
    hermite_hash_counts = {}

    for eq in equations:
        tokens = tokenize(eq["name"])
        M = build_adjacency_matrix(tokens, vocab)
        K = build_hermite_kernel(M)
        mh = matrix_hash(M)
        kh = hermite_hash(K)
        hash_counts[mh] = hash_counts.get(mh, 0) + 1
        hermite_hash_counts[kh] = hermite_hash_counts.get(kh, 0) + 1
        predictions.append({
            "equation_id": eq["equation_id"],
            "proxy_pred": None,
            "exact_pred": None,
            "matrix_hash": mh,
            "matrix_8x8": M,
            "hermite_kernel_hash": kh,
            "hermite_kernel_8x8": K,
            "hermite_kernel_rho": _HERMITE_RHO,
            "source_records": eq["source_records"],
            "notes": "generated from equation_record.equation_id token adjacency; "
                     "hermite kernel via generalized bilinear generating function (ρ=0.5)",
        })

    n_unique_adj = len(hash_counts)
    n_unique_hermite = len(hermite_hash_counts)
    n_total = len(predictions)
    n_collisions_adj = sum(1 for c in hash_counts.values() if c > 1)
    n_collisions_hermite = sum(1 for c in hermite_hash_counts.values() if c > 1)
    print(f"Matrix hash collisions (adjacency): {n_collisions_adj} groups "
          f"({n_unique_adj} unique / {n_total} total)", flush=True)
    print(f"Matrix hash collisions (hermite):   {n_collisions_hermite} groups "
          f"({n_unique_hermite} unique / {n_total} total)", flush=True)
    if n_collisions_hermite < n_collisions_adj:
        print("Hermite kernel DISAMBIGUATES more equations than raw adjacency: "
              "the algebraic variety metric is strictly finer.", flush=True)
    elif n_collisions_hermite == n_collisions_adj:
        print("Hermite kernel matches adjacency resolution at ρ={}.".format(_HERMITE_RHO),
              flush=True)
    else:
        print("Hermite kernel merges some distinct equations — "
              "consider tuning ρ.", flush=True)

    artifact = {
        "schema": "rrc_pist_predictions_250_v1",
        "claim_boundary": "matrix-only;hermite-kernel-added;no-classifier;no-lean-spectral",
        "matrix_schema": "token_strand_adjacency_8x8_v1",
        "hermite_kernel_schema": "mehler_kernel_8x8_v1_generalized_bilinear_generating_function",
        "hermite_kernel_rho": _HERMITE_RHO,
        "global_vocab_hash": gvh,
        "summary": {
            "total_source_records": total_source_records,
            "unique_equation_ids": n_total,
            "unique_matrix_hashes": n_unique_adj,
            "unique_hermite_hashes": n_unique_hermite,
            "hermite_disambiguation": n_collisions_adj - n_collisions_hermite,
        },
        "predictions": predictions,
    }

    with open(OUTPUT_FILE, "w") as f:
        json.dump(artifact, f, indent=2, sort_keys=True)
    print(f"\nWrote {OUTPUT_FILE}", flush=True)

    return 0


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