#!/usr/bin/env python3
"""Does dyn12 leak the future into the past?

Omega is defined in section 3.4 as attention RECEIVED -- summed over the QUERY axis:

    omega = a.mean(1).sum(-2)        # [B, T], indexed by KEY

Under a causal mask a_ij is nonzero only for j <= i, so

    Omega_j = sum over i >= j of a_ij

which includes queries from tokens AFTER j. That Omega updates state_j, state_j feeds
the next layer's Hebbian kernel, the kernel shapes attention, and attention shapes the
logits at position j. If that chain is live, position j can see its own future and the
held-out loss is measuring a leak rather than a mechanism.

THE TEST, which does not care about any of that reasoning:
  Run the model on a sequence. Change ONLY the LAST token. If the logits at an EARLY
  position move, information flowed backwards. A correct causal LM cannot do this --
  its prediction at position j is a function of tokens 0..j alone.

A standard transformer is included as the control, so a positive result cannot be
blamed on the harness.

  python tools/causality_probe.py
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

import torch

_root = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_root / 'architecture'))
sys.path.insert(0, str(_root))
sys.stdout.reconfigure(encoding="utf-8", errors="replace")

import cosmos_state_ladder as L  # noqa: E402

DEV = "cpu"
torch.manual_seed(0)


def logits_for(model, ids):
    model.eval()
    with torch.no_grad():
        out = model(ids)
        return out[0] if isinstance(out, tuple) else out


def probe(rung: str, T: int = 24, trials: int = 3, gate: float | None = None,
          seed: int | None = None) -> dict:
    """gate=None leaves the gate at init. PHOS's trained layer-0 gate is 0.562, and the
    leak is proportional to it -- Omega reaches the logits only through g*H, so measuring
    at initialisation (g near zero) understates a trained model by orders of magnitude.

    seed=None draws a FRESH RANDOM seed, and that is the default on purpose. A causality
    result that only holds for one hand-picked seed is worth nothing -- the reader cannot
    tell a real property from a lucky draw. Random weights and random sequences every run
    mean each person who executes this gets independent evidence, and a leak that only
    appears sometimes still gets caught. Pass --seed N when you need to reproduce a
    specific run exactly."""
    if seed is None:
        seed = int.from_bytes(os.urandom(4), "little")
    torch.manual_seed(seed)
    V = 96
    model = L.Ladder(vocab=V, rung=rung, ffn="harmonic").to(DEV)
    if gate is not None:
        import math as _m
        raw = _m.log(gate / (1.0 - gate))          # invert sigmoid
        with torch.no_grad():
            for b in model.blocks:
                if hasattr(b.attn, "gate"):
                    b.attn.gate.fill_(raw)
    worst = 0.0
    for t in range(trials):
        g = torch.Generator().manual_seed(seed + 1 + t)
        ids = torch.randint(0, V, (1, T), generator=g)
        base = logits_for(model, ids)

        alt = ids.clone()
        # change ONLY the final token
        alt[0, -1] = (alt[0, -1] + 1 + t) % V
        moved = logits_for(model, alt)

        # compare every position EXCEPT the last: none may change
        early = (base[:, :-1, :] - moved[:, :-1, :]).abs().max().item()
        worst = max(worst, early)
    return {"rung": rung, "max_early_logit_change": worst}


def main() -> int:
    seed = None
    if "--seed" in sys.argv:
        seed = int(sys.argv[sys.argv.index("--seed") + 1])
    run_seed = seed if seed is not None else int.from_bytes(os.urandom(4), "little")

    print("  Changing ONLY the last token. Logits at earlier positions must not move.")
    print("  The control rungs give EXACTLY 0.0, so any nonzero value is a real")
    print("  dependency, not float noise -- identical ops on identical inputs.")
    print(f"  Omega mode: {'CAUSAL (per-query entropy)' if L.CAUSAL_OMEGA else 'ORIGINAL query-sum'}")
    print(f"  seed: {run_seed}"
          + ("  [fixed]" if seed is not None else "  [random -- your run is independent evidence]"))
    # Derive the path from argv rather than hardcoding it: this file ships inside the
    # kit as benchmarks/causality_probe.py, and telling a reader to run a path that does
    # not exist on their disk is how a reproducible result stops being reproducible.
    print(f"  reproduce this exact run:  python {os.path.relpath(sys.argv[0])} --seed {run_seed}\n")

    cases = [("none", None), ("static54", None), ("dyn12", None),
             ("dyn12", 0.562), ("tri", 0.562)]
    rows = []
    for rung, gate in cases:
        try:
            r = probe(rung, gate=gate, seed=run_seed)
            r["gate"] = gate
            rows.append(r)
        except Exception as e:
            rows.append({"rung": rung, "gate": gate, "error": f"{type(e).__name__}: {e}"})

    for r in rows:
        label = r["rung"] + (f" @g={r['gate']}" if r.get("gate") else " @init")
        if "error" in r:
            print(f"  {label:<20} ERROR {r['error']}")
            continue
        d = r["max_early_logit_change"]
        verdict = "causal" if d == 0.0 else "LEAKS THE FUTURE"
        print(f"  {label:<20} max early-position change: {d:.3e}   {verdict}")

    ok = [r for r in rows if "error" not in r]
    base = next((r for r in ok if r["rung"] == "none"), None)
    if base and base["max_early_logit_change"] != 0.0:
        print("\n  The BASELINE moved too -- the harness is wrong, not the architecture.")
        return 2
    leaky = [r for r in ok if r["max_early_logit_change"] > 0.0]
    print()
    if leaky:
        print("  Omega is summed over QUERIES, so Omega_j counts attention from tokens")
        print("  AFTER j. Position j therefore sees its own future. Held-out loss won")
        print("  this way is not comparable to a rung that cannot do it.")
        return 1
    print("  No rung leaks. Omega's query-sum stays inside the causal boundary.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())