#!/usr/bin/env python3
"""
rrc_affine_conservation_probe.py — Does math notation obey the affine Ã₂ δ law?

The operator grammar's 3-cycle relation–binary_op–arrow is the affine Ã₂
extended-Dynkin diagram. Its Cartan matrix [[2,-1,-1],[-1,2,-1],[-1,-1,2]] has
null vector δ=(1,1,1); the associated quadratic form is

    Q(r,b,a) = (r−b)² + (b−a)² + (a−r)²        (imbalance from the δ direction)

CONSERVATION HYPOTHESIS: well-formed notation holds a *conserved ratio* among the
three cycle-roles — i.e. each equation sits near a fixed point of the cycle
simplex, so Q_norm = Q/T² is small and tightly distributed, and outliers flag
malformation. This is FALSIFIABLE: we test the observed Q against two null models.

  Null-δ      : multinomial(T, (1/3,1/3,1/3))  — tests if equations sit at δ.
  Null-margin : multinomial(T, corpus marginal) — tests if per-equation ratios
                are TIGHTER than random draws at the population average (i.e.
                whether there is a per-equation conservation constraint at all).

Verdict: conservation holds iff observed dispersion ≪ Null-margin dispersion.

Usage:  python3 4-Infrastructure/shim/rrc_affine_conservation_probe.py
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import rrc_root_system_probe as P            # noqa: E402
from math_symbols import CHAR_INFO           # noqa: E402

ROOT = Path("/home/allaun/Research Stack")
OUT = ROOT / "shared-data/data/rrc_affine_conservation_receipt.json"
CYCLE = ["relation", "binary_op", "arrow"]   # the affine Ã₂ cycle roles
RNG = np.random.default_rng(20260618)


def q_norm(vec: np.ndarray) -> float:
    r, b, a = vec
    T = r + b + a
    if T <= 0:
        return 0.0
    return ((r - b) ** 2 + (b - a) ** 2 + (a - r) ** 2) / (T * T)


def null_q(T: int, p: np.ndarray, n: int = 300) -> float:
    """Mean Q_norm of n multinomial(T, p) draws."""
    draws = RNG.multinomial(T, p, size=n).astype(float)
    return float(np.mean([q_norm(d) for d in draws]))


def main() -> None:
    eqs = P.load_equations()
    roles = sorted({i["role"] for i in CHAR_INFO.values()})
    idx = [roles.index(c) for c in CYCLE]
    M = np.array([P.role_vector(e, roles)[idx] for e in eqs])  # (N,3) cycle counts

    T = M.sum(axis=1)
    keep = T >= 2                       # operator-bearing equations only
    Mk, eqk = M[keep], [e for e, k in zip(eqs, keep) if k]
    Tk = Mk.sum(axis=1)
    N = len(Mk)

    marginal = Mk.sum(axis=0) / Mk.sum()
    arrow_prev = float((Mk[:, 2] > 0).mean())
    centroid = (Mk / Tk[:, None]).mean(axis=0)

    obs_Q = np.array([q_norm(v) for v in Mk])
    null_delta = np.array([null_q(int(t), np.array([1/3, 1/3, 1/3])) for t in Tk])
    null_marg = np.array([null_q(int(t), marginal) for t in Tk])

    cons_strength = float(null_marg.mean() / obs_Q.mean()) if obs_Q.mean() > 0 else float("inf")

    # anomaly = how far an equation's Q sits above the corpus median (robust z)
    med, mad = np.median(obs_Q), np.median(np.abs(obs_Q - np.median(obs_Q))) + 1e-9
    z = (obs_Q - med) / (1.4826 * mad)
    order = np.argsort(-z)

    print("=" * 68)
    print(f"AFFINE Ã₂ δ-CONSERVATION PROBE — {N} operator-bearing equations")
    print("=" * 68)
    print(f"cycle roles (relation, binary_op, arrow)")
    print(f"  corpus marginal ratio : {np.round(marginal,3)}")
    print(f"  simplex centroid      : {np.round(centroid,3)}   (δ = [0.333 0.333 0.333])")
    print(f"  ‖centroid − δ‖        : {np.linalg.norm(centroid-np.array([1/3]*3)):.3f}")
    print(f"  arrow prevalence      : {arrow_prev:.1%} of equations have any arrow")
    print(f"\n  observed   mean Q_norm: {obs_Q.mean():.4f}  (std {obs_Q.std():.4f})")
    print(f"  Null-δ     mean Q_norm: {null_delta.mean():.4f}")
    print(f"  Null-margin mean Q_norm: {null_marg.mean():.4f}")
    print(f"\n  >>> conservation strength (Null-margin / observed): {cons_strength:.2f}")
    verdict = ("CONSERVED: equations hold a tighter ratio than chance"
               if cons_strength > 1.15 else
               "NOT CONSERVED: observed ≈ random at the marginal — no per-eq law")
    print(f"  >>> VERDICT: {verdict}")
    print(f"\n  Top-5 δ-imbalance anomalies (validity-check candidates):")
    for i in order[:5]:
        r, b, a = Mk[i].astype(int)
        print(f"    z={z[i]:5.1f}  (rel={r},bin={b},arr={a})  {eqk[i][:54]}")

    OUT.write_text(json.dumps({
        "schema": "rrc_affine_conservation_v1",
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "n_equations": N,
        "cycle_roles": CYCLE,
        "marginal_ratio": [float(x) for x in marginal],
        "simplex_centroid": [float(x) for x in centroid],
        "centroid_to_delta": float(np.linalg.norm(centroid - np.array([1/3]*3))),
        "arrow_prevalence": arrow_prev,
        "observed_mean_Q": float(obs_Q.mean()),
        "null_delta_mean_Q": float(null_delta.mean()),
        "null_marginal_mean_Q": float(null_marg.mean()),
        "conservation_strength": cons_strength,
        "verdict": verdict,
        "top_anomalies": [
            {"z": float(z[i]), "cycle": Mk[i].astype(int).tolist(), "eq": eqk[i][:80]}
            for i in order[:10]
        ],
    }, indent=2, ensure_ascii=False))
    print(f"\nReceipt: {OUT}")


if __name__ == "__main__":
    main()