import torch
import json
import hashlib
import os
import contextlib
import io
from pathlib import Path
from transformer_lens import HookedTransformer
from sae import SparseAutoencoder
from model_config import get_config, get_sae_source, get_sae_path, get_norm_path

SAE_LAYER       = 8
SAE_SAVE_PATH   = Path.home() / "inspection-backend" / "sae" / "sae_layer8.pt"
NORM_CACHE_PATH = Path.home() / "inspection-backend" / "sae" / "acts_norm.pt"
TOP_K_FEATURES  = 10
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_sae_cache: dict[str, SparseAutoencoder] = {}  # model_id -> loaded SAE
_norm_cache: dict[str, dict | None] = {}        # model_id -> norm tensors
_session_label_cache = {}  # (feature_idx, prompt_hash, model_id) -> label

# Legacy globals kept so callers that don't pass model_id still work
_sae = None
_norm = None


def _load_sae_native(model_id: str) -> SparseAutoencoder:
    sae_path = get_sae_path(model_id)
    if not sae_path.exists():
        raise FileNotFoundError(f"SAE not found at {sae_path}. Run train_sae.py first.")
    return SparseAutoencoder.load(sae_path, device=DEVICE)


def _load_sae_from_hf(model_id: str) -> SparseAutoencoder:
    """Download SAE weights from HuggingFace Hub and load into SparseAutoencoder.
    Used by both 'neuronpedia' and 'huggingface' sources â€” same weight origin,
    different label strategy."""
    cfg = get_config(model_id)
    sae_path = get_sae_path(model_id)

    if not sae_path.exists():
        print(f"[SAE] downloading {model_id} weights from HuggingFace...", flush=True)
        try:
            from huggingface_hub import hf_hub_download
        except ImportError:
            raise RuntimeError("huggingface_hub not installed. Run: pip install huggingface-hub")

        # neuronpedia source uses neuronpedia_hf_* keys; huggingface source uses hf_* keys
        hf_repo = cfg.get("hf_repo") or cfg.get("neuronpedia_hf_repo")
        hf_file = cfg.get("hf_file") or cfg.get("neuronpedia_hf_file")
        if not hf_repo or not hf_file:
            raise ValueError(f"[SAE] no HF repo/file configured for model '{model_id}'")

        sae_path.parent.mkdir(parents=True, exist_ok=True)
        with contextlib.redirect_stderr(io.StringIO()):
            local = hf_hub_download(
                repo_id=hf_repo,
                filename=hf_file,
                repo_type="model",
            )
        import shutil
        shutil.copy(local, sae_path)
        print(f"[SAE] saved to {sae_path}", flush=True)

    # Load safetensors and remap keys into SparseAutoencoder
    from safetensors import safe_open
    tensors = {}
    with safe_open(str(sae_path), framework="pt", device="cpu") as f:
        for k in f.keys():
            tensors[k] = f.get_tensor(k)

    # Neuronpedia/SAELens key mapping -> Aquin SparseAutoencoder keys
    KEY_MAP = {
        "encoder.weight": "W_enc",   # shape [d_model, n_features] after transpose
        "encoder.bias":   "b_enc",
        "W_dec":          "W_dec",
        "b_dec":          "b_dec",
        "b_pre":          "b_pre",
    }

    d_model   = cfg["d_model"]
    # infer n_features from the encoder weight
    enc_w = tensors["encoder.weight"] if "encoder.weight" in tensors else tensors.get("W_enc")
    if enc_w is None:
        raise KeyError(f"Cannot find encoder weight in SAE checkpoint for {model_id}. Keys: {list(tensors.keys())}")
    n_features = enc_w.shape[0]  # [n_features, d_model] in SAELens convention

    sae = SparseAutoencoder(d_model=d_model, n_features=n_features)
    sd = sae.state_dict()

    for src_key, dst_key in KEY_MAP.items():
        if src_key not in tensors:
            continue
        t = tensors[src_key].float()
        if src_key == "encoder.weight":
            # SAELens: [n_features, d_model]; Aquin W_enc: [d_model, n_features]
            t = t.T
        sd[dst_key] = t

    sae.load_state_dict(sd)
    sae.to(DEVICE)
    sae.eval()
    return sae


