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, FNOSeqMixer, GLAMixer, gpt2_small_config


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

        bucket, key = ckpt[5:].split("/", 1)
        local = "/tmp/interp_ckpt.pt"
        print(f"[interp] Downloading s3://{bucket}/{key} ...", flush=True)
        boto3.client("s3").download_file(bucket, key, local)
        ckpt = local
    elif ckpt.startswith("hf://"):
        from huggingface_hub import hf_hub_download

        path = ckpt[5:]
        filename = path.rsplit("/", 1)[-1]
        repo = path.rsplit("/", 1)[0]
        print(f"[interp] Downloading hf://{repo}/{filename} ...", flush=True)
        ckpt = hf_hub_download(
            repo_id=repo, filename=filename, token=os.environ.get("HF_TOKEN")
        )
    print(f"[interp] Loading checkpoint from {ckpt}", flush=True)
    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_n_embd = state["wte.weight"].shape[1]
        if ckpt_n_embd != cfg.n_embd:
            print(
                f"[interp] Checkpoint n_embd={ckpt_n_embd} differs from default {cfg.n_embd}; "
                f"rebuilding model to match checkpoint",
                flush=True,
            )
            cfg = gpt2_small_config(seq_len=1024, n_embd=ckpt_n_embd)
            cfg.n_head = max(1, ckpt_n_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()

    model_state = model.state_dict()
    filtered = {
        k: v
        for k, v in state.items()
        if k in model_state and v.shape == model_state[k].shape
    }
    skipped = [k for k in state if k not in filtered and k in model_state]
    if skipped:
        print(f"[interp] Skipped {len(skipped)} size-mismatched keys", flush=True)
    model.load_state_dict(filtered, strict=False)
    print(
        f"[interp] Checkpoint loaded — {model.param_count() / 1e6:.1f}M params",
        flush=True,
    )
    return model


def analyze_fno_spectrum(model: CPUGPT, out_dir: Path) -> dict:
    print("[interp] Computing FNO filter spectra ...", flush=True)
    spectra = {}
    fno_layers = [(i, b) for i, b in enumerate(model.blocks) if not b.is_gla]

    for layer_idx, block in fno_layers:
        with torch.no_grad():
            h = block.mixer.filter_td.float().cpu()
            C, M = h.shape
            h_pad = F.pad(h, (0, M))
            Hf = torch.fft.rfft(h_pad, dim=1)
            mag = Hf.abs().mean(0).numpy().tolist()
            spectra[f"layer_{layer_idx}"] = {
                "magnitudes": mag,
                "n_modes": M,
                "n_channels": C,
            }

    out_json = out_dir / "fno_filter_spectrum.json"
    with open(out_json, "w") as f:
        json.dump(spectra, f, indent=2)
    print(f"[interp] FNO spectra saved → {out_json}", flush=True)

    try:
        import matplotlib

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

        fig, axes = plt.subplots(3, 3, figsize=(12, 9))
        axes = axes.flatten()
        for ax_i, (name, data) in enumerate(spectra.items()):
            if ax_i >= len(axes):
                break
            mag = np.array(data["magnitudes"])
            freqs = np.arange(len(mag))
            axes[ax_i].plot(freqs, mag, lw=1.2, color="#2196F3")
            axes[ax_i].set_title(name, fontsize=9)
            axes[ax_i].set_xlabel("freq index", fontsize=7)
            axes[ax_i].set_ylabel("|H(f)|", fontsize=7)
            axes[ax_i].set_yscale("log")
            axes[ax_i].grid(True, alpha=0.3)
        for ax_i in range(len(spectra), len(axes)):
            axes[ax_i].set_visible(False)
        fig.suptitle(
            "FNO Learned Filter Spectra (mean |H(f)| across channels)", fontsize=11
        )
        plt.tight_layout()
        out_png = out_dir / "fno_filter_spectrum.png"
        plt.savefig(out_png, dpi=150, bbox_inches="tight")
        plt.close()
        print(f"[interp] FNO spectrum plot saved → {out_png}", flush=True)
    except ImportError:
        print("[interp] matplotlib not available — JSON only", flush=True)

    return spectra


SAMPLE_TEXTS = [
    "The quick brown fox jumps over the lazy dog near the river bank.",
    "Machine learning models can now process language at superhuman speeds.",
    "Fourier analysis decomposes signals into their constituent frequencies.",
]


def analyze_gla_gates(model: CPUGPT, device: str, out_dir: Path) -> dict:
    print("[interp] Capturing GLA gate activations ...", flush=True)
    from transformers import GPT2Tokenizer

    enc = GPT2Tokenizer.from_pretrained("gpt2")

    gate_captures: dict[str, list] = {}
    hooks = []

    def make_hook(layer_name):
        def hook(module, args, output):
            x = args[0]
            with torch.no_grad():
                log_g = -F.softplus(module.g_proj(x).float())
                gate_captures[layer_name] = log_g[0].cpu().numpy().tolist()

        return hook

    gla_layers = [(i, b) for i, b in enumerate(model.blocks) if b.is_gla]
    for layer_idx, block in gla_layers:
        name = f"layer_{layer_idx}_gla"
        h = block.mixer.register_forward_hook(make_hook(name))
        hooks.append(h)

    gla_chunk = getattr(model.cfg, "gla_chunk", 64)

    results = {}
    for text in SAMPLE_TEXTS:
        gate_captures.clear()
        tokens = enc.encode(text, add_special_tokens=False)[: model.cfg.seq_len]
        if len(tokens) % gla_chunk != 0 or len(tokens) == 0:
            pad_len = gla_chunk - (len(tokens) % gla_chunk)
            if pad_len == gla_chunk:
                pad_len = 0
            tokens = tokens + [0] * pad_len
        if len(tokens) < gla_chunk:
            tokens = tokens + [0] * (gla_chunk - len(tokens))
        idx = torch.tensor([tokens], dtype=torch.long, device=device)
        with torch.no_grad():
            _ = model(idx)
        results[text[:40]] = {name: vals for name, vals in gate_captures.items()}

    for h in hooks:
        h.remove()

    out_json = out_dir / "gla_gate_heatmap.json"
    with open(out_json, "w") as f:
        json.dump(results, f, indent=2)
    print(f"[interp] GLA gate data saved → {out_json}", flush=True)

    try:
        import matplotlib

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

        first_text = list(results.keys())[0]
        first_layer = list(results[first_text].keys())[0]
        gates = np.array(results[first_text][first_layer])
        decay = np.exp(gates)

        enc2 = GPT2Tokenizer.from_pretrained("gpt2")
        tokens = enc2.encode(SAMPLE_TEXTS[0], add_special_tokens=False)[
            : gates.shape[0]
        ]
        tok_labels = [enc2.decode([t]) for t in tokens]

        fig, ax = plt.subplots(figsize=(14, 4))
        im = ax.imshow(
            decay.T,
            aspect="auto",
            cmap="RdYlGn",
            vmin=0,
            vmax=1,
            interpolation="nearest",
        )
        ax.set_xticks(range(len(tok_labels)))
        ax.set_xticklabels(tok_labels, rotation=45, ha="right", fontsize=7)
        ax.set_ylabel("Head index", fontsize=9)
        ax.set_title(
            f"GLA Gate Decay Rates — {first_layer}\n(green=remember, red=forget)",
            fontsize=10,
        )
        plt.colorbar(im, ax=ax, fraction=0.02, pad=0.01)
        plt.tight_layout()
        out_png = out_dir / "gla_gate_heatmap.png"
        plt.savefig(out_png, dpi=150, bbox_inches="tight")
        plt.close()
        print(f"[interp] GLA gate heatmap saved → {out_png}", flush=True)
    except ImportError:
        print("[interp] matplotlib not available — JSON only", flush=True)

    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--ckpt",
        default="",
        help="Checkpoint path: local .pt, s3://bucket/key, or hf://repo/file",
    )
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--out-dir", default="/tmp/interpretability")
    args = ap.parse_args()

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

    model = load_checkpoint(args.ckpt, args.device)

    analyze_fno_spectrum(model, out_dir)
    analyze_gla_gates(model, args.device, out_dir)

    print(f"\n[interp] All 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()