"""
FGAA (Feature Guided Activation Addition) steering method.

Implements the method from "Interpretable Steering of Large Language Models
with Feature Guided Activation Additions" (arXiv:2501.09929).

Uses Gemma Scope SAEs and pre-computed effect approximators.
"""

from __future__ import annotations

import json
import tempfile
from argparse import Namespace
from pathlib import Path
from typing import TYPE_CHECKING

import torch
from huggingface_hub import hf_hub_download

from wisent.comparison.utils import apply_steering_to_model, remove_steering, convert_to_lm_eval_format

if TYPE_CHECKING:
    from wisent.core.models.wisent_model import WisentModel

__all__ = ["generate_steering_vector", "apply_steering_to_model", "remove_steering", "convert_to_lm_eval_format"]


# Known BOS feature indices from paper (Appendix G)
# These features activate most strongly on the BOS token
# Run detect_bos_features.py to verify or find features for other models
BOS_FEATURES = {
    "google/gemma-2-2b": [11087, 3220, 11752, 12160, 11498],
    "google/gemma-2-9b": [],  # Not listed in paper
}


# SAE configurations for supported Gemma models
SAE_CONFIGS = {
    "google/gemma-2-2b": {
        "sae_release": "gemma-scope-2b-pt-res-canonical",
        "sae_id_template": "layer_{layer}/width_16k/canonical",
        "adapter_file": "adapter_2b_layer_12.pt",
        "num_layers": 26,
        "default_layer": 12,  # Paper uses layer 12 (effect approximator only available for this layer)
        "d_model": 2304,
        "d_sae": 16384,
    },
    "google/gemma-2-9b": {
        "sae_release": "gemma-scope-9b-pt-res-canonical",
        "sae_id_template": "layer_{layer}/width_16k/canonical",
        "adapter_file": "adapter_9b_layer_12.pt",
        "num_layers": 42,
        "default_layer": 12,  # Paper uses layer 12 (effect approximator only available for this layer)
        "d_model": 3584,
        "d_sae": 16384,
    },
}


def load_sae(model_name: str, layer_idx: int, device: str = "cuda:0"):
    """
    Load Gemma Scope SAE for a specific layer.

    Args:
        model_name: HuggingFace model name (e.g., 'google/gemma-2-2b')
        layer_idx: Layer index to load SAE for
        device: Device to load SAE on

    Returns:
        SAE object from sae_lens
    """
    from sae_lens import SAE

    if model_name not in SAE_CONFIGS:
        raise ValueError(f"No SAE config for model '{model_name}'. Supported: {list(SAE_CONFIGS.keys())}")

    config = SAE_CONFIGS[model_name]
    sae_id = config["sae_id_template"].format(layer=layer_idx)

    print(f"   Loading SAE from {config['sae_release']} / {sae_id}")
    sae, _, sparsity = SAE.from_pretrained(
        release=config["sae_release"],
        sae_id=sae_id,
        device=device,
    )

    return sae, sparsity


def load_effect_approximator(model_name: str, device: str = "cuda:0") -> tuple[torch.Tensor, torch.Tensor]:
    """
    Load the pre-trained effect approximator (adapter) from HuggingFace.

    The adapter contains:
    - W: [d_model, d_sae] - maps SAE feature space to model activation space
    - b: [d_sae] - bias term

    Args:
        model_name: HuggingFace model name
        device: Device to load on

    Returns:
        Tuple of (W, b) tensors
    """
    if model_name not in SAE_CONFIGS:
        raise ValueError(f"No effect approximator for model '{model_name}'")

    config = SAE_CONFIGS[model_name]
    adapter_file = config["adapter_file"]

    print(f"   Loading adapter from schalnev/sae-ts-effects / {adapter_file}")
    path = hf_hub_download(
        repo_id="schalnev/sae-ts-effects",
        filename=adapter_file,
        repo_type="dataset",
    )

    adapter = torch.load(path, map_location=device, weights_only=False)

    # Adapter is OrderedDict with 'W' and 'b'
    W = adapter["W"].to(device)  # [d_model, d_sae]
    b = adapter["b"].to(device)  # [d_sae]

    print(f"   Adapter W shape: {W.shape}, b shape: {b.shape}")

    return W, b


