#!/usr/bin/env python3
"""Characterize the +-1 off-diagonal entries in the latent confusion matrices.

Question: when latent slot m decodes a node at TRUE depth m-1 or m+1, is that
  (a) a PHASE DRIFT -- the BFS chain slipped a step and stays shifted, or
  (b) an ISOLATED SLIP -- slot m errs but slot m+1 is back on the diagonal, or
  (c) a STALL -- slot m decodes the *same node* slot m-1 already decoded
      (the chain failed to advance a hop)?

For each arm we report:
  * histogram of offset = (true depth of decoded node) - m
  * repeat rate: decoded(m) == decoded(m-1)  [stall signature]
  * drift persistence: P(slot m+1 also off by the same offset | slot m off by d)
    vs the base rate. High persistence => real drift; low => isolated noise.
"""
import argparse
import json
from collections import Counter

import torch
from transformers import AutoModelForCausalLM, AutoConfig

from stokenizer import STokenizer
from coconut import Coconut
from scripts.probe_latents import build_prefix_tokens, node_depth_maps

ARMS = [
    ("Backtracking", "backtrack"),
    ("Current-stage-only", "curstage"),
    ("Retention-gated (no repair)", "accstage-nobt"),
]


def latest_ckpt(slug):
    import os
    d = f"ckpts/star-coconut-L10-bfs-{slug}"
    cks = sorted((f for f in os.listdir(d) if f.startswith("checkpoint_")),
                 key=lambda x: int(x.split("_")[1]))
    return os.path.join(d, cks[-1])


@torch.no_grad()
def decode_chains(ckpt, val_path, model_id, L, device, batch_size=64):
    """Return list of per-sample dicts: decoded node id + its true depth per slot."""
    tok = STokenizer()
    latent_id = tok.convert_tokens_to_ids("<|latent|>")
    base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(model_id))
    model = Coconut(base, latent_id,
                    tok.convert_tokens_to_ids("<|start-latent|>"),
                    tok.convert_tokens_to_ids("<|end-latent|>"),
                    tok.eos_token_id)
    model.load_state_dict(torch.load(ckpt, map_location="cpu"), strict=False)
    model.to(device).eval()

    data = json.load(open(val_path))
    chains = []
    for i in range(0, len(data), batch_size):
        batch = data[i:i + batch_size]
        input_ids = torch.tensor([build_prefix_tokens(s, tok) + [latent_id] * L
                                  for s in batch], device=device)
        attn = torch.ones_like(input_ids)
        pos = torch.arange(input_ids.shape[1], device=device).unsqueeze(0).expand(len(batch), -1)
        logits = model.forward(input_ids, attn, input_ids.clone(), pos).logits
        for bi, s in enumerate(batch):
            role = node_depth_maps(s, L)
            root_pos = len(build_prefix_tokens(s, tok)) - 1
            nodes, depths = [], []
            for m in range(1, L + 1):
                pred = int(torch.argmax(logits[bi, root_pos + (m - 1)]).item())
                r = role.get(pred)
                nodes.append(pred)
                # depth is None for decoy/off-graph (we only study in-component drift)
                depths.append(r[1] if (r is not None and r[0] == "pos") else None)
            chains.append({"nodes": nodes, "depths": depths})
    return chains


def analyze(chains, L):
    offsets = Counter()
    repeat = 0
    slot_total = 0
    # drift persistence
    off_then_off_same = 0
    off_then_diag = 0
    off_events = 0
    for c in chains:
        for m in range(1, L + 1):
            d = c["depths"][m - 1]
            slot_total += 1
            if d is None:
                offsets["decoy/off"] += 1
                continue
            offsets[d - m] += 1
            if m >= 2 and c["nodes"][m - 1] == c["nodes"][m - 2]:
                repeat += 1
        # persistence: slot m off by o, what does slot m+1 do?
        for m in range(1, L):
            d, d2 = c["depths"][m - 1], c["depths"][m]
            if d is None or d2 is None:
                continue
            o = d - m
            if o == 0:
                continue
            off_events += 1
            if d2 - (m + 1) == o:
                off_then_off_same += 1
            elif d2 - (m + 1) == 0:
                off_then_diag += 1
    return offsets, repeat, slot_total, off_events, off_then_off_same, off_then_diag


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--val", default="data/star_2arm_L10_valid_fo_bfs.json")
    ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json")
    ap.add_argument("--L", type=int, default=10)
    ap.add_argument("--device", default="cuda:0")
    args = ap.parse_args()

    for label, slug in ARMS:
        ck = latest_ckpt(slug)
        chains = decode_chains(ck, args.val, args.model_id, args.L, args.device)
        offs, repeat, tot, oe, same, diag = analyze(chains, args.L)
        print(f"\n===== {label}  ({ck}) =====")
        print(f"total latent slots decoded: {tot}")
        print("offset (true_depth - slot) histogram:")
        for k in sorted([k for k in offs if isinstance(k, int)]):
            print(f"   {k:+d}: {offs[k]:>6}  ({offs[k]/tot:.4%})")
        print(f"   decoy/off-graph: {offs['decoy/off']:>6}  ({offs['decoy/off']/tot:.4%})")
        print(f"stall (decoded node == previous slot's node): {repeat} "
              f"({repeat/tot:.4%})")
        if oe:
            print(f"drift persistence: of {oe} off-diagonal in-component events, "
                  f"{same} ({same/oe:.1%}) kept the SAME offset at slot m+1, "
                  f"{diag} ({diag/oe:.1%}) snapped back to the diagonal")
        else:
            print("drift persistence: no off-diagonal in-component events")


if __name__ == "__main__":
    main()