"""
Detect BOS (Beginning of Sequence) features in Gemma Scope SAEs.

BOS features are SAE features that activate most strongly at the BOS token position.
These should be filtered out when computing steering vectors as they introduce
artifacts without contributing to steering.

Reference: "Interpretable Steering of Large Language Models with Feature Guided
Activation Additions" (arXiv:2501.09929), Appendix G.

Known BOS features from paper (Gemma-2-2B, layer 12, 16k SAE):
- 11087, 3220, 11752, 12160, 11498

Usage:
    python -m wisent.comparison.detect_bos_features --model google/gemma-2-2b
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from datasets import load_dataset
from tqdm import tqdm


# Known BOS feature indices from paper (Appendix G)
KNOWN_BOS_FEATURES = {
    "google/gemma-2-2b": [11087, 3220, 11752, 12160, 11498],
    "google/gemma-2-9b": [],  # Not listed in paper
}


def load_sample_texts(num_samples: int = 5000, min_length: int = 50) -> list[str]:
    """Load sample texts from WikiText dataset."""
    print(f"Loading up to {num_samples} sample texts from WikiText...")

    dataset = load_dataset("wikitext", "wikitext-103-v1", split="train")

    texts = []
    for item in dataset:
        if len(texts) >= num_samples:
            break
        text = item["text"].strip()
        if len(text) >= min_length:
            texts.append(text)

    print(f"   Loaded {len(texts)} texts")
    return texts


def detect_bos_features(
    model,
    tokenizer,
    sae,
    layer_idx: int,
    device: str,
    texts: list[str],
    threshold: float = 1.5,
    batch_size: int = 8,
) -> tuple[list[int], torch.Tensor]:
    """
    Detect BOS features by finding features that activate most strongly at position 0.

    Uses mean activation approach: computes average activation at BOS position vs
    other positions. A feature is BOS if its mean activation at position 0 is
    significantly higher than at other positions.

    Reference: FGAA paper (arXiv:2501.09929) identifies BOS features as those that
    "exclusively had the strongest activation on the BOS token".

    Args:
        model: HuggingFace model
        tokenizer: Tokenizer
        sae: SAE object from sae_lens
        layer_idx: Layer index (0-indexed)
        device: Device
        texts: List of sample texts to analyze
        threshold: Ratio of (mean BOS activation) / (mean other activation) to qualify
                   as BOS feature. Default 1.5 means BOS activation must be 1.5x higher.
        batch_size: Batch size for processing

    Returns:
        Tuple of (bos_feature_indices, bos_ratio_tensor, mean_bos_activation_tensor)
    """
    d_sae = sae.cfg.d_sae

    # Accumulate sum of activations at BOS position and other positions
    bos_activation_sum = torch.zeros(d_sae, device=device)
    other_activation_sum = torch.zeros(d_sae, device=device)
    bos_count = 0
    other_count = 0

    print(f"Detecting BOS features from {len(texts)} samples...")
    print(f"   Layer: {layer_idx}, d_sae: {d_sae}")

    for i in tqdm(range(0, len(texts), batch_size), desc="Processing"):
        batch_texts = texts[i:i + batch_size]

        inputs = tokenizer(
            batch_texts,
            return_tensors="pt",
            truncation=True,
            max_length=128,
            padding=True,
        )
        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)

        hs = out.hidden_states
        acts = hs[layer_idx + 1].to(sae.W_enc.dtype)
        latents = sae.encode(acts)

        attention_mask = inputs.get("attention_mask")
        for j in range(latents.shape[0]):
            seq_len = int(attention_mask[j].sum().item()) if attention_mask is not None else latents.shape[1]
            sample_latents = latents[j, :seq_len, :]  # [seq_len, d_sae]

            # Accumulate BOS activation (position 0)
            bos_activation_sum += sample_latents[0]
            bos_count += 1

            # Accumulate other positions (1 to seq_len-1)
            if seq_len > 1:
                other_activation_sum += sample_latents[1:].sum(dim=0)
                other_count += seq_len - 1

    # Compute mean activations
    mean_bos = bos_activation_sum / bos_count
    mean_other = other_activation_sum / max(other_count, 1)

    # Compute ratio: how much stronger is BOS activation vs other positions
    # Add small epsilon to avoid division by zero
    bos_ratio = mean_bos / (mean_other + 1e-8)

    # Features where BOS activation is significantly higher than other positions
    bos_features = (bos_ratio >= threshold).nonzero().squeeze(-1).tolist()

    if isinstance(bos_features, int):
        bos_features = [bos_features]

    print(f"\nDetected {len(bos_features)} BOS features (threshold={threshold})")
    print(f"   (Features where mean BOS activation >= {threshold}x mean other activation)")
    return bos_features, bos_ratio, mean_bos


def compare_with_known(model_name: str, detected: list[int], bos_ratios: torch.Tensor, mean_bos: torch.Tensor) -> None:
    """Compare detected BOS features with known list from paper."""
    known = KNOWN_BOS_FEATURES.get(model_name, [])
    detected_set = set(detected)
    known_set = set(known)

    print(f"\n{'='*60}")
    print(f"BOS Feature Comparison for {model_name}")
    print(f"{'='*60}")
    print(f"Known (paper): {sorted(known_set)}")
    print(f"Detected:      {len(detected_set)} features")
    if len(detected_set) <= 20:
        print(f"               {sorted(detected_set)}")
    print(f"{'='*60}")
    print(f"Common:        {sorted(detected_set & known_set)}")
    print(f"Only in paper: {sorted(known_set - detected_set)}")
    if len(detected_set - known_set) <= 20:
        print(f"Only detected: {sorted(detected_set - known_set)}")
    else:
        print(f"Only detected: {len(detected_set - known_set)} additional features")

    if known:
        print(f"\nKnown features stats (mean_bos / mean_other, mean BOS activation):")
        for idx in sorted(known):
            ratio = bos_ratios[idx].item()
            activation = mean_bos[idx].item()
            status = "detected" if idx in detected_set else "missed"
            print(f"   Feature {idx}: ratio={ratio:.2f}x, bos_act={activation:.4f} ({status})")

    print(f"\nTop 10 features by BOS activation ratio:")
    for idx in bos_ratios.topk(10).indices.tolist():
        ratio = bos_ratios[idx].item()
        activation = mean_bos[idx].item()
        marker = " (known)" if idx in known_set else ""
        print(f"   Feature {idx}: ratio={ratio:.2f}x, bos_act={activation:.4f}{marker}")

    print(f"\nTop 10 features by mean BOS activation:")
    for idx in mean_bos.topk(10).indices.tolist():
        ratio = bos_ratios[idx].item()
        activation = mean_bos[idx].item()
        marker = " (known)" if idx in known_set else ""
        detected_marker = " [detected]" if idx in detected_set else ""
        print(f"   Feature {idx}: bos_act={activation:.4f}, ratio={ratio:.2f}x{marker}{detected_marker}")


def main():
    parser = argparse.ArgumentParser(description="Detect BOS features in Gemma Scope SAEs")
    parser.add_argument("--model", default="google/gemma-2-2b", help="Model name")
    parser.add_argument("--layer", type=int, default=12, help="Layer index")
    parser.add_argument("--num-samples", type=int, default=5000, help="Number of text samples")
    parser.add_argument("--threshold", type=float, default=1.5, help="BOS detection threshold: ratio of mean BOS activation to mean other activation")
    parser.add_argument("--batch-size", type=int, default=8, help="Batch size")
    parser.add_argument("--device", default="cuda:0", help="Device")
    parser.add_argument("--output-dir", default="wisent/comparison/comparison_results", help="Output directory")
    args = parser.parse_args()

    print(f"Model: {args.model}")
    print(f"Layer: {args.layer}")
    print(f"Threshold: {args.threshold}")

    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"\nLoading model...")
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        torch_dtype=torch.float32,
        device_map=args.device,
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.eval()

    print(f"\nLoading SAE...")
    from sae_lens import SAE

    release = "gemma-scope-2b-pt-res-canonical" if "2b" in args.model.lower() else "gemma-scope-9b-pt-res-canonical"
    sae_id = f"layer_{args.layer}/width_16k/canonical"
    sae, _, _ = SAE.from_pretrained(release=release, sae_id=sae_id, device=args.device)

    texts = load_sample_texts(args.num_samples)

    bos_features, bos_ratios, mean_bos = detect_bos_features(
        model=model,
        tokenizer=tokenizer,
        sae=sae,
        layer_idx=args.layer,
        device=args.device,
        texts=texts,
        threshold=args.threshold,
        batch_size=args.batch_size,
    )

    compare_with_known(args.model, bos_features, bos_ratios, mean_bos)

    output = {
        "model": args.model,
        "layer": args.layer,
        "threshold": args.threshold,
        "num_samples": len(texts),
        "detected_bos_features": bos_features,
        "known_bos_features": KNOWN_BOS_FEATURES.get(args.model, []),
    }

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"bos_features_{args.model.replace('/', '_')}_layer{args.layer}.json"
    with open(output_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nSaved to {output_path}")


if __name__ == "__main__":
    main()
