"""geolip_vitals.py — the shared diagnostic harness (progression plan Tree 7a).
One implementation, every tree imports it. ALL functions are READOUTS: no gradients,
no losses. CV is a readout, never a force (discovery_catalog #3). Judge addressing by
drift->0.29154 and CV->0.20, never recon cosine (MANIFEST).

Vitals provided:
  anchor_drift        — geodesic drift of anchors from init; binding fraction @0.29154
  pentachoron_cv      — CM 4-volume CV over random 5-row subsets (geovocab2 import)
  axis_aliveness      — oriented-address usage: axes alive, hppl, collapse flag
  gate_stats          — gate means vs the 0.012-0.03 band
  path_diversity      — unique-path counting, FIXED high-bits hash (low-16 bug is the
                        retracted artifact — never use the low bits)
  grad_norm_spread    — gradient democracy monitor (orders-of-magnitude spread)
  CVScreen            — CV@1000-batch early band screen (<0.30 LOW / .35-.50 MID / >.80 HIGH)

Smoke on a torch-capable env:  python geolip_vitals.py
"""
from __future__ import annotations
import math
import torch

BINDING = 0.29154            # radians; the binding/separation constant (MANIFEST)
CV_BAND = (0.13, 0.30)       # CM CV band (discovery_catalog #4)
GATE_BAND = (0.012, 0.03)    # live invariant candidate (acd_campaign)
KNUTH32 = 2654435761


# ----------------------------------------------------------------------------- drift
@torch.no_grad()
def anchor_drift(current: torch.Tensor, init: torch.Tensor, tol: float = 0.05) -> dict:
    """Geodesic drift (radians) of each row of `current` from its row in `init`,
    both row-normalized. Returns mean/std/per-row drift and the fraction of rows
    within +/-tol of BINDING (the GLFM '46%' readout)."""
    a = torch.nn.functional.normalize(current.float(), dim=-1)
    b = torch.nn.functional.normalize(init.float(), dim=-1)
    cos = (a * b).sum(-1).clamp(-1.0, 1.0)
    drift = torch.arccos(cos)
    frac = ((drift - BINDING).abs() <= tol).float().mean()
    return {"mean": drift.mean().item(), "std": drift.std().item(),
            "per_row": drift, "binding_fraction": frac.item()}


# -------------------------------------------------------------------------------- cv
@torch.no_grad()
def _pentachoron_volumes(pts: torch.Tensor) -> torch.Tensor:
    """Batched Cayley-Menger 4-simplex volumes. pts: (B, 5, D) -> (B,) volumes.
    One float64 det over all samples (vol^2 = -det(CM)/9216 for n=4). Built-in
    per the 2026-07-11 rider amendment (the operator: the installed geovocab2 path is a
    per-sample class call — too slow for a vitals loop); geovocab2 stays the
    reference implementation, parity-checked via cv_reference_check()."""
    B = pts.shape[0]
    d2 = torch.cdist(pts.double(), pts.double()).pow(2)          # (B,5,5)
    cm = torch.ones(B, 6, 6, dtype=torch.float64, device=pts.device)
    cm[:, 0, 0] = 0.0
    cm[:, 1:, 1:] = d2
    det = torch.linalg.det(cm)
    return (-det / 9216.0).clamp_min(0.0).sqrt().float()


@torch.no_grad()
def pentachoron_cv(rows: torch.Tensor, n_samples: int = 200,
                   generator: torch.Generator | None = None) -> float:
    """CV (std/mean) of Cayley-Menger 4-simplex volumes over n_samples random
    5-row subsets. Rows are row-normalized before measurement. Uses the built-in
    batched CM (float64 det); validate against geovocab2 with
    cv_reference_check() after any change to the volume math."""
    x = torch.nn.functional.normalize(rows.float(), dim=-1)
    n = x.shape[0]
    if n < 5:
        raise ValueError(f"pentachoron_cv needs >=5 rows, got {n}")
    g = generator or torch.Generator(device="cpu").manual_seed(0)
    idx = torch.stack([torch.randperm(n, generator=g)[:5]
                       for _ in range(n_samples)])               # (B,5)
    v = _pentachoron_volumes(x[idx].cpu())
    return (v.std() / v.mean().clamp_min(1e-12)).item()