def compute_v_diff(
    model,
    tokenizer,
    sae,
    pairs: list[dict],
    layer_idx: int,
    device: str,
) -> torch.Tensor:
    """
    Compute v_diff: the difference vector between positive and negative examples in SAE space.

    v_diff = mean(f(h_l(x+))) - mean(f(h_l(x-)))

    Args:
        model: HuggingFace model
        tokenizer: Tokenizer
        sae: SAE object from sae_lens
        pairs: List of contrastive pairs
        layer_idx: Layer to extract activations from
        device: Device

    Returns:
        v_diff tensor of shape [d_sae]
    """
    pos_features_list = []
    neg_features_list = []

    print(f"   Computing v_diff from {len(pairs)} pairs...")

    for i, pair in enumerate(pairs):
        prompt = pair["prompt"]
        pos_response = pair["positive_response"]["model_response"]
        neg_response = pair["negative_response"]["model_response"]

        pos_text = f"{prompt} {pos_response}"
        neg_text = f"{prompt} {neg_response}"

        # Get activations and encode through SAE
        pos_acts = _get_residual_stream_activations(model, tokenizer, pos_text, layer_idx, device)
        pos_acts = pos_acts.to(device).to(sae.W_enc.dtype)
        # SAE encode: latents = (x - b_dec) @ W_enc + b_enc
        pos_latents = sae.encode(pos_acts)
        # Mean over sequence dimension
        pos_features_list.append(pos_latents.mean(dim=1).detach())  # [1, d_sae]

        neg_acts = _get_residual_stream_activations(model, tokenizer, neg_text, layer_idx, device)
        neg_acts = neg_acts.to(device).to(sae.W_enc.dtype)
        neg_latents = sae.encode(neg_acts)
        neg_features_list.append(neg_latents.mean(dim=1).detach())

        if (i + 1) % 10 == 0:
            print(f"      Processed {i + 1}/{len(pairs)} pairs")

    # Stack and compute mean
    pos_features = torch.cat(pos_features_list, dim=0)  # [num_pairs, d_sae]
    neg_features = torch.cat(neg_features_list, dim=0)

    v_diff = pos_features.mean(dim=0) - neg_features.mean(dim=0)  # [d_sae]

    print(f"   v_diff computed, shape: {v_diff.shape}")
    print(f"   v_diff stats: mean={v_diff.mean():.6f}, std={v_diff.std():.6f}, "
          f"min={v_diff.min():.6f}, max={v_diff.max():.6f}")

    return v_diff


def compute_v_target(
    v_diff: torch.Tensor,
    sparsity: torch.Tensor,
    model_name: str,
    density_threshold: float = 0.01,
    top_k_positive: int = 50,
    top_k_negative: int = 0,
) -> torch.Tensor:
    """
    Compute v_target by filtering v_diff.

    Three filtering stages:
    1. Density filtering: zero out features with activation density > threshold
    2. BOS token filtering: zero out features that activate mainly on BOS token
    3. Top-k selection: keep top positive and negative features

    Args:
        v_diff: Difference vector in SAE space [d_sae]
        sparsity: Feature sparsity/density values from SAE [d_sae]
        model_name: Model name to look up BOS features
        density_threshold: Zero out features with density above this (default 0.01)
        top_k_positive: Number of top positive features to keep
        top_k_negative: Number of top negative features to keep (paper uses 0)

    Returns:
        v_target tensor of shape [d_sae]
    """
    v_filtered = v_diff.clone()

    # Stage 1: Density filtering
    # Zero out features that are too commonly activated (not specific enough)
    if sparsity is not None:
        density_mask = sparsity > density_threshold
        num_filtered = density_mask.sum().item()
        v_filtered[density_mask] = 0
        print(f"   Density filtering: zeroed {num_filtered} features (density > {density_threshold})")

    # Stage 2: BOS filtering
    # Zero out features that activate mainly on BOS tokens
    bos_features = BOS_FEATURES.get(model_name, [])
    if bos_features:
        for idx in bos_features:
            v_filtered[idx] = 0
        print(f"   BOS filtering: zeroed {len(bos_features)} features {bos_features}")
    else:
        print(f"   BOS filtering: no known BOS features for {model_name}")

    # Stage 3: Top-k selection
    v_target = torch.zeros_like(v_filtered)

    # Get top positive features
    if top_k_positive > 0:
        pos_values = v_filtered.clone()
        pos_values[pos_values < 0] = 0
        top_pos_values, top_pos_indices = pos_values.topk(min(top_k_positive, (pos_values > 0).sum().item()))
        v_target[top_pos_indices] = v_filtered[top_pos_indices]
        print(f"   Selected top {len(top_pos_indices)} positive features")

    # Get top negative features (paper uses 0)
    if top_k_negative > 0:
        neg_values = -v_filtered.clone()
        neg_values[neg_values < 0] = 0
        top_neg_values, top_neg_indices = neg_values.topk(min(top_k_negative, (neg_values > 0).sum().item()))
        v_target[top_neg_indices] = v_filtered[top_neg_indices]
        print(f"   Selected top {len(top_neg_indices)} negative features")

    num_nonzero = (v_target != 0).sum().item()
    print(f"   v_target: {num_nonzero} non-zero features")

    return v_target


