import torch
import torch.nn.functional as F
import json
from pathlib import Path
from transformer_lens import HookedTransformer
from sae import SparseAutoencoder
from feature_analysis import get_causal_label, load_sae, load_norm, normalize
from model_config import get_sae_layer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
N_SAMPLES = 10


def _generate_sentences(label: str, client, n: int = N_SAMPLES) -> dict:
    """Ask GPT-4o-mini to generate positive and negative sentences for a feature label."""
    prompt = f"""You are helping evaluate a sparse autoencoder feature with the label: "{label}"

Generate {n} short sentences (5-15 words each) where this feature SHOULD fire strongly,
and {n} short sentences where this feature should NOT fire.

Reply ONLY with a JSON object in this exact format, no markdown:
{{
  "positive": ["sentence1", "sentence2", ...],
  "negative": ["sentence1", "sentence2", ...]
}}"""

    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=800,
        temperature=0.7,
    )
    raw = resp.choices[0].message.content.strip()
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
    return json.loads(raw.strip())


def _get_feature_activation(sentence: str, feature_idx: int, model: HookedTransformer, sae: SparseAutoencoder, model_id: str = "llama-3.2-1b") -> float:
    """Run a sentence through the model and return the max activation of a feature across all token positions."""
    sae_layer = get_sae_layer(model_id)
    tokens = model.to_tokens(sentence)
    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]  # [seq, d_model]
    acts = sae.encode(normalize(resid, model_id))              # [seq, n_features]
    return float(acts[:, feature_idx].max().item())


def _cohen_d_score(pos_acts: list[float], neg_acts: list[float]) -> float:
    """Cohen's d between positive and negative activation distributions, clipped to [0, 1]."""
    if not pos_acts or not neg_acts:
        return 0.0
    pos_t = torch.tensor(pos_acts, dtype=torch.float32)
    neg_t = torch.tensor(neg_acts, dtype=torch.float32)
    mu_pos = pos_t.mean().item()
    mu_neg = neg_t.mean().item()
    pooled_std = torch.cat([pos_t, neg_t]).std().item() + 1e-6
    raw = (mu_pos - mu_neg) / pooled_std
    return round(float(max(0.0, min(1.0, raw))), 4)


def _feature_purity_score(sentences: list[str], client) -> float | None:
    """
    Measure how monosemantic a feature is by computing the mean pairwise cosine
    similarity of embeddings for the sentences that activate it.

    High score  => feature fires for one coherent concept (monosemantic).
    Low score   => feature fires for unrelated contexts (polysemantic).

    Uses OpenAI text-embedding-3-large. Remaps cosine sim from [-1, 1] to [0, 1].
    Returns None if client is unavailable or embedding fails.
    """
    if not sentences or client is None:
        return None

    try:
        resp = client.embeddings.create(
            model="text-embedding-3-large",
            input=sentences[:10],
        )
        vecs = torch.tensor(
            [e.embedding for e in resp.data], dtype=torch.float32
        )  # [n, d_embed]
        vecs = F.normalize(vecs, dim=-1)

        sim_matrix = vecs @ vecs.T  # [n, n], cosine similarities

        n = vecs.shape[0]
        if n < 2:
            return None

        # Mean of upper triangle only (exclude diagonal self-similarity)
        upper_mask = torch.ones(n, n, dtype=torch.bool).triu(diagonal=1)
        mean_sim = sim_matrix[upper_mask].mean().item()

        # Remap [-1, 1] -> [0, 1]
        purity = (mean_sim + 1.0) / 2.0
        return round(float(purity), 4)

    except Exception as e:
        print(f"[purity] embedding error: {e}", flush=True)
        return None


# ── MUI helpers ───────────────────────────────────────────────────────────────

def _kl_div(p: torch.Tensor, q: torch.Tensor) -> float:
    """KL(p || q) where both are probability distributions."""
    p = p.clamp(min=1e-10)
    q = q.clamp(min=1e-10)
    return float((p * (p / q).log()).sum().item())


