"""Shared measurement helpers for docs/DECODE_FIX_AND_BPTT.md (reuse the root_cause tools)."""
import random, re, sys
from pathlib import Path
import numpy as np
import torch, torch.nn.functional as F
HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2]
RC = ROOT / "docs/artifacts/root_cause"
for p in (str(ROOT), str(ROOT / "tools"), str(RC), str(ROOT / "docs/artifacts/field_leak_spectrum")):
    if p not in sys.path:
        sys.path.insert(0, p)
from chat_atom_native import load_model, generate_reply  # noqa: E402
from common import PROMPTS, val_windows_bytes  # noqa: E402  (root_cause)
import eval_windows as EW  # noqa: E402  (the 2048-transition val)

SP = set(b" ,.-'?!:;")


def val2048(m):
    """normal + field-alone byte CE on the exact eval_windows.py val (4 x 512)."""
    normal, alone = [], []
    for w in EW.windows():
        normal += EW.run(m, w["pairs"], False); alone += EW.run(m, w["pairs"], True)
    return {"normal_byte_ce": float(np.mean(normal)), "field_alone_byte_ce": float(np.mean(alone))}


@torch.no_grad()
def run_logp(m, pk):
    m.eval(); m.reset_state(reset_atomizer=False)
    return [torch.log_softmax(m.forward_packet(p)["surface"]["byte_logits"][0], -1) for p in pk[:-1]]


def perturbation_kl(m, lags=(1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64), per_window=6, seed=0):
    """Same protocol as root_cause/diag_atom_more.py: flip one ASCII letter at p in a val window,
    KL(next-byte | orig || perturbed) L bytes later; normal and field-alone (last_atom off)."""
    wins = val_windows_bytes(); rng = random.Random(seed); sens = {}
    for mode in ("normal", "field_alone"):
        m.surface.last_atom_force_off = mode == "field_alone"
        kl = {L: [] for L in lags}
        for w in wins:
            cands = [p for p in range(64, 440) if 97 <= w[p] <= 122]
            for p in rng.sample(cands, per_window):
                wp = bytearray(w); wp[p] = rng.choice([c for c in range(97, 123) if c != w[p]])
                seg, segp = bytes(w[: p + 66]), bytes(wp[: p + 66])
                a = run_logp(m, m.atomizer.encode_bytes(seg, reset=True))
                b = run_logp(m, m.atomizer.encode_bytes(segp, reset=True))
                for L in lags:
                    i = p + L - 1
                    kl[L].append(float((a[i].exp() * (a[i] - b[i])).sum()))
        sens[mode] = {int(L): float(np.mean(v)) for L, v in kl.items()}
    m.surface.last_atom_force_off = False
    return sens


def grad_by_lag(m, K=None, lags=(0, 1, 2, 4, 8, 16, 32, 64), T=200):
    """||d loss_T / d atom_{T-L}|| (r,phi,omega,E,kappa) in train mode on val window 1,
    root_cause protocol; K = field_bptt to set (None = keep the model's)."""
    old_k = m.field_bptt
    if K is not None:
        m.set_field_bptt(K)
    m.train(); drop = m.surface.last_atom_dropout; m.surface.last_atom_dropout = 0.0
    stash = []; orig = m.compiler.forward

    def hooked(features, atom_count=0):
        atom, op, conf = orig(features, atom_count)
        for n in ("r", "phi", "omega", "E", "kappa"):
            getattr(atom, n).retain_grad()
        stash.append(atom); return atom, op, conf
    m.compiler.forward = hooked
    try:
        wins = val_windows_bytes()
        m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(wins[1][: T + 1], reset=True)
        for q in pk[: T - 1]:
            m.forward_packet(q)
        out = m.forward_packet(pk[T - 1])
        loss = F.cross_entropy(out["surface"]["byte_logits"], torch.tensor([pk[T].payload[0]]))
        loss.backward()
    finally:
        m.compiler.forward = orig
        m.surface.last_atom_dropout = drop
        m.zero_grad(set_to_none=True)
        m.set_field_bptt(old_k)
        m.eval()
    return {int(L): sum(float(getattr(stash[T - 1 - L], n).grad.norm()) if getattr(stash[T - 1 - L], n).grad is not None else 0.0
                        for n in ("r", "phi", "omega", "E", "kappa")) for L in lags}


def chat_samples(m, seeds=range(8), gate=True):
    """diag_decode.py protocol: generate_packets(prompt) seeds 0..7, T .8, top-k 12, 96 pkts/100 bytes."""
    out = {}
    for p in PROMPTS:
        out[p] = []
        for s in seeds:
            torch.manual_seed(s)
            raw = m.generate_packets(p, max_packets=96, temperature=0.8, top_k=12, deterministic=False,
                                     max_length=100, prefer_printable=True, reset=True, merge_enabled=False,
                                     speech_gate=gate, anti_repeat=gate)
            out[p].append(raw.decode("utf-8", "replace"))
    return out


def doubled_rate(samples):
    txt = "".join(s for v in samples.values() for s in v)
    return {"chars": len(txt), "punct": sum(1 for c in txt if c in ",.-'?!:;"),
            "doubled_punct_pairs": len(re.findall(r"([,.\-'?!:;])\1", txt)),
            "space_runs_ge3": len(re.findall(r"   +", txt))}


def probe_chats(m, seeds=(1, 2, 3)):
    """tools/chat_atom_native.py --no-role-prime --max-packets 96 --max-length 100 -T .8 -k 12;
    unseeded = one draw per prompt (as the probe runs), plus chat_seeds_probe seeds."""
    rows = []
    for p in PROMPTS:
        rows.append({"seed": None, "prompt": p, "resp": generate_reply(m, p, max_packets=96, temperature=0.8, top_k=12,
                                                                      max_length=100, role_prime=False)})
    for s in seeds:
        for p in PROMPTS:
            torch.manual_seed(s)
            rows.append({"seed": s, "prompt": p, "resp": generate_reply(m, p, max_packets=96, temperature=0.8, top_k=12,
                                                                       max_length=100, role_prime=False)})
    return rows


_JUDGE = {}


def judge_nll(samples):
    """GRU f100 judge (val .10) NLL/byte on the first reply line (root_cause/judge.py protocol)."""
    from judge import gru_nll
    from gru_baseline import GRULM
    if "m" not in _JUDGE:
        j = GRULM(512, 2); j.load_state_dict(torch.load(RC / "gru_f100.pt")); j.eval(); _JUDGE["m"] = j
    vals = []
    for p, vs in samples.items():
        for v in vs:
            first = v.split("\n")[0] + ("\n" if "\n" in v else "")
            x = gru_nll(_JUDGE["m"], p, first)
            if x is not None:
                vals.append(x)
    return float(np.mean(vals)) if vals else None


def prompt_identical_frac(samples):
    vals = list(samples.values())
    return sum(len({vals[k][i] for k in range(len(vals))}) == 1 for i in range(len(vals[0]))) / len(vals[0])