@torch.no_grad()
def cv_reference_check(n_trials: int = 50, tol: float = 1e-5) -> float:
    """Parity check of the built-in batched CM against geovocab2's reference
    implementation (the formula's source of truth). Returns max |rel diff|;
    raises if geovocab2 is absent or parity fails. Run after touching
    _pentachoron_volumes."""
    try:
        from geovocab2.shapes.formula.symbolic.cayley_menger import (
            CayleyMengerFromSimplex)
    except Exception as e:  # pragma: no cover
        raise ImportError(
            "cv_reference_check requires geovocab2 (install via the geolip-svae "
            "umbrella: pip install git+https://github.com/AbstractEyes/"
            "geolip-svae).") from e
    ref = CayleyMengerFromSimplex()
    g = torch.Generator().manual_seed(0)
    pts = torch.nn.functional.normalize(
        torch.randn(n_trials, 5, 4, generator=g), dim=-1)
    mine = _pentachoron_volumes(pts)
    # compare at float64: the reference computes in the INPUT dtype, and fp32
    # dets lose up to ~4% on near-degenerate pentachora (measured 2026-07-11)
    tsuccessors = torch.stack([ref.forward(p.double())["volume"].float() for p in pts])
    rel = ((mine - tsuccessors).abs() / tsuccessors.abs().clamp_min(1e-12)).max().item()
    if rel > tol:
        raise AssertionError(f"CM parity vs geovocab2 failed: max rel {rel}")
    return rel


# ------------------------------------------------------------------------- aliveness
@torch.no_grad()
def axis_aliveness(oriented_weights: torch.Tensor, alive_thresh: float = 1e-3) -> dict:
    """`oriented_weights`: (..., 2K) nonnegative oriented-softmax address rows
    (sum to 1 on the last dim). Returns axes-alive count, mean-usage perplexity
    (hppl analogue; healthy hosted reference 125-126/128), and a collapse flag.
    Reference behavior: near-uniform aliveness at div_weight=0 (discovery #22)."""
    w = oriented_weights.reshape(-1, oriented_weights.shape[-1]).float()
    usage = w.mean(0)
    usage = usage / usage.sum().clamp_min(1e-12)
    # an axis is alive if its mean usage exceeds alive_thresh x the uniform share
    alive = int((usage > alive_thresh * (1.0 / usage.numel())).sum())
    ent = -(usage.clamp_min(1e-12) * usage.clamp_min(1e-12).log()).sum()
    ppl = float(ent.exp())
    return {"axes_total": usage.numel(), "axes_alive": alive, "usage_ppl": ppl,
            "collapsed": ppl < 0.05 * usage.numel()}


# ------------------------------------------------------------------------------ gates
@torch.no_grad()
def gate_stats(gates: torch.Tensor) -> dict:
    """Gate values (post-sigmoid/clamp). Reports mean and whether it sits in the
    0.012-0.03 band (read-only — the band is a candidate invariant, never a target)."""
    g = gates.float().flatten()
    m = g.mean().item()
    return {"mean": m, "std": g.std().item(),
            "in_band": GATE_BAND[0] <= m <= GATE_BAND[1]}


# ------------------------------------------------------------------------------ paths
@torch.no_grad()
def path_diversity(ids: torch.Tensor) -> dict:
    """Unique-path counting with the FIXED multiplicative hash:
    ((ids * 2654435761) % 2^32) >> 16 — Knuth needs the HIGH bits; the low-16
    variant produced the retracted ~1,500 path ceiling (sessions/2026-07-07).
    `ids`: integer tensor, one composed path id per row (any shape)."""
    x = ids.reshape(-1).to(torch.int64)
    hashed = ((x * KNUTH32) % (1 << 32)) >> 16
    return {"n": int(x.numel()),
            "unique_raw": int(torch.unique(x).numel()),
            "unique_hashed": int(torch.unique(hashed).numel())}


