from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).parent.parent))
from model_cpu_gpt2 import CPUGPT, CPUGPTBlock, CPUGPTConfig, gpt2_small_config

SAMPLE_TEXTS = [
    "The quick brown fox jumps over the lazy dog",
    "In the beginning was the Word, and the Word was",
    "The capital of France is Paris and the capital of Germany",
]


def load_model(ckpt: str, device: str) -> CPUGPT:
    cfg = gpt2_small_config(seq_len=1024)
    model = CPUGPT(cfg).to(device).eval()
    if not ckpt:
        print("[explain] No checkpoint — random init", flush=True)
        return model
    if ckpt.startswith("s3://"):
        import boto3

        bucket, key = ckpt[5:].split("/", 1)
        local = "/tmp/explain_ckpt.pt"
        boto3.client("s3").download_file(bucket, key, local)
        ckpt = local
    elif ckpt.startswith("hf://"):
        from huggingface_hub import hf_hub_download

        path = ckpt[5:]
        repo, filename = path.rsplit("/", 1)
        ckpt = hf_hub_download(
            repo_id=repo, filename=filename, token=os.environ.get("HF_TOKEN")
        )
    raw = torch.load(ckpt, map_location="cpu", weights_only=False)
    state = raw.get("model", raw)
    state = {
        k.replace("module.", "").replace("_orig_mod.", ""): v for k, v in state.items()
    }
    if "wte.weight" in state:
        ckpt_embd = state["wte.weight"].shape[1]
        if ckpt_embd != cfg.n_embd:
            cfg = gpt2_small_config(seq_len=1024, n_embd=ckpt_embd)
            cfg.n_head = max(1, ckpt_embd // 64)
            for k, v in state.items():
                if "filter_td" in k:
                    cfg.fno_modes = v.shape[1]
                    break
            model = CPUGPT(cfg).to(device).eval()
    ms = model.state_dict()
    model.load_state_dict(
        {k: v for k, v in state.items() if k in ms and v.shape == ms[k].shape},
        strict=False,
    )
    print(f"[explain] loaded {model.param_count() / 1e6:.1f}M params", flush=True)
    return model


def get_tokenizer():
    try:
        from transformers import GPT2Tokenizer

        return GPT2Tokenizer.from_pretrained("gpt2")
    except Exception:
        import tiktoken

        return tiktoken.get_encoding("r50k_base")


def encode(tok, text: str) -> list[int]:
    if hasattr(tok, "encode") and hasattr(tok, "decode"):
        try:
            return tok.encode(text)
        except TypeError:
            return tok.encode(text, add_special_tokens=False)
    return tok.encode(text, add_special_tokens=False)


def decode_token(tok, tid: int) -> str:
    try:
        return repr(tok.decode([tid]))
    except Exception:
        return f"<{tid}>"


def logit_lens(
    model: CPUGPT, tok, texts: list[str], device: str, top_k: int = 5
) -> dict:
    print("[explain] Running logit lens ...", flush=True)
    results = {}
    unembed = model.lm_head.weight

    for text in texts:
        ids = encode(tok, text)[: model.cfg.seq_len - 1]
        idx = torch.tensor([ids], dtype=torch.long, device=device)
        layer_preds = []

        residuals = []
        hooks = []

        def make_hook(i):
            def h(module, inp, out):
                residuals.append((i, out.detach().clone()))

            return h

        for i, block in enumerate(model.blocks):
            hooks.append(block.register_forward_hook(make_hook(i)))

        with torch.no_grad():
            _ = model(idx)
        for h in hooks:
            h.remove()

        for layer_i, resid in residuals:
            normed = F.rms_norm(resid[0], (resid.shape[-1],))
            logits = (normed @ unembed.T).float()
            logits = 15.0 * torch.tanh(logits / 15.0)
            top = logits[-1].topk(top_k)
            preds = [
                {"token": decode_token(tok, int(t)), "logit": round(float(v), 3)}
                for t, v in zip(top.indices, top.values)
            ]
            layer_preds.append({"layer": layer_i, "top_k": preds})
            residuals_ref = residuals

        input_tokens = [decode_token(tok, t) for t in ids]
        results[text] = {
            "input_tokens": input_tokens,
            "n_tokens": len(ids),
            "layers": layer_preds,
        }
        print(
            f"  '{text[:40]}' — layer-by-layer prediction at last position:", flush=True
        )
        for lp in layer_preds:
            top1 = lp["top_k"][0]["token"]
            top1_logit = lp["top_k"][0]["logit"]
            print(
                f"    layer {lp['layer']:2d}: {top1:15s} (logit {top1_logit:.2f})",
                flush=True,
            )

    return results


def gla_state_probe(
    model: CPUGPT, tok, texts: list[str], device: str, top_k: int = 10
) -> dict:
    print("[explain] Running GLA state probing ...", flush=True)
    results = {}
    unembed = model.lm_head.weight.float()

    for text in texts:
        ids = encode(tok, text)[: model.cfg.seq_len - 1]
        input_tokens = [decode_token(tok, t) for t in ids]
        text_results = {"input_tokens": input_tokens, "positions": []}

        probe_positions = sorted({0, len(ids) // 2, len(ids) - 1})

        for pos in probe_positions:
            prefix_ids = ids[: pos + 1]
            prefix_idx = torch.tensor([prefix_ids], dtype=torch.long, device=device)

            state_captures = {}
            hooks = []
            gla_blocks = [(i, b) for i, b in enumerate(model.blocks) if b.is_gla]

            def make_state_hook(layer_i):
                pass

            def make_gla_hook(layer_i):
                def h(module, inp, out):
                    x = inp[0]
                    with torch.no_grad():
                        B, T, C = x.shape
                        H, D = module.n_head, module.d_head
                        CS = module.chunk
                        if T < CS:
                            return
                        q = module.q_proj(x).reshape(B, T, H, D).transpose(1, 2)
                        k = module.k_proj(x).reshape(B, T, H, D).transpose(1, 2)
                        v = module.v_proj(x).reshape(B, T, H, D).transpose(1, 2)
                        log_g = -F.softplus(module.g_proj(x).float()).transpose(1, 2)
                        n_chunks = T // CS
                        chunk_log_g = log_g.reshape(B, H, n_chunks, CS).sum(-1)
                        chunk_gate = chunk_log_g.exp()
                        k_k = F.elu(k.float()) + 1.0
                        v_f = v.float()
                        state = torch.zeros(
                            B, H, D, D, device=x.device, dtype=torch.float32
                        )
                        for c in range(n_chunks):
                            g = chunk_gate[:, :, c, None, None]
                            k_c = k_k[:, :, c * CS : (c + 1) * CS, :]
                            v_c = v_f[:, :, c * CS : (c + 1) * CS, :]
                            state = g * state + k_c.transpose(-2, -1) @ v_c
                        state_captures[layer_i] = state[0].cpu()

                return h

            for layer_i, block in gla_blocks:
                hooks.append(block.mixer.register_forward_hook(make_gla_hook(layer_i)))

            with torch.no_grad():
                _ = model(prefix_idx)
            for h in hooks:
                h.remove()

            pos_result = {
                "pos": pos,
                "token": decode_token(tok, ids[pos]),
                "layers": {},
            }
            for layer_i, state in state_captures.items():
                H, D, _ = state.shape
                mean_state = state.mean(0)
                U, S, Vh = torch.linalg.svd(mean_state)
                top_dir = U[:, 0]

                H_idx = 0
                v_proj_w = model.blocks[layer_i].mixer.v_proj.weight
                head_slice = v_proj_w.reshape(H, D, -1)[H_idx]
                dir_full = (top_dir.to(device) @ v_proj_w[:D, :].float()).cpu()
                dir_full = F.normalize(dir_full, dim=0)

                sims = (unembed.cpu() @ dir_full).float()
                top_vocab = sims.topk(top_k)
                concepts = [
                    {"token": decode_token(tok, int(t)), "sim": round(float(s), 3)}
                    for t, s in zip(top_vocab.indices, top_vocab.values)
                ]
                pos_result["layers"][f"layer_{layer_i}"] = {
                    "top_singular_value": round(float(S[0]), 4),
                    "dominant_concepts": concepts,
                }
                print(
                    f"  pos={pos} ('{decode_token(tok, ids[pos])}') layer_{layer_i}: "
                    f"state SV={S[0]:.2f}, concepts: "
                    f"{[c['token'] for c in concepts[:3]]}",
                    flush=True,
                )

            text_results["positions"].append(pos_result)
        results[text] = text_results

    return results


def integrated_gradients(
    model: CPUGPT, tok, texts: list[str], device: str, steps: int = 20
) -> dict:
    print("[explain] Running integrated gradients ...", flush=True)
    results = {}

    for text in texts:
        ids = encode(tok, text)[: model.cfg.seq_len - 1]
        if len(ids) < 2:
            continue
        idx = torch.tensor([ids], dtype=torch.long, device=device)

        with torch.no_grad():
            logits = model(idx)
        with torch.no_grad():
            x = model.wte(idx)
            x = F.rms_norm(x, (x.size(-1),))
            for block in model.blocks:
                x = block(x)
            x = model.ln_out(x)
            logits = model.lm_head(x).float()
            logits = 15.0 * torch.tanh(logits / 15.0)

        target_id = int(logits[0, -1].argmax())
        target_token = decode_token(tok, target_id)

        emb = model.wte(idx).detach()
        baseline = torch.zeros_like(emb)
        grads_accum = torch.zeros_like(emb)

        for step in range(steps):
            alpha = (step + 0.5) / steps
            interp = baseline + alpha * (emb - baseline)
            interp = interp.requires_grad_(True)

            x = F.rms_norm(interp, (interp.size(-1),))
            for block in model.blocks:
                x = block(x)
            x = model.ln_out(x)
            logit_target = (model.lm_head(x).float())[0, -1, target_id]
            logit_target = 15.0 * torch.tanh(logit_target / 15.0)
            logit_target.backward()
            grads_accum += interp.grad.detach()

        ig = ((emb - baseline) * grads_accum / steps).norm(dim=-1)[0]
        ig_norm = (ig / (ig.sum() + 1e-8)).tolist()

        input_tokens = [decode_token(tok, t) for t in ids]
        top_idx = sorted(range(len(ig_norm)), key=lambda i: ig_norm[i], reverse=True)[
            :5
        ]

        results[text] = {
            "predicted_next": target_token,
            "input_tokens": input_tokens,
            "attributions": [round(v, 4) for v in ig_norm],
            "top_influential": [
                {
                    "pos": i,
                    "token": input_tokens[i],
                    "attribution": round(ig_norm[i], 4),
                }
                for i in top_idx
            ],
        }

        print(f"  '{text[:40]}' → predicts {target_token}", flush=True)
        for item in results[text]["top_influential"][:3]:
            print(
                f"    [{item['pos']}] {item['token']:15s}  attr={item['attribution']:.3f}",
                flush=True,
            )

    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--out-dir", default="/tmp/explain")
    ap.add_argument("--ig-steps", type=int, default=20)
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    model = load_model(args.ckpt, args.device)
    tok = get_tokenizer()

    print(f"\n{'=' * 60}", flush=True)
    print("1. LOGIT LENS", flush=True)
    print("=" * 60, flush=True)
    ll = logit_lens(model, tok, SAMPLE_TEXTS, args.device)
    with open(out_dir / "logit_lens.json", "w") as f:
        json.dump(ll, f, indent=2)

    print(f"\n{'=' * 60}", flush=True)
    print("2. GLA STATE PROBING", flush=True)
    print("=" * 60, flush=True)
    gp = gla_state_probe(model, tok, SAMPLE_TEXTS, args.device)
    with open(out_dir / "gla_state_probing.json", "w") as f:
        json.dump(gp, f, indent=2)

    print(f"\n{'=' * 60}", flush=True)
    print("3. INTEGRATED GRADIENTS", flush=True)
    print("=" * 60, flush=True)
    ig = integrated_gradients(
        model, tok, SAMPLE_TEXTS, args.device, steps=args.ig_steps
    )
    with open(out_dir / "integrated_gradients.json", "w") as f:
        json.dump(ig, f, indent=2)

    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import numpy as np

        fig, axes = plt.subplots(1, 3, figsize=(18, 5))

        text0 = SAMPLE_TEXTS[0]
        n_layers = len(ll[text0]["layers"])
        n_tok = ll[text0]["n_tokens"]
        lens_matrix = np.zeros((n_layers, 1))
        layer_labels, top1_tokens = [], []
        for lp in ll[text0]["layers"]:
            layer_labels.append(f"L{lp['layer']}")
            top1_tokens.append(
                f"{lp['top_k'][0]['token']}\n({lp['top_k'][0]['logit']:.1f})"
            )
            lens_matrix[lp["layer"], 0] = lp["top_k"][0]["logit"]
        axes[0].barh(range(n_layers), lens_matrix[:, 0], color="#2196F3")
        axes[0].set_yticks(range(n_layers))
        axes[0].set_yticklabels(
            [f"L{i}: {t}" for i, t in zip(range(n_layers), top1_tokens)], fontsize=6
        )
        axes[0].set_xlabel("Top-1 logit at last position")
        axes[0].set_title(f"Logit Lens\n'{text0[:35]}...'", fontsize=9)

        if text0 in ig:
            attrs = ig[text0]["attributions"]
            toks = ig[text0]["input_tokens"]
            axes[1].bar(range(len(attrs)), attrs, color="#4CAF50")
            axes[1].set_xticks(range(len(toks)))
            axes[1].set_xticklabels(toks, rotation=45, ha="right", fontsize=6)
            axes[1].set_ylabel("Attribution (normalized)")
            axes[1].set_title(
                f"Integrated Gradients → {ig[text0]['predicted_next']}", fontsize=9
            )

        gp0 = gp.get(text0, {})
        if gp0.get("positions"):
            last_pos = gp0["positions"][-1]
            layers = list(last_pos["layers"].keys())
            svs = [last_pos["layers"][l]["top_singular_value"] for l in layers]
            axes[2].bar(range(len(layers)), svs, color="#FF5722")
            axes[2].set_xticks(range(len(layers)))
            axes[2].set_xticklabels(layers, rotation=45, ha="right", fontsize=8)
            axes[2].set_ylabel("Top singular value")
            axes[2].set_title(
                f"GLA State Magnitude\npos={last_pos['pos']} ('{last_pos['token']}')",
                fontsize=9,
            )

        plt.suptitle("FELA Mechanistic Interpretability", fontsize=11, y=1.02)
        plt.tight_layout()
        plt.savefig(out_dir / "explain_summary.png", dpi=150, bbox_inches="tight")
        plt.close()
        print(f"\n[explain] Summary plot → {out_dir}/explain_summary.png", flush=True)
    except ImportError:
        print("[explain] matplotlib not available — JSON only", flush=True)

    print(f"\n[explain] Done. Outputs in {out_dir}/", flush=True)
    for p in sorted(out_dir.iterdir()):
        print(f"  {p.name}  ({p.stat().st_size // 1024} KB)", flush=True)


if __name__ == "__main__":
    main()