def compute_v_opt(
    v_target: torch.Tensor,
    W: torch.Tensor,
    b: torch.Tensor,
) -> torch.Tensor:
    """
    Compute v_opt using the effect approximator.

    From paper: v_opt = (W @ v_target_norm) / ||W @ v_target_norm|| - (W @ b) / ||W @ b||

    Args:
        v_target: Target vector in SAE space [d_sae]
        W: Effect approximator weight matrix [d_model, d_sae]
        b: Effect approximator bias [d_sae]

    Returns:
        v_opt tensor of shape [d_model]
    """
    # L1 normalize v_target (as specified in paper)
    v_target_norm = v_target / (v_target.abs().sum() + 1e-8)

    # W is [d_model, d_sae], v_target_norm is [d_sae]
    # W @ v_target_norm -> [d_model]
    Wv = W @ v_target_norm
    Wv_normalized = Wv / (Wv.norm() + 1e-8)

    # Bias term: W @ b -> [d_model]
    Wb = W @ b
    Wb_normalized = Wb / (Wb.norm() + 1e-8)

    # Final v_opt (paper formula)
    v_opt = Wv_normalized - Wb_normalized

    print(f"   v_opt computed, shape: {v_opt.shape}, norm: {v_opt.norm():.6f}")

    return v_opt


def _get_residual_stream_activations(
    model,
    tokenizer,
    text: str,
    layer_idx: int,
    device: str,
) -> torch.Tensor:
    """
    Get residual stream activations from a specific layer.

    Uses output_hidden_states=True (same as wisent's ActivationCollector).

    Args:
        model: HuggingFace model
        tokenizer: Tokenizer
        text: Input text
        layer_idx: Layer index (0-indexed)
        device: Device

    Returns:
        Tensor of shape (1, seq_len, d_model)
    """
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        out = model(**inputs, output_hidden_states=True, use_cache=False)

    # hidden_states is tuple: (embedding, layer0, layer1, ..., layerN)
    # layer_idx=0 -> hs[1], layer_idx=12 -> hs[13]
    hs = out.hidden_states
    return hs[layer_idx + 1]  # +1 because hs[0] is embedding layer