def _entropy(p: torch.Tensor) -> float:
    """Shannon entropy of a probability distribution."""
    p = p.clamp(min=1e-10)
    return float(-(p * p.log()).sum().item())


def run_mui_score(
    feature_idx: int,
    prompt: str,
    model: HookedTransformer,
    sae: SparseAutoencoder,
    n_positions: int = 8,
    model_id: str = "llama-3.2-1b",
) -> dict:
    """
    Model Utilization Index (MUI): measures how much ablating a feature shifts
    the model's final output distribution, normalized by baseline entropy.

    MUI = KL(p_baseline || p_ablated) / H(p_baseline), clipped to [0, 1]

    High MUI => feature causally matters to the output.
    Low MUI  => feature fires but is largely ignored downstream.

    Returns:
        {
            feature_idx,
            score,           # MUI clipped to [0, 1]
            per_position,    # list of { position, activation, kl_divergence, context }
            mean_kl,
            baseline_entropy,
        }
    """
    sae_layer = get_sae_layer(model_id)
    tokens = model.to_tokens(prompt)

    with torch.no_grad():
        baseline_logits, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
        )

    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]  # [seq, d_model]
    acts = sae.encode(normalize(resid, model_id))              # [seq, n_features]

    feat_acts = acts[:, feature_idx]
    top_k = min(n_positions, feat_acts.shape[0])
    top_positions = feat_acts.topk(top_k).indices.tolist()
    top_positions = [p for p in top_positions if feat_acts[p].item() > 0.01]

    if not top_positions:
        return {
            "feature_idx": feature_idx,
            "score": 0.0,
            "per_position": [],
            "mean_kl": 0.0,
            "baseline_entropy": None,
        }

    baseline_probs = torch.softmax(baseline_logits[0, -1], dim=-1)
    baseline_H = _entropy(baseline_probs)

    per_position = []
    kl_vals = []

    for pos in top_positions:
        act_val = feat_acts[pos].item()

        def ablate(value, hook, pos=pos):
            resid_here = value[0, pos]
            f = sae.encode(normalize(resid_here.unsqueeze(0), model_id))[0]
            f_abl = f.clone()
            f_abl[feature_idx] = 0.0
            orig = sae.decode(f.unsqueeze(0))[0]
            abl  = sae.decode(f_abl.unsqueeze(0))[0]
            value[0, pos] = resid_here + (abl - orig)
            return value

        with torch.no_grad():
            abl_logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[(f"blocks.{sae_layer}.hook_resid_post", ablate)]
            )

        abl_probs = torch.softmax(abl_logits[0, -1], dim=-1)
        kl = _kl_div(baseline_probs, abl_probs)
        kl_vals.append(kl)

        context = model.tokenizer.decode(
            tokens[0, max(0, pos - 3):pos + 4].tolist()
        )
        per_position.append({
            "position": pos,
            "activation": round(act_val, 3),
            "kl_divergence": round(kl, 4),
            "context": context,
        })

    mean_kl = sum(kl_vals) / len(kl_vals)
    mui = round(float(min(mean_kl / max(baseline_H, 1e-6), 1.0)), 4)

    torch.cuda.synchronize()

    return {
        "feature_idx": feature_idx,
        "score": mui,
        "per_position": sorted(per_position, key=lambda x: x["kl_divergence"], reverse=True),
        "mean_kl": round(mean_kl, 4),
        "baseline_entropy": round(baseline_H, 4),
    }


# ── Main entry point ──────────────────────────────────────────────────────────