def load_sae(model_id: str = "llama-3.2-1b") -> SparseAutoencoder:
    global _sae
    if model_id not in _sae_cache:
        source = get_sae_source(model_id)
        if source in ("neuronpedia", "huggingface"):
            _sae_cache[model_id] = _load_sae_from_hf(model_id)
        else:
            _sae_cache[model_id] = _load_sae_native(model_id)
        # keep legacy global in sync for the default model
        if model_id == "llama-3.2-1b":
            _sae = _sae_cache[model_id]
    return _sae_cache[model_id]


def load_norm(model_id: str = "llama-3.2-1b") -> dict | None:
    global _norm
    if model_id not in _norm_cache:
        norm_path = get_norm_path(model_id)
        if norm_path and norm_path.exists():
            d = torch.load(norm_path, map_location=DEVICE)
            _norm_cache[model_id] = {"mean": d["mean"].to(DEVICE), "std": d["std"].to(DEVICE)}
        else:
            _norm_cache[model_id] = None
        if model_id == "llama-3.2-1b":
            _norm = _norm_cache[model_id]
    return _norm_cache[model_id]


def normalize(x: torch.Tensor, model_id: str = "llama-3.2-1b") -> torch.Tensor:
    n = load_norm(model_id)
    if n is None:
        return x
    return (x - n["mean"]) / n["std"]


def label_feature_causally(fi: int, prompt: str, model: HookedTransformer, sae: SparseAutoencoder, client, top_k: int = 5, model_id: str = "llama-3.2-1b") -> str:
    from model_config import get_sae_layer
    sae_layer = get_sae_layer(model_id)

    tokens = model.to_tokens(prompt[:512])
    if tokens.shape[1] < 4:
        return f"feature_{fi}"

    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
            return_type=None,
        )
    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]
    acts = sae.encode(normalize(resid, model_id))  # [seq, n_features]

    firing_pos = (acts[:, fi] > 0.5).nonzero(as_tuple=True)[0].tolist()
    if not firing_pos:
        top_pos = int(acts[:, fi].argmax().item())
        if acts[top_pos, fi].item() < 0.01:
            return f"feature_{fi}"
        firing_pos = [top_pos]

    causal_examples = []

    for pos in firing_pos[:3]:
        act_val = acts[pos, fi].item()

        with torch.no_grad():
            baseline_logits = model(tokens)[0, -1]
        baseline_probs = torch.softmax(baseline_logits, dim=-1)

        def ablate_feature(value, hook, pos=pos):
            resid_here = value[0, pos]
            feat_acts = sae.encode(normalize(resid_here.unsqueeze(0), model_id))[0]
            feat_acts_ablated = feat_acts.clone()
            feat_acts_ablated[fi] = 0.0
            original_contribution = sae.decode(feat_acts)
            ablated_contribution = sae.decode(feat_acts_ablated)
            value[0, pos] = resid_here + (ablated_contribution - original_contribution)
            return value

        with torch.no_grad():
            ablated_logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[(f"blocks.{sae_layer}.hook_resid_post", ablate_feature)]
            )[0, -1]
        ablated_probs = torch.softmax(ablated_logits, dim=-1)

        delta = baseline_probs - ablated_probs
        topk_boosted = delta.topk(top_k)
        topk_suppressed = (-delta).topk(top_k)

        boosted = [
            (model.tokenizer.decode([idx.item()]).strip(), round(val.item(), 4))
            for idx, val in zip(topk_boosted.indices, topk_boosted.values)
            if val > 0.001
        ]
        suppressed = [
            (model.tokenizer.decode([idx.item()]).strip(), round(val.item(), 4))
            for idx, val in zip(topk_suppressed.indices, topk_suppressed.values)
            if val > 0.001
        ]

        context = model.tokenizer.decode(tokens[0, max(0, pos - 4):pos + 5].tolist())
        pivot = model.tokenizer.decode([tokens[0, pos].item()]).strip()

        causal_examples.append({
            "activation": round(act_val, 3),
            "context": context,
            "pivot": pivot,
            "boosts": boosted,
            "suppresses": suppressed,
        })

    if not causal_examples:
        return f"feature_{fi}"

    causal_examples.sort(key=lambda x: x["activation"], reverse=True)

    formatted = "\n".join(
        f'  Context: "...{ex["context"]}..." (token: "{ex["pivot"]}", activation: {ex["activation"]})\n'
        f'  Causally boosts: {ex["boosts"]}\n'
        f'  Causally suppresses: {ex["suppresses"]}'
        for ex in causal_examples
    )

    prompt_text = f"""A sparse autoencoder feature causally influences a language model's predictions as follows:

{formatted}

Based on what this feature causally promotes and suppresses in this context, give a concise 2-6 word label for its functional role. Good label examples: "promotes plural nouns", "suppresses hedging language", "boosts location names after prepositions".

Reply with ONLY the label."""

    try:
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt_text}],
            max_tokens=20,
            temperature=0.0,
        )
        return resp.choices[0].message.content.strip().strip('"')
    except Exception as e:
        print(f"[label] error on feature {fi}: {e}", flush=True)
        return f"feature_{fi}"


