#!/usr/bin/env python3
"""Causal attention ablation at the hop-m query.

Zero out attention from q_m onto a chosen key set, renormalize the row, and
measure the drop in frontier accuracy / ce_score. Two primary ablations:

  latents   — previous latent positions (and optionally self)
  edges_m   — '|' separators of the two depth-m reachable edges
  random    — matched-size random control among non-special tokens

If attending to latents is load-bearing, ablating them should collapse metrics.
If attending to hop-m edge separators is load-bearing, that ablation should.
"""
import argparse
import json
import math
import os
import sys
from collections import defaultdict

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from scripts.attention_atlas import build_prompt, load_model, graph_roles


def ce_score_from_logits(logits, frontier):
    """s = min(1, exp(log|F| - CE)) with CE vs Unif(F)."""
    Fset = list(frontier)
    if not Fset:
        return 0.0
    logp = F.log_softmax(logits, dim=-1)
    ce = -sum(logp[n].item() for n in Fset) / len(Fset)
    return min(1.0, math.exp(math.log(len(Fset)) - ce))


def patch_gpt2_attn(model, mask_fn):
    """Replace each block's _attn so that after softmax, mask_fn can zero keys.

    mask_fn(layer_idx, attn_weights) -> modified attn_weights
    attn_weights shape: (B, H, Tq, Tk)
    """
    originals = []
    for li, block in enumerate(model.base_causallm.transformer.h):
        attn = block.attn
        # Keep the unbound original function; instance may hold a bound method.
        orig = attn._attn.__func__ if hasattr(attn._attn, "__func__") else attn._attn
        originals.append((attn, attn._attn))

        def make(layer_idx, attn_mod, orig_fn):
            def _attn(query, key, value, attention_mask=None, head_mask=None):
                attn_weights = torch.matmul(query, key.transpose(-1, -2))
                if attn_mod.scale_attn_weights:
                    attn_weights = attn_weights / torch.full(
                        [], value.size(-1) ** 0.5, dtype=attn_weights.dtype,
                        device=attn_weights.device,
                    )
                if attn_mod.scale_attn_by_inverse_layer_idx:
                    attn_weights = attn_weights / float(attn_mod.layer_idx + 1)
                if not attn_mod.is_cross_attention:
                    query_length, key_length = query.size(-2), key.size(-2)
                    causal_mask = attn_mod.bias[
                        :, :, key_length - query_length: key_length, :key_length
                    ]
                    mask_value = torch.full(
                        [], torch.finfo(attn_weights.dtype).min,
                        dtype=attn_weights.dtype, device=attn_weights.device,
                    )
                    attn_weights = torch.where(
                        causal_mask, attn_weights.to(attn_weights.dtype), mask_value
                    )
                if attention_mask is not None:
                    attn_weights = attn_weights + attention_mask
                attn_weights = F.softmax(attn_weights, dim=-1)
                attn_weights = attn_weights.type(value.dtype)
                # Ablate AFTER softmax, BEFORE dropout / value mul, then renorm.
                attn_weights = mask_fn(layer_idx, attn_weights)
                attn_weights = attn_mod.attn_dropout(attn_weights)
                if head_mask is not None:
                    attn_weights = attn_weights * head_mask
                attn_output = torch.matmul(attn_weights, value)
                return attn_output, attn_weights
            return _attn

        attn._attn = make(li, attn, orig)
    return originals


def restore_attn(model, originals):
    for attn, orig in originals:
        attn._attn = orig


@torch.no_grad()
def evaluate(model, batch_meta, input_ids, device, mask_fn=None):
    """One filled Coconut forward + replay; optional attention mask_fn during replay."""
    B, T = input_ids.shape
    am = torch.ones_like(input_ids)
    pos = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
    out = model.forward(input_ids, am, input_ids.clone(), pos)

    originals = None
    if mask_fn is not None:
        originals = patch_gpt2_attn(model, mask_fn)
    try:
        rep = model.base_causallm(
            inputs_embeds=out.inputs_embeds,
            attention_mask=am,
            position_ids=pos,
        )
        logits = rep.logits  # (B, T, V)
    finally:
        if originals is not None:
            restore_attn(model, originals)

    stats = defaultdict(lambda: {"front": 0, "ce": 0.0, "n": 0})
    for bi, (s, roles, root_pos, L) in enumerate(batch_meta):
        for m in range(1, L + 1):
            q = root_pos + (m - 1)
            lg = logits[bi, q]
            frontier = [int(x) for x in s["neighbor_k"].get(str(m), [])]
            pred = int(lg[:100].argmax().item())  # node vocab
            stats[m]["n"] += 1
            stats[m]["front"] += int(pred in frontier)
            stats[m]["ce"] += ce_score_from_logits(lg, frontier)
    return stats


def merge(dst, src):
    for m, v in src.items():
        dst[m]["n"] += v["n"]
        dst[m]["front"] += v["front"]
        dst[m]["ce"] += v["ce"]


def summarize(stats):
    out = {}
    for m in sorted(stats):
        n = max(stats[m]["n"], 1)
        out[m] = {
            "frontier": stats[m]["front"] / n,
            "ce_score": stats[m]["ce"] / n,
            "n": n,
        }
    return out