@torch.no_grad()
def compose_path_ids(stage_indices: list[torch.Tensor], radix: int) -> torch.Tensor:
    """Compose per-stage discrete indices (each (...,) int in [0, radix)) into a
    single path id, positional base-`radix` — construction, not hashing."""
    out = torch.zeros_like(stage_indices[0], dtype=torch.int64)
    for s in stage_indices:
        out = out * radix + s.to(torch.int64)
    return out


# --------------------------------------------------------------------- grad democracy
@torch.no_grad()
def grad_norm_spread(groups: dict[str, list[torch.nn.Parameter]]) -> dict:
    """Gradient-democracy monitor. `groups`: name -> params of one parallel member
    (tower/expert). Reports per-group grad norms and the orders-of-magnitude spread.
    Reference: unequalized heterogeneous towers spread ~20 orders (fibonacci dead at
    2.25e-21 under helix); equalized ~0.0 (canon/fibonacci_systems.md)."""
    norms = {}
    for name, params in groups.items():
        gs = [p.grad for p in params if p.grad is not None]
        norms[name] = float(torch.sqrt(sum((g.float() ** 2).sum() for g in gs)).item()) \
            if gs else 0.0
    vals = [v for v in norms.values() if v > 0]
    spread = (math.log10(max(vals)) - math.log10(min(vals))) if len(vals) >= 2 else 0.0
    return {"norms": norms, "spread_orders": spread, "dead": [k for k, v in norms.items() if v == 0.0]}


# ----------------------------------------------------------------------------- screen
class CVScreen:
    """CV@N early band screen (tri-band ft1): record pentachoron CV at `step_mark`
    batches; classify <0.30 LOW / 0.35-0.50 MID / >0.80 HIGH. Turns ~2h/config
    into ~7min. Readout only."""
    def __init__(self, step_mark: int = 1000):
        self.step_mark = step_mark
        self.recorded: float | None = None

    def maybe_record(self, step: int, rows: torch.Tensor) -> float | None:
        if self.recorded is None and step >= self.step_mark:
            self.recorded = pentachoron_cv(rows)
        return self.recorded

    @property
    def band(self) -> str | None:
        c = self.recorded
        if c is None:
            return None
        if c < 0.30:
            return "LOW"
        if 0.35 <= c <= 0.50:
            return "MID"
        if c > 0.80:
            return "HIGH"
        return "BETWEEN"


# ------------------------------------------------------------------------------ smoke
if __name__ == "__main__":  # shapes/parse smoke ONLY — no training, ever.
    g = torch.Generator().manual_seed(0)
    K, D = 64, 4
    init = torch.nn.functional.normalize(torch.randn(K, D, generator=g), dim=-1)
    cur = torch.nn.functional.normalize(init + 0.29 * torch.randn(K, D, generator=g), dim=-1)
    print("drift:", {k: v for k, v in anchor_drift(cur, init).items() if k != "per_row"})
    w = torch.softmax(torch.randn(32, 2 * K, generator=g), dim=-1)
    print("aliveness:", axis_aliveness(w))
    print("gates:", gate_stats(torch.full((8,), 0.024)))
    ids = compose_path_ids([torch.randint(0, 16, (4096,), generator=g) for _ in range(4)], 16)
    print("paths:", path_diversity(ids))
    lin = torch.nn.Linear(8, 8)
    lin(torch.randn(4, 8)).sum().backward()
    print("democracy:", grad_norm_spread({"a": list(lin.parameters())}))
    print("OK — vitals smoke passed (pentachoron_cv needs geovocab2; run on GPU env)")