def _fetch_neuronpedia_label(fi: int, model_id: str) -> str | None:
    """Hit the Neuronpedia explanation API. Returns label string or None on miss/error."""
    try:
        import requests
        cfg = get_config(model_id)
        slug = cfg.get("neuronpedia_slug")
        if not slug:
            return None
        url = f"https://www.neuronpedia.org/api/feature/{slug}/{fi}"
        r = requests.get(url, timeout=8)
        if r.status_code != 200:
            return None
        data = r.json()
        # Neuronpedia returns explanations as a list, ranked by score
        explanations = data.get("explanations") or []
        if explanations:
            return explanations[0].get("description") or explanations[0].get("label")
        return None
    except Exception as e:
        print(f"[label/neuronpedia] feature {fi}: {e}", flush=True)
        return None


def get_causal_label(fi: int, prompt: str, model: HookedTransformer, sae: SparseAutoencoder, client, model_id: str = "llama-3.2-1b") -> str:
    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()[:8]
    key = (fi, prompt_hash, model_id)
    if key in _session_label_cache:
        return _session_label_cache[key]

    source = get_sae_source(model_id)
    if source == "neuronpedia":
        label = _fetch_neuronpedia_label(fi, model_id) or f"feature_{fi}"
    else:
        # native + huggingface: run causal ablation, label with GPT-4o-mini
        label = label_feature_causally(fi, prompt, model, sae, client, model_id=model_id)

    _session_label_cache[key] = label
    return label