def build_key_sets(roles, root_pos, m, node_role):
    q = root_pos + (m - 1)
    latents = [p for p in range(q) if roles[p][0] == "latent"]
    # hop-m reachable edge separators (tagged with dest node)
    edges_m = [
        p for p in range(q + 1)
        if roles[p][0] == "sep"
        and node_role.get(roles[p][1], ("off", -1))[0] == "pos"
        and node_role.get(roles[p][1])[1] == m
    ]
    # all seps
    all_sep = [p for p in range(q + 1) if roles[p][0] == "sep"]
    # random matched to |edges_m| among non-latent non-self content
    pool = [
        p for p in range(q)
        if roles[p][0] not in ("latent", "bos") and p not in edges_m
    ]
    return q, latents, edges_m, all_sep, pool


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpts/star-coconut-L10-bfs-backtrack-ce095/checkpoint_2000")
    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("--n", type=int, default=128)
    ap.add_argument("--batch_size", type=int, default=16)
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--layers", default="both", choices=["0", "1", "both"])
    args = ap.parse_args()

    rng = np.random.default_rng(args.seed)
    model, tok = load_model(args.ckpt, args.model_id, args.device)
    data = json.load(open(args.val))[: args.n]
    L = args.L
    order = np.arange(len(data[0]["edges"]))

    # Prebuild per-sample role info (fixed edge order)
    metas_all = []
    for s in data:
        ids, roles, rp = build_prompt(s, tok, L, order, False)
        nr = graph_roles(s, L)
        metas_all.append((s, roles, rp, nr, ids))

    ablations = ["none", "latents", "edges_m", "random_match", "all_sep"]
    results = {name: defaultdict(lambda: {"front": 0, "ce": 0.0, "n": 0})
               for name in ablations}

    which = ({0, 1} if args.layers == "both"
             else {int(args.layers)})

    for start in range(0, len(metas_all), args.batch_size):
        chunk = metas_all[start:start + args.batch_size]
        seqs = [c[4] for c in chunk]
        input_ids = torch.tensor(seqs, device=args.device)
        batch_meta = [(c[0], c[1], c[2], L) for c in chunk]

        # per-batch key sets for each (bi, m)
        keysets = []
        for bi, (s, roles, rp, nr, _) in enumerate(chunk):
            per_m = {}
            for m in range(1, L + 1):
                q, lat, em, asep, pool = build_key_sets(roles, rp, m, nr)
                n_em = max(len(em), 1)
                rnd = list(rng.choice(pool, size=min(n_em, len(pool)), replace=False)) if pool else []
                per_m[m] = {"q": q, "latents": lat, "edges_m": em,
                            "random": rnd, "all_sep": asep}
            keysets.append(per_m)

        def make_mask(mode):
            if mode == "none":
                return None

            def mask_fn(layer_idx, attn_weights):
                if layer_idx not in which:
                    return attn_weights
                # attn_weights: (B, H, T, T)
                w = attn_weights.clone()
                for bi, per_m in enumerate(keysets):
                    for m, ks in per_m.items():
                        q = ks["q"]
                        if mode == "latents":
                            keys = ks["latents"]
                        elif mode == "edges_m":
                            keys = ks["edges_m"]
                        elif mode == "random_match":
                            keys = ks["random"]
                        elif mode == "all_sep":
                            keys = ks["all_sep"]
                        else:
                            keys = []
                        if not keys:
                            continue
                        w[bi, :, q, keys] = 0.0
                        row = w[bi, :, q, : q + 1]
                        denom = row.sum(dim=-1, keepdim=True).clamp_min(1e-12)
                        w[bi, :, q, : q + 1] = row / denom
                return w
            return mask_fn

        for mode in ablations:
            st = evaluate(model, batch_meta, input_ids, args.device, make_mask(mode))
            merge(results[mode], st)
        print(f"  {min(start+args.batch_size, len(metas_all))}/{len(metas_all)}", flush=True)

    print(f"\n=== ablation on {args.ckpt} | n={len(data)} | layers={args.layers} ===")
    print(f"{'hop':>3} | {'clean F / CE':>18} | {'-latents ΔF / ΔCE':>20} | "
          f"{'-edges_m ΔF / ΔCE':>20} | {'-random ΔF / ΔCE':>18} | {'-all_sep ΔF / ΔCE':>18}")
    clean = summarize(results["none"])
    for m in range(1, L + 1):
        c = clean[m]
        def d(mode):
            s = summarize(results[mode])[m]
            return s["frontier"] - c["frontier"], s["ce_score"] - c["ce_score"]
        dl = d("latents"); de = d("edges_m"); dr = d("random_match"); da = d("all_sep")
        print(f"{m:3d} | {c['frontier']:.3f} / {c['ce_score']:.3f}     | "
              f"{dl[0]:+.3f} / {dl[1]:+.3f}         | "
              f"{de[0]:+.3f} / {de[1]:+.3f}         | "
              f"{dr[0]:+.3f} / {dr[1]:+.3f}       | "
              f"{da[0]:+.3f} / {da[1]:+.3f}")

    # means over hops 2..10 (hop 1 has no previous latents)
    def mean_delta(mode, metric, hops):
        c = summarize(results["none"])
        s = summarize(results[mode])
        return float(np.mean([s[m][metric] - c[m][metric] for m in hops]))

    hops = list(range(2, L + 1))
    print("\n--- mean Δ over hops 2..10 ---")
    for mode in ["latents", "edges_m", "random_match", "all_sep"]:
        print(f"  {mode:14s}  Δfrontier={mean_delta(mode,'frontier',hops):+.4f}  "
              f"Δce_score={mean_delta(mode,'ce_score',hops):+.4f}")

    out = {
        "ckpt": args.ckpt, "n": len(data), "layers": args.layers,
        "per_hop": {mode: summarize(results[mode]) for mode in ablations},
    }
    os.makedirs("figs/attention_atlas", exist_ok=True)
    path = f"figs/attention_atlas/ablation_layers_{args.layers}.json"
    # json-safe
    out["per_hop"] = {
        mode: {str(k): v for k, v in summarize(results[mode]).items()}
        for mode in ablations
    }
    with open(path, "w") as f:
        json.dump(out, f, indent=2)
    print(f"\nwrote {path}")


if __name__ == "__main__":
    main()