def generate_steering_vector(
    task: str,
    model_name: str,
    output_path: str | Path,
    trait_label: str = "correctness",
    num_pairs: int = 50,
    method: str = "fgaa",
    layers: str | None = None,
    device: str = "cuda:0",
    keep_intermediate: bool = False,
    density_threshold: float = 0.01,
    top_k_positive: int = 50,
    top_k_negative: int = 0,
    extraction_strategy: str = "chat_mean",  # Accepted for API compatibility but not used (FGAA has its own method)
    **kwargs,  # Accept additional kwargs for compatibility
) -> Path:
    """
    Generate a steering vector using the FGAA method.

    Args:
        task: lm-eval task name (e.g., 'boolq', 'cb')
        model_name: HuggingFace model name (must be Gemma 2B or 9B)
        output_path: Where to save the steering vector
        trait_label: Label for the trait being steered
        num_pairs: Number of contrastive pairs to use
        method: Method name (should be 'fgaa')
        layers: Layer(s) to use (e.g., '12' or '10,11,12')
        device: Device to run on
        keep_intermediate: Whether to keep intermediate files
        density_threshold: Density threshold for filtering (default 0.01)
        top_k_positive: Number of top positive features to keep
        top_k_negative: Number of top negative features to keep

    Returns:
        Path to the saved steering vector
    """
    import gc

    output_path = Path(output_path)

    if model_name not in SAE_CONFIGS:
        raise ValueError(
            f"No SAE config for model '{model_name}'. "
            f"Supported models: {list(SAE_CONFIGS.keys())}"
        )

    config = SAE_CONFIGS[model_name]

    # Parse layers
    if layers is None:
        layer_indices = [config["default_layer"]]
    elif layers == "all":
        layer_indices = list(range(config["num_layers"]))
    else:
        layer_indices = [int(l.strip()) for l in layers.split(",")]

    # Step 1: Generate contrastive pairs
    print(f"Step 1: Generating contrastive pairs from task: {task}")
    from wisent.core.cli.generate_pairs_from_task import execute_generate_pairs_from_task

    pairs_file = tempfile.NamedTemporaryFile(mode='w', suffix='_pairs.json', delete=False).name
    pairs_args = Namespace(
        task_name=task,
        limit=num_pairs,
        output=pairs_file,
        seed=42,
        verbose=False,
    )
    execute_generate_pairs_from_task(pairs_args)

    with open(pairs_file) as f:
        pairs_data = json.load(f)
    pairs = pairs_data["pairs"]
    print(f"   Loaded {len(pairs)} contrastive pairs")

    # Step 2: Load model
    print(f"\nStep 2: Loading model {model_name}...")
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map=device,
        torch_dtype=torch.float32,
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.eval()

    # Step 3: Load effect approximator (shared across layers)
    print(f"\nStep 3: Loading effect approximator...")
    W, b = load_effect_approximator(model_name, device=device)

    steering_vectors = {}
    feature_info = {}

    for layer_idx in layer_indices:
        print(f"\nStep 4: Processing layer {layer_idx}")

        # Load SAE for this layer
        sae, sparsity = load_sae(model_name, layer_idx, device=device)

        # Compute v_diff
        print(f"\nStep 5: Computing v_diff for layer {layer_idx}...")
        v_diff = compute_v_diff(model, tokenizer, sae, pairs, layer_idx, device)

        # Compute v_target
        print(f"\nStep 6: Computing v_target for layer {layer_idx}...")
        v_target = compute_v_target(
            v_diff,
            sparsity,
            model_name,
            density_threshold=density_threshold,
            top_k_positive=top_k_positive,
            top_k_negative=top_k_negative,
        )

        # Compute v_opt
        print(f"\nStep 7: Computing v_opt for layer {layer_idx}...")
        v_opt = compute_v_opt(v_target, W, b)

        steering_vectors[str(layer_idx)] = v_opt.cpu().float().tolist()

        # Store feature info
        nonzero_mask = v_target != 0
        nonzero_indices = nonzero_mask.nonzero().squeeze(-1).tolist()
        feature_info[str(layer_idx)] = {
            "num_selected_features": len(nonzero_indices) if isinstance(nonzero_indices, list) else 1,
            "selected_feature_indices": nonzero_indices[:20] if isinstance(nonzero_indices, list) else [nonzero_indices],
            "v_diff_stats": {
                "mean": v_diff.mean().item(),
                "std": v_diff.std().item(),
                "min": v_diff.min().item(),
                "max": v_diff.max().item(),
            },
        }

        # Cleanup SAE
        del sae, sparsity, v_diff, v_target
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # Cleanup
    del model, W, b
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    if not keep_intermediate:
        import os
        os.unlink(pairs_file)

    # Save results
    result = {
        "steering_vectors": steering_vectors,
        "layers": [str(l) for l in layer_indices],
        "model": model_name,
        "method": "fgaa",
        "trait_label": trait_label,
        "task": task,
        "num_pairs": len(pairs),
        "fgaa_params": {
            "density_threshold": density_threshold,
            "top_k_positive": top_k_positive,
            "top_k_negative": top_k_negative,
        },
        "feature_info": feature_info,
    }

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(result, f, indent=2)

    print(f"\nSaved FGAA steering vector to {output_path}")
    return output_path