def _run_sae_pass(prompt: str, response: str, model: HookedTransformer, top_k: int = TOP_K_FEATURES, model_id: str = "llama-3.2-1b") -> dict:
    """
    Pure SAE forward pass. No labeling at all.
    Returns raw feature indices and activations for all positions,
    plus token strings and attribution structure (with feature_idx but no labels).
    """
    from model_config import get_sae_layer
    sae_layer = get_sae_layer(model_id)
    sae = load_sae(model_id)

    full_ctx = f"{prompt}\n{response}"
    tokens = model.to_tokens(full_ctx)

    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
            return_type=None,
        )
    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]

    with torch.no_grad():
        feature_acts = sae.encode(normalize(resid, model_id))

    seq_len = resid.shape[0]
    prompt_ctx = f"{prompt}\n"
    prompt_ctx_len = model.to_tokens(prompt_ctx, prepend_bos=True).shape[1]

    all_strs = [model.to_string([tokens[0, i].item()]) for i in range(seq_len)]
    prompt_strs = all_strs[1:prompt_ctx_len]
    response_strs = all_strs[prompt_ctx_len:]
    prompt_idxs = list(range(1, prompt_ctx_len))
    response_idxs = list(range(prompt_ctx_len, seq_len))

    def feats_for_unlabeled(positions):
        out = []
        for pos in positions:
            acts = feature_acts[pos]
            topk = acts.topk(top_k)
            out.append([
                {
                    "feature_idx": int(i),
                    "activation": round(float(v), 3),
                    "label": f"feature_{int(i)}",  # placeholder, filled in later
                }
                for i, v in zip(topk.indices, topk.values) if v > 0.001
            ])
        return out

    prompt_features = feats_for_unlabeled(prompt_idxs)
    response_features = feats_for_unlabeled(response_idxs)

    resp_acts = feature_acts[prompt_ctx_len:]
    resp_top = resp_acts.max(0).values.topk(20)
    top_response_features = []
    for idx, val in zip(resp_top.indices.tolist(), resp_top.values.tolist()):
        if val < 0.001:
            continue
        best_pos = int(resp_acts[:, idx].argmax().item())
        top_response_features.append({
            "feature_idx": idx,
            "activation": round(val, 3),
            "label": f"feature_{idx}",  # placeholder
            "token": response_strs[best_pos].strip() if best_pos < len(response_strs) else "",
            "token_idx": best_pos,
        })

    attribution = []
    for ri, rpos in enumerate(response_idxs):
        resp_tok = response_strs[ri].strip() if ri < len(response_strs) else ""
        if not resp_tok or resp_tok in ("the", "a", "an", "is", "of", ".", ","):
            continue
        resp_feats = feature_acts[rpos]
        top_resp = resp_feats.topk(top_k)
        driven_by = []
        for fidx, fval in zip(top_resp.indices.tolist(), top_resp.values.tolist()):
            if fval < 0.001:
                continue
            prompt_feat_acts = feature_acts[prompt_idxs, fidx]
            active = (prompt_feat_acts > 0.001).nonzero(as_tuple=True)[0].tolist()
            if not active:
                continue
            driven_by.append({
                "feature_idx": fidx,
                "label": f"feature_{fidx}",  # placeholder
                "activation": round(fval, 3),
                "also_in_prompt_positions": active,
                "also_in_prompt_tokens": [prompt_strs[p].strip() for p in active if p < len(prompt_strs)],
            })
        if driven_by:
            attribution.append({
                "response_token": resp_tok,
                "response_ti": ri,
                "driven_by_features": sorted(driven_by, key=lambda x: x["activation"], reverse=True)[:5],
            })

    torch.cuda.synchronize()

    return {
        "prompt_tokens": [t.strip() for t in prompt_strs],
        "response_tokens": [t.strip() for t in response_strs],
        "prompt_features": prompt_features,
        "response_features": response_features,
        "top_response_features": top_response_features,
        "attribution": attribution,
    }


def run_feature_analysis_unlabeled(prompt: str, response: str, model: HookedTransformer, model_id: str = "llama-3.2-1b") -> dict:
    """
    Public entry point for the unlabeled SAE pass.
    Returns the full result structure with feature_idx placeholders for labels.
    Fast -- no GPT calls.
    """
    return _run_sae_pass(prompt, response, model, model_id=model_id)


def run_feature_analysis(prompt: str, response: str, model: HookedTransformer, client, top_k: int = TOP_K_FEATURES, model_id: str = "llama-3.2-1b") -> dict:
    """
    Legacy full analysis with labeling inline.
    Kept for backwards compatibility -- prefer the streaming approach via
    run_feature_analysis_unlabeled + get_causal_label per feature.
    """
    sae = load_sae(model_id)
    result = _run_sae_pass(prompt, response, model, top_k, model_id=model_id)

    def fill_labels(features_list):
        for pos_feats in features_list:
            for f in pos_feats:
                f["label"] = get_causal_label(f["feature_idx"], prompt, model, sae, client, model_id=model_id)

    fill_labels(result["prompt_features"])
    fill_labels(result["response_features"])

    for f in result["top_response_features"]:
        f["label"] = get_causal_label(f["feature_idx"], prompt, model, sae, client, model_id=model_id)

    for attr in result["attribution"]:
        for f in attr["driven_by_features"]:
            f["label"] = get_causal_label(f["feature_idx"], prompt, model, sae, client, model_id=model_id)

    return result