def run_interp_score(
    feature_idx: int,
    prompt: str,
    model: HookedTransformer,
    sae: SparseAutoencoder,
    client,
    n_samples: int = N_SAMPLES,
    model_id: str = "llama-3.2-1b",
) -> dict:
    """
    Evaluate how well a feature label predicts where the feature fires (InterpScore),
    how monosemantic the feature is (FeaturePurityScore), and how much it causally
    influences the model's output (MUI).

    Returns:
        {
            feature_idx,
            label,
            score,              # InterpScore: Cohen's d clipped to [0, 1]
            purity_score,       # FeaturePurityScore: mean pairwise cosine sim remapped to [0, 1]
            mui_score,          # Model Utilization Index: KL / H clipped to [0, 1]
            mui_per_position,   # list of { position, activation, kl_divergence, context }
            mui_mean_kl,
            baseline_entropy,
            positive_mean,
            negative_mean,
            positive_examples:  [{ sentence, activation }],
            negative_examples:  [{ sentence, activation }],
        }
    """
    label = get_causal_label(feature_idx, prompt, model, sae, client, model_id=model_id)

    # Always run MUI — it doesn't need the OpenAI client
    mui_result = run_mui_score(feature_idx, prompt, model, sae, model_id=model_id)

    if client is None:
        return {
            "feature_idx": feature_idx,
            "label": label,
            "score": None,
            "purity_score": None,
            "mui_score": mui_result["score"],
            "mui_per_position": mui_result["per_position"],
            "mui_mean_kl": mui_result["mean_kl"],
            "baseline_entropy": mui_result["baseline_entropy"],
            "error": "OpenAI client not configured. Set OPENAI_API_KEY.",
            "positive_examples": [],
            "negative_examples": [],
            "positive_mean": None,
            "negative_mean": None,
        }

    try:
        sentences = _generate_sentences(label, client, n=n_samples)
    except Exception as e:
        return {
            "feature_idx": feature_idx,
            "label": label,
            "score": None,
            "purity_score": None,
            "mui_score": mui_result["score"],
            "mui_per_position": mui_result["per_position"],
            "mui_mean_kl": mui_result["mean_kl"],
            "baseline_entropy": mui_result["baseline_entropy"],
            "error": f"Sentence generation failed: {e}",
            "positive_examples": [],
            "negative_examples": [],
            "positive_mean": None,
            "negative_mean": None,
        }

    positive_results = []
    for sent in sentences.get("positive", []):
        try:
            act = _get_feature_activation(sent, feature_idx, model, sae, model_id=model_id)
            positive_results.append({"sentence": sent, "activation": round(act, 4)})
        except Exception as e:
            print(f"[interp_score] positive sentence failed: {e}", flush=True)

    negative_results = []
    for sent in sentences.get("negative", []):
        try:
            act = _get_feature_activation(sent, feature_idx, model, sae, model_id=model_id)
            negative_results.append({"sentence": sent, "activation": round(act, 4)})
        except Exception as e:
            print(f"[interp_score] negative sentence failed: {e}", flush=True)

    pos_acts = [r["activation"] for r in positive_results]
    neg_acts = [r["activation"] for r in negative_results]

    score      = _cohen_d_score(pos_acts, neg_acts)
    pos_mean   = round(sum(pos_acts) / len(pos_acts), 4) if pos_acts else None
    neg_mean   = round(sum(neg_acts) / len(neg_acts), 4) if neg_acts else None

    purity_score = _feature_purity_score(
        [r["sentence"] for r in positive_results], client
    )

    torch.cuda.synchronize()

    return {
        "feature_idx": feature_idx,
        "label": label,
        "score": score,
        "purity_score": purity_score,
        "mui_score": mui_result["score"],
        "mui_per_position": mui_result["per_position"],
        "mui_mean_kl": mui_result["mean_kl"],
        "baseline_entropy": mui_result["baseline_entropy"],
        "positive_mean": pos_mean,
        "negative_mean": neg_mean,
        "positive_examples": sorted(positive_results, key=lambda x: x["activation"], reverse=True),
        "negative_examples": sorted(negative_results, key=lambda x: x["activation"], reverse=True),
    }
