#!/usr/bin/env python3
"""Does a latent encode ONE node at depth m, or BOTH frontier nodes at once?

At stage k the training target is random.choice(neighbor_k[k]); in a 2-arm star
that set has 2 members (one per arm), so the supervision is ambiguous. Two
hypotheses for what the model learns:

  (H1) COMMIT: the latent picks one arm and puts nearly all mass on it.
       -> p(top1) >> p(top2), and top2 is often not the sibling frontier node.
  (H2) SUPERPOSITION: the latent represents BOTH depth-m nodes simultaneously;
       argmax then breaks the tie ~arbitrarily.
       -> top-2 tokens are exactly the two frontier nodes, with comparable mass.

We also ask whether the model has any PREFERENCE for the arm that leads to the
true target leaf (i.e. does it secretly know the answer early?):
  -> compare p(target-arm node) vs p(other-arm node) at each hop.

Reported per hop m:
  frontier_mass   mean total prob on the 2 true depth-m nodes
  top2_is_frontier  fraction where the top-2 tokens are exactly those 2 nodes
  p1/p2 ratio     mean ratio of larger to smaller of the two frontier probs
  target_arm_win  fraction where the target-arm node outranks the other arm
  p_target_share  mean p(target arm) / (p(target)+p(other))  [0.5 = no preference]
"""
import argparse
import json

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoConfig

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

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 run(ckpt, val_path, model_id, L, device, batch_size=64):
    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))
    acc = {m: {"mass": [], "top2": 0, "ratio": [], "twin": 0, "share": [], "n": 0}
           for m in range(1, L + 1)}
    n_frontier_sizes = {}

    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
        probs = F.softmax(logits.float(), dim=-1)

        for bi, s in enumerate(batch):
            root_pos = len(build_prefix_tokens(s, tok)) - 1
            steps = s["steps"]
            for m in range(1, L + 1):
                front = s["neighbor_k"].get(str(m), [])
                n_frontier_sizes[len(front)] = n_frontier_sizes.get(len(front), 0) + 1
                if len(front) != 2:
                    continue
                p = probs[bi, root_pos + (m - 1)]
                pf = [p[int(n)].item() for n in front]
                a = acc[m]
                a["n"] += 1
                a["mass"].append(sum(pf))
                top2 = set(torch.topk(p, 2).indices.tolist())
                if top2 == {int(front[0]), int(front[1])}:
                    a["top2"] += 1
                hi, lo = max(pf), min(pf)
                a["ratio"].append(hi / lo if lo > 0 else float("inf"))
                # which of the two is on the true shortest path to the target leaf?
                if m - 1 < len(steps):
                    tgt = int(steps[m - 1])
                    other = [int(n) for n in front if int(n) != tgt]
                    if other:
                        pt, po = p[tgt].item(), p[other[0]].item()
                        if pt > po:
                            a["twin"] += 1
                        if pt + po > 0:
                            a["share"].append(pt / (pt + po))
    return acc, n_frontier_sizes


def mean(x):
    return sum(x) / len(x) if x else float("nan")


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")
    ap.add_argument("--only", default=None, help="restrict to one arm slug")
    args = ap.parse_args()

    for label, slug in ARMS:
        if args.only and slug != args.only:
            continue
        ck = latest_ckpt(slug)
        acc, sizes = run(ck, args.val, args.model_id, args.L, args.device)
        print(f"\n===== {label}  ({ck}) =====")
        print(f"frontier-set sizes seen: {sizes}")
        print(f"{'hop':>4} {'frontier_mass':>14} {'top2_is_frontier':>17} "
              f"{'p_hi/p_lo':>10} {'target_arm_win':>15} {'p_target_share':>15}")
        for m in range(1, args.L + 1):
            a = acc[m]
            if not a["n"]:
                continue
            print(f"{m:>4} {mean(a['mass']):>14.3f} {a['top2']/a['n']:>17.3f} "
                  f"{mean(a['ratio']):>10.2f} {a['twin']/a['n']:>15.3f} "
                  f"{mean(a['share']):>15.3f}")


if __name__ == "__main__":
    main()