app.py

"""
Local SAE Feature Analysis Service

This Flask service runs SAE inference locally using SAELens and TransformerLens,
bypassing Neuronpedia's rate limits for corpus analysis.

Supports:
- GPT-2 Small with res-jb SAE (Layer 6, residual stream)
- Pythia-70M with EleutherAI SAEs (MLP outputs, all layers)
- Gemma 2 2B with Gemma Scope SAEs (Layer 12):
  - Residual stream (gemma-scope-2b-pt-res)
  - Attention output (gemma-scope-2b-pt-att)
  - MLP output (gemma-scope-2b-pt-mlp)
- Batched feature analysis for corpus workflows

Model Selection Rationale:
- GPT-2 Small: The "fruit fly" of interpretability - small, fast, well-documented
- Pythia-70M: EleutherAI's interpretability-focused model with full training checkpoints
- Gemma 2 2B: Google's official Gemma Scope SAEs with residual/attention/MLP variants
Larger models (Llama, Mistral, Pythia 1B+) require GPU and are available via 
Neuronpedia API, not local inference.
"""

import os

import json
import math
import threading
from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import numpy as np


class NumpyEncoder(json.JSONEncoder):
    """JSON encoder that handles numpy types transparently."""
    def default(self, obj):
        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating,)):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)

# Use cached models without network checks when possible
# This prevents hangs when HuggingFace has connectivity issues
os.environ.setdefault("HF_HUB_OFFLINE", "0")  # Set to "1" after first successful load
os.environ.setdefault("TRANSFORMERS_OFFLINE", "0")
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")

# Lazy loading to reduce startup time
_model = None
_saes = {}
_current_model_id = None

import threading
import functools
_model_swap_lock = threading.RLock()

def gpu_exclusive(f):
    """Decorator: hold the model lock for the entire request lifecycle.
    Prevents concurrent requests from swapping the model mid-inference."""
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        with _model_swap_lock:
            return f(*args, **kwargs)
    return wrapper

def _cleanup_model_vram():
    """Free VRAM from current model, ALL cached SAEs, AND OpenReviewer."""
    global _model, _saes, _current_model_id, _openreviewer_model, _openreviewer_tokenizer
    import gc
    if _model is not None:
        try:
            _model.cpu()
        except Exception:
            pass
        del _model
        _model = None
    all_sae_keys = list(_saes.keys())
    for k in all_sae_keys:
        sae = _saes.pop(k, None)
        if sae is not None:
            try:
                if hasattr(sae, 'cpu'):
                    sae.cpu()
                del sae
            except Exception:
                pass
    _current_model_id = None
    if _openreviewer_model is not None:
        print("[SAE Service] Unloading OpenReviewer during VRAM cleanup")
        try:
            _openreviewer_model.cpu()
        except Exception:
            pass
        del _openreviewer_model
        _openreviewer_model = None
    if _openreviewer_tokenizer is not None:
        del _openreviewer_tokenizer
        _openreviewer_tokenizer = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.empty_cache()
        free_mb = torch.cuda.mem_get_info()[0] // (1024 * 1024)
        print(f"[SAE Service] VRAM cleanup done. Free: {free_mb} MB")

GEMMA_SCOPE_L0_MAP = {
    "gemma-2-2b": {
        0: 46, 1: 40, 2: 53, 3: 59, 4: 60, 5: 68,
        6: 70, 7: 69, 8: 71, 9: 73, 10: 77, 11: 80,
        12: 82, 13: 83, 14: 83, 15: 78, 16: 78, 17: 77,
        18: 74, 19: 73, 20: 71, 21: 70, 22: 72, 23: 74,
        24: 73, 25: 55,
    },
    "gemma-2-9b": {
        0: 35, 1: 31, 2: 29, 3: 37, 4: 37, 5: 37,
        6: 47, 7: 46, 8: 51, 9: 51, 10: 57, 11: 60,
        12: 64, 13: 65, 14: 67, 15: 65, 16: 75, 17: 73,
        18: 71, 19: 67, 20: 68, 21: 66, 22: 65, 23: 63,
        24: 61, 25: 61, 26: 63, 27: 65, 28: 65, 29: 66,
        30: 66, 31: 63, 32: 61, 33: 63, 34: 60, 35: 61,
        36: 61, 37: 63, 38: 64, 39: 64, 40: 61, 41: 52,
    },
}

DEFAULT_LAYERS = {
    "gpt2-small": 6,
    "gemma-2-2b": 12,
    "gemma-2-2b-res": 12,
    "gemma-2-2b-att": 12,
    "gemma-2-2b-mlp": 12,
    "gemma-2-9b": 20,
    "pythia-70m": 3,
    "salamandra-2b": 12,
}

MODEL_LAYER_COUNTS = {
    "gpt2-small": 12,
    "gemma-2-2b": 26, "gemma-2-2b-res": 26, "gemma-2-2b-att": 26, "gemma-2-2b-mlp": 26,
    "gemma-2-9b": 42,
    "pythia-70m": 6,
    "salamandra-2b": 24,
}

def clamp_layer(model_id: str, layer: int) -> int:
    max_layers = MODEL_LAYER_COUNTS.get(model_id, 26)
    return max(0, min(layer, max_layers - 1))

def resolve_gemma_sae_id(model_id: str, layer: int) -> str:
    base = "gemma-2-2b" if model_id.startswith("gemma-2-2b") else model_id
    l0_map = GEMMA_SCOPE_L0_MAP.get(base, {})
    default_l0 = 82 if "2b" in model_id else 68
    l0 = l0_map.get(layer, default_l0)
    return f"layer_{layer}/width_16k/average_l0_{l0}"

# Serialize model access — TransformerLens hooks use shared mutable state on the
# global model object.  When Flask handles concurrent requests (threaded=True by
# default), multiple threads calling model.run_with_cache() simultaneously corrupt
# each other's hook activations and return identical results.  This lock ensures
# only one request uses the model at a time.
_model_lock = threading.Lock()

def sanitize_nan(obj):
    """Recursively replace NaN/Inf floats with 0.0 to produce valid JSON."""
    if isinstance(obj, float):
        if math.isnan(obj) or math.isinf(obj):
            return 0.0
        return obj
    elif isinstance(obj, dict):
        return {k: sanitize_nan(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [sanitize_nan(v) for v in obj]
    return obj

app = Flask(__name__)
CORS(app)


def sanitize_for_json(obj):
    """Recursively convert numpy types to native Python types for JSON serialization."""
    if isinstance(obj, dict):
        return {k: sanitize_for_json(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [sanitize_for_json(v) for v in obj]
    elif isinstance(obj, (np.integer,)):
        return int(obj)
    elif isinstance(obj, (np.floating,)):
        v = float(obj)
        if math.isnan(v) or math.isinf(v):
            return 0
        return v
    elif isinstance(obj, np.ndarray):
        return sanitize_for_json(obj.tolist())
    elif isinstance(obj, float):
        if math.isnan(obj) or math.isinf(obj):
            return 0
        return obj
    return obj

@app.errorhandler(TypeError)
def handle_type_error(error):
    """Catch TypeError from jsonify when numpy types slip through — return a clean error."""
    error_msg = str(error)
    if 'JSON serializable' in error_msg:
        print(f"[JSON] TypeError caught: {error_msg} — this usually means numpy types weren't converted")
        return jsonify({"error": "Internal serialization error", "details": error_msg}), 500
    raise error

_GPU_EXEMPT_ROUTES = {'/health', '/cleanup-vram', '/cleanup-tensors'}

@app.before_request
def acquire_gpu_lock():
    """Serialize all GPU-using requests so model swaps can't corrupt running inference."""
    from flask import g, request as req
    if req.path not in _GPU_EXEMPT_ROUTES:
        _model_swap_lock.acquire()
        g._gpu_lock_held = True

@app.teardown_request
def release_gpu_lock(exc=None):
    from flask import g
    if getattr(g, '_gpu_lock_held', False):
        _model_swap_lock.release()
        g._gpu_lock_held = False

@app.after_request
def cleanup_gpu_tensors(response):
    """Free intermediate GPU tensors after every request.
    Route handlers create cache objects from run_with_cache() that hold ALL
    intermediate activations on GPU. Once the handler returns, those caches are
    out of scope but not yet collected by Python's GC. Force collection here
    so CUDA memory is available for the next request."""
    from flask import request as req
    if req.path not in _GPU_EXEMPT_ROUTES:
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return response

@app.after_request
def fix_nan_json(response):
    """Intercept JSON responses and sanitize NaN/Inf values that break JSON parsing."""
    if response.content_type and 'application/json' in response.content_type:
        try:
            data = response.get_json(silent=True)
            if data is not None:
                cleaned = sanitize_for_json(data)
                response.set_data(json.dumps(cleaned, allow_nan=False))
        except Exception:
            pass
    return response

def get_hook_name(model_id: str, layer: int = None) -> str:
    """Get the correct hook name for a given model/SAE type"""
    default_layer = DEFAULT_LAYERS.get(model_id, 12)
    l = layer if layer is not None else default_layer
    if model_id == "gpt2-small":
        return f"blocks.{l}.hook_resid_pre"
    elif model_id == "pythia-70m":
        return f"blocks.{l}.hook_mlp_out"
    elif model_id == "gemma-2-2b-att":
        return f"blocks.{l}.attn.hook_result"
    elif model_id == "gemma-2-2b-mlp":
        return f"blocks.{l}.hook_mlp_out"
    elif model_id in ("gemma-2-2b", "gemma-2-2b-res", "gemma-2-9b"):
        return f"blocks.{l}.hook_resid_pre"
    else:
        return f"blocks.{l}.hook_resid_pre"


def safe_get_cache(cache, hook_name):
    """Safely get activations from cache, trying alternative hook names if needed."""
    if hook_name in cache:
        return cache[hook_name]
    alt = hook_name.replace("hook_resid_pre", "hook_resid_post") if "hook_resid_pre" in hook_name else hook_name.replace("hook_resid_post", "hook_resid_pre")
    if alt in cache:
        print(f"[SAE Service] Hook '{hook_name}' not found, using '{alt}'")
        return cache[alt]
    attn_fallbacks = [
        ("attn.hook_result", "attn.hook_z"),
        ("attn.hook_z", "attn.hook_result"),
        ("attn.hook_result", "attn.hook_attn_out"),
        ("attn.hook_z", "attn.hook_attn_out"),
    ]
    for src, dst in attn_fallbacks:
        if src in hook_name:
            attn_alt = hook_name.replace(src, dst)
            if attn_alt in cache:
                print(f"[SAE Service] Hook '{hook_name}' not found, using attention fallback '{attn_alt}'")
                return cache[attn_alt]
    parts = hook_name.split(".")
    block_prefix = ".".join(parts[:2]) if len(parts) >= 2 else hook_name
    available = [k for k in cache.keys() if block_prefix in k]
    if available:
        print(f"[SAE Service] Hook '{hook_name}' not found, using '{available[0]}' from {len(available)} options")
        return cache[available[0]]
    raise KeyError(f"'{hook_name}' not found in cache. Available: {list(cache.keys())[:10]}")


GATED_MODELS = {"meta-llama/Meta-Llama-3.1-8B"}

ACTIVATION_ONLY_MODELS = {"qwen2.5-1.5b", "bloom-3b", "bloom-560m", "salamandra-2b"}

LLAMA_SAE_LAYERS = 32


def _load_model_only(model_id: str):
    """Load just the HookedTransformer model (no SAE) for activation-only analysis.
    Uses global _model cache, handles VRAM cleanup when switching models."""
    global _model, _saes, _current_model_id
    with _model_swap_lock:
        if _openreviewer_model is not None:
            print("[SAE Service] Unloading OpenReviewer to free VRAM for SAE model")
            _unload_openreviewer()
        base_model = get_base_model_id(model_id)
        if _model is not None and get_base_model_id(_current_model_id) == base_model:
            _model.reset_hooks()
            return _model
        if _model is not None:
            _cleanup_model_vram()
        _model = _load_hooked_transformer(base_model)
        _saes = {}
        _current_model_id = model_id
        return _model

def get_base_model_id(model_id: str) -> str:
    """Get the base transformer model for loading"""
    if model_id is None:
        return None
    if model_id.startswith("gemma-2-2b"):
        return "gemma-2-2b"
    elif model_id == "gemma-2-9b":
        return "gemma-2-9b"
    elif model_id == "pythia-70m":
        return "EleutherAI/pythia-70m"
    elif model_id in ("llama-3.1-8b", "llama-3.1-8B", "Llama-3.1-8B"):
        return "meta-llama/Meta-Llama-3.1-8B"
    elif model_id in ("salamandra-2b", "salamandra-2B"):
        return "BSC-LT/salamandra-2b"
    elif model_id.startswith("gpt2"):
        return "gpt2-small"
    return model_id


def _ensure_qwen_rope_theta():
    """Monkey-patch transformers Qwen2Config to provide rope_theta default if missing."""
    try:
        from transformers import Qwen2Config
        _orig_getattr = Qwen2Config.__getattribute__

        def _patched_getattr(self, key):
            if key == "rope_theta":
                try:
                    return _orig_getattr(self, key)
                except AttributeError:
                    return 1000000.0
            return _orig_getattr(self, key)

        Qwen2Config.__getattribute__ = _patched_getattr
        print("[Model Loader] Patched Qwen2Config to provide rope_theta default")
    except ImportError:
        pass
    except Exception as e:
        print(f"[Model Loader] Qwen patch warning: {e}")

_ensure_qwen_rope_theta()


def _load_hooked_transformer(base_model: str, device: str = None):
    """Load a HookedTransformer model, handling gated models and HF token.
    
    For models not in the TransformerLens registry (e.g. Llama 3.1),
    we use from_pretrained_no_processing which loads directly from HuggingFace
    without requiring a built-in model config.
    """
    from transformer_lens import HookedTransformer
    if device is None:
        device = _get_device()
    kwargs = dict(device=device, dtype=torch.float16)
    if base_model in GATED_MODELS:
        hf_token = os.environ.get("HF_TOKEN")
        if not hf_token:
            raise RuntimeError(f"Model {base_model} is gated and requires HF_TOKEN environment variable")
        kwargs["token"] = hf_token
    print(f"[Model Loader] Loading {base_model} on {device}")
    try:
        return HookedTransformer.from_pretrained(base_model, **kwargs)
    except Exception as e:
        error_msg = str(e)
        if "not found" in error_msg.lower() or "not in" in error_msg.lower() or "valid" in error_msg.lower():
            print(f"[Model Loader] {base_model} not in TransformerLens registry, trying from_pretrained_no_processing...")
            try:
                return HookedTransformer.from_pretrained_no_processing(base_model, **kwargs)
            except Exception as e2:
                print(f"[Model Loader] from_pretrained_no_processing also failed: {e2}")
                raise e2
        raise e

GEMMA_MLP_L0 = 60


def _get_device():
    """Auto-detect best available device: CUDA GPU if available, else CPU"""
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def get_model_and_sae(model_id: str = "gpt2-small", layer: int = None):
    """Lazy load model and SAE on first request, using GPU if available.
    
    SAEs are cached by model_id:layer so switching layers doesn't re-download.
    If layer is None, uses the default layer for the model.
    """
    global _model, _saes, _current_model_id
    
    with _model_swap_lock:
        if _openreviewer_model is not None:
            print("[SAE Service] Unloading OpenReviewer to free VRAM for SAE model")
            _unload_openreviewer()

        if _model is not None:
            _model.reset_hooks()
        
        default_layer = DEFAULT_LAYERS.get(model_id, 12)
        actual_layer = layer if layer is not None else default_layer
        sae_cache_key = f"{model_id}:{actual_layer}"
        
        if _model is not None and _current_model_id == model_id and sae_cache_key in _saes:
            return _model, _saes[sae_cache_key]
        
        device = _get_device()
        print(f"[SAE Service] Loading model: {model_id} layer {actual_layer} on {device}")
        
        from transformer_lens import HookedTransformer
        from sae_lens import SAE
        
        base_model = get_base_model_id(model_id)
        
        current_base = get_base_model_id(_current_model_id) if _current_model_id else None
        if _model is None or current_base != base_model:
            if _model is not None:
                print(f"[SAE Service] Swapping model: {_current_model_id} -> {model_id}, cleaning up VRAM")
                _cleanup_model_vram()
            print(f"[SAE Service] Loading transformer: {base_model} on {device}")
            try:
                _model = _load_hooked_transformer(base_model, device=device)
            except torch.cuda.OutOfMemoryError:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                raise RuntimeError(f"CUDA out of memory loading {model_id}. GPU does not have enough VRAM for this model. Try restarting the Space to clear stale allocations.")
        
        if sae_cache_key not in _saes:
            if model_id == "gpt2-small":
                sae, _, _ = SAE.from_pretrained(
                    release="gpt2-small-res-jb",
                    sae_id=f"blocks.{actual_layer}.hook_resid_pre",
                    device=device
                )
            elif model_id in ("gemma-2-2b", "gemma-2-2b-res"):
                sae_id = resolve_gemma_sae_id(model_id, actual_layer)
                print(f"[SAE Service] Gemma SAE ID: {sae_id}")
                sae, _, _ = SAE.from_pretrained(
                    release="gemma-scope-2b-pt-res",
                    sae_id=sae_id,
                    device=device
                )
            elif model_id == "gemma-2-2b-att":
                sae, _, _ = SAE.from_pretrained(
                    release="gemma-scope-2b-pt-att",
                    sae_id=f"layer_{actual_layer}/width_16k/average_l0_77",
                    device=device
                )
            elif model_id == "gemma-2-2b-mlp":
                sae, _, _ = SAE.from_pretrained(
                    release="gemma-scope-2b-pt-mlp",
                    sae_id=f"layer_{actual_layer}/width_16k/average_l0_{GEMMA_MLP_L0}",
                    device=device
                )
            elif model_id == "gemma-2-9b":
                sae_id = resolve_gemma_sae_id(model_id, actual_layer)
                print(f"[SAE Service] Gemma 9B SAE ID: {sae_id}")
                sae, _, _ = SAE.from_pretrained(
                    release="gemma-scope-9b-pt-res",
                    sae_id=sae_id,
                    device=device
                )
            elif model_id == "pythia-70m":
                sae, _, _ = SAE.from_pretrained(
                    release="EleutherAI/sae-pythia-70m-32k",
                    sae_id=f"blocks.{actual_layer}.hook_mlp_out",
                    device=device
                )
            elif model_id == "llama-3.1-8b":
                llama_layer = min(actual_layer, LLAMA_SAE_LAYERS - 1)
                llama_sae_id = f"blocks.{llama_layer}.hook_resid_post"
                print(f"[SAE Service] Llama Scope SAE: layer {llama_layer}, id {llama_sae_id}")
                try:
                    sae, _, _ = SAE.from_pretrained(
                        release="llama_scope_lxr_32x",
                        sae_id=llama_sae_id,
                        device=device
                    )
                except Exception as llama_sae_err:
                    print(f"[SAE Service] Llama Scope 32x failed ({llama_sae_err}), trying 8x release...")
                    try:
                        sae, _, _ = SAE.from_pretrained(
                            release="llama_scope_lxr_8x",
                            sae_id=llama_sae_id,
                            device=device
                        )
                    except Exception as llama_sae_err2:
                        print(f"[SAE Service] Llama Scope 8x also failed ({llama_sae_err2}), falling back to activation-only")
                        sae = None
            elif model_id in ACTIVATION_ONLY_MODELS:
                sae = None
                print(f"[SAE Service] {model_id} is activation-only (no SAE)")
            else:
                raise ValueError(f"Unsupported model: {model_id}")
            
            _saes[sae_cache_key] = sae
            print(f"[SAE Service] SAE cached as {sae_cache_key}")
        
        _current_model_id = model_id
        print(f"[SAE Service] Model and SAE loaded successfully on {device}")
        return _model, _saes[sae_cache_key]


@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint"""
    free_mb = 0
    model_loaded = _model is not None
    if torch.cuda.is_available():
        free_mb = torch.cuda.mem_get_info()[0] // (1024 * 1024)
    return jsonify({"status": "healthy", "service": "sae-local", "free_mb": free_mb, "model_loaded": model_loaded})


@app.route('/cleanup-vram', methods=['POST'])
def cleanup_vram():
    """Force full VRAM cleanup — unloads model + SAEs. Used before new protocol rounds."""
    global _model, _saes, _current_model_id
    try:
        with _model_swap_lock:
            _cleanup_model_vram()
            free_mb = 0
            if torch.cuda.is_available():
                torch.cuda.synchronize()
                import gc
                gc.collect()
                torch.cuda.empty_cache()
                gc.collect()
                torch.cuda.empty_cache()
                free_mb = torch.cuda.mem_get_info()[0] // (1024 * 1024)
            return jsonify({"status": "ok", "free_mb": free_mb})
    except Exception as e:
        return jsonify({"status": "error", "error": str(e)}), 500

@app.route('/cleanup-tensors', methods=['POST'])
def cleanup_tensors():
    """Lightweight VRAM cleanup — frees intermediate tensors but keeps model loaded.
    Used between consecutive GPU tool calls within a protocol run."""
    try:
        import gc
        gc.collect()
        free_mb = 0
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
            gc.collect()
            torch.cuda.empty_cache()
            free_mb = torch.cuda.mem_get_info()[0] // (1024 * 1024)
        return jsonify({"status": "ok", "free_mb": free_mb, "model_loaded": _model is not None})
    except Exception as e:
        return jsonify({"status": "error", "error": str(e)}), 500


@app.route('/gpu-analyze', methods=['POST'])
def gpu_analyze():
    """
    GPU-compatible analyze: get top SAE features at a specific layer.
    
    Request body:
    {
        "prompt": "The Eiffel Tower is located in",
        "model_id": "gpt2-small",
        "layer": 6,
        "top_k": 32
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model_id', data.get('model', 'gpt2-small'))
        layer = data.get('layer', 0)
        top_k = data.get('top_k', 32)

        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400

        with _model_lock:
            actual_layer = clamp_layer(model_id, layer)
            model, sae = get_model_and_sae(model_id, layer=actual_layer)

            tokens = model.to_tokens(prompt)
            str_tokens = [t.replace('\u0120', ' ').replace('\u010a', '\n') for t in model.to_str_tokens(prompt)]

            hook_name = f"blocks.{actual_layer}.hook_resid_post"

            if sae is None:
                return jsonify({"error": f"SAE not available for {model_id} (activation-only mode). SAE weights could not be loaded."}), 400

            with torch.no_grad():
                _, cache = model.run_with_cache(tokens, names_filter=[hook_name])
                resid = safe_get_cache(cache, hook_name)
                last_pos_resid = resid[0, -1, :]

                sae_input = last_pos_resid.unsqueeze(0)
                feature_acts = sae.encode(sae_input)
                acts = feature_acts[0]

                top_values, top_indices = torch.topk(acts, min(top_k, acts.shape[0]))

                features = []
                for i in range(top_values.shape[0]):
                    val = top_values[i].item()
                    if val > 0:
                        features.append({
                            "feature": top_indices[i].item(),
                            "index": top_indices[i].item(),
                            "activation": round(val, 6)
                        })

                del cache

            target_token = str_tokens[-1] if str_tokens else ""
            device_name = "cuda" if torch.cuda.is_available() else "cpu"

            return jsonify({
                "success": True,
                "model": model_id,
                "layer": actual_layer,
                "tokens": str_tokens,
                "targetToken": target_token,
                "features": features,
                "device": device_name
            })

    except Exception as e:
        print(f"[SAE Service] GPU-analyze error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-activation-map', methods=['POST'])
def feature_activation_map():
    """
    Get per-token activations for a specific SAE feature across all token positions.
    Returns a diverging activation map suitable for Anthropic-style token coloring.

    Request body:
    {
        "prompt": "The model is reasoning about this problem...",
        "model": "gemma-2-2b",
        "layer": 20,
        "feature_index": 1234,
        "normalize": true
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', data.get('model_id', 'gemma-2-2b'))
        layer = data.get('layer', 20)
        feature_index = data.get('feature_index', 0)
        normalize = data.get('normalize', True)

        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400

        with _model_lock:
            actual_layer = clamp_layer(model_id, layer)
            model, sae = get_model_and_sae(model_id, layer=actual_layer)

            if sae is None:
                return jsonify({"error": f"SAE not available for {model_id}. Feature activation map requires SAE support."}), 400

            tokens = model.to_tokens(prompt)
            str_tokens = [t.replace('\u0120', ' ').replace('\u010a', '\n') for t in model.to_str_tokens(prompt)]

            hook_name = f"blocks.{actual_layer}.hook_resid_post"

            with torch.no_grad():
                _, cache = model.run_with_cache(tokens, names_filter=[hook_name])
                resid = safe_get_cache(cache, hook_name)
                seq_len = resid.shape[1]

                all_acts = sae.encode(resid[0])
                feature_acts = all_acts[:, feature_index].float().cpu().tolist()

                top_features_last = []
                last_acts = all_acts[-1]
                top_vals, top_idxs = torch.topk(last_acts, min(20, last_acts.shape[0]))
                for i in range(top_vals.shape[0]):
                    v = top_vals[i].item()
                    if v > 0:
                        top_features_last.append({
                            "feature": top_idxs[i].item(),
                            "activation": round(v, 6),
                        })

                del cache

            max_abs = max(abs(v) for v in feature_acts) if feature_acts else 1.0
            if normalize and max_abs > 0:
                normalized = [round(v / max_abs, 6) for v in feature_acts]
            else:
                normalized = [round(v, 6) for v in feature_acts]

            return jsonify({
                "success": True,
                "model": model_id,
                "layer": actual_layer,
                "featureIndex": feature_index,
                "tokens": str_tokens,
                "activations": [round(v, 6) for v in feature_acts],
                "normalizedActivations": normalized,
                "maxActivation": round(max_abs, 6),
                "topFeaturesAtLastToken": top_features_last,
            })

    except Exception as e:
        print(f"[SAE Service] Feature activation map error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/gpu-analyze-batch', methods=['POST'])
def gpu_analyze_batch():
    """
    Batch GPU analyze: get top SAE features across multiple layers in one request.
    Acquires the model lock once, runs the model once caching all layers, then
    extracts features at each layer. Much faster than N individual /gpu-analyze calls.
    
    Request body:
    {
        "prompts": ["The Eiffel Tower is in", "A calculator is a"],
        "model_id": "gpt2-small",
        "layers": [0, 1, 2, ..., 11],
        "top_k": 20
    }
    
    Returns:
    {
        "success": true,
        "results": [
            {"prompt_index": 0, "layer": 0, "features": [...]},
            {"prompt_index": 0, "layer": 1, "features": [...]},
            {"prompt_index": 1, "layer": 0, "features": [...]},
            ...
        ]
    }
    """
    try:
        data = request.json
        prompts = data.get('prompts', [])
        model_id = data.get('model_id', data.get('model', 'gpt2-small'))
        layers = data.get('layers', list(range(12)))
        top_k = data.get('top_k', 20)

        if not prompts:
            return jsonify({"error": "No prompts provided"}), 400

        import time
        start_time = time.time()

        with _model_lock:
            model, sae = get_model_and_sae(model_id)

            if sae is None:
                return jsonify({"error": f"SAE not available for {model_id} (activation-only mode). SAE weights could not be loaded."}), 400

            actual_layers = [min(l, model.cfg.n_layers - 1) for l in layers]
            unique_layers = sorted(set(actual_layers))
            hook_names = [f"blocks.{l}.hook_resid_post" for l in unique_layers]

            all_results = []

            with torch.no_grad():
                for pi, prompt in enumerate(prompts):
                    tokens = model.to_tokens(prompt)
                    _, cache = model.run_with_cache(tokens, names_filter=hook_names)

                    for layer in unique_layers:
                        hook_name = f"blocks.{layer}.hook_resid_post"
                        resid = safe_get_cache(cache, hook_name)
                        last_pos_resid = resid[0, -1, :]

                        sae_input = last_pos_resid.unsqueeze(0)
                        feature_acts = sae.encode(sae_input)
                        acts = feature_acts[0]

                        top_values, top_indices = torch.topk(acts, min(top_k, acts.shape[0]))

                        features = []
                        for i in range(top_values.shape[0]):
                            val = top_values[i].item()
                            if val > 0:
                                features.append({
                                    "feature": top_indices[i].item(),
                                    "index": top_indices[i].item(),
                                    "activation": round(val, 6)
                                })

                        all_results.append({
                            "prompt_index": pi,
                            "layer": layer,
                            "features": features
                        })

                    del cache

            elapsed = round(time.time() - start_time, 2)
            device_name = "cuda" if torch.cuda.is_available() else "cpu"
            print(f"[SAE Service] Batch analyze: {len(prompts)} prompts x {len(unique_layers)} layers in {elapsed}s on {device_name}")

            return jsonify({
                "success": True,
                "results": all_results,
                "model": model_id,
                "device": device_name,
                "elapsed": elapsed
            })

    except Exception as e:
        print(f"[SAE Service] Batch analyze error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/analyze', methods=['POST'])
def analyze():
    """
    Analyze text with SAE features locally.
    
    Request body:
    {
        "text": "Hello world",
        "model": "gpt2-small",
        "features": [100, 200, 300, ...]  // Feature indices to check
    }
    
    Response:
    {
        "tokens": ["Hello", " world"],
        "activations": {
            "100": [0.0, 0.5, ...],  // Activation per token
            "200": [0.1, 0.0, ...],
            ...
        },
        "model": "gpt2-small"
    }
    """
    try:
        data = request.json
        text = data.get('text', '')
        model_id = data.get('model', 'gpt2-small')
        requested_features = data.get('features', [])
        
        if not text:
            return jsonify({"error": "No text provided"}), 400
        
        with _model_lock:
            model, sae = get_model_and_sae(model_id)
            
            tokens = model.to_tokens(text)
            str_tokens = model.to_str_tokens(text)
            
            with torch.no_grad():
                _, cache = model.run_with_cache(tokens)
                
                hook_name = get_hook_name(model_id)
                
                activations = safe_get_cache(cache, hook_name)
                
                feature_acts = sae.encode(activations)
            
            result_activations = {}
            for feat_idx in requested_features:
                if feat_idx < feature_acts.shape[-1]:
                    feat_values = feature_acts[0, :, feat_idx].tolist()
                    result_activations[str(feat_idx)] = feat_values
            
            clean_tokens = [t.replace('Ġ', ' ').replace('Ċ', '\n') for t in str_tokens]
        
        return jsonify({
            "tokens": clean_tokens,
            "activations": result_activations,
            "model": model_id,
            "tokenCount": len(clean_tokens),
            "featuresAnalyzed": len(requested_features)
        })
        
    except Exception as e:
        print(f"[SAE Service] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/batch-analyze', methods=['POST'])
def batch_analyze():
    """
    Analyze multiple texts in a batch for corpus analysis.
    
    Request body:
    {
        "texts": ["text1", "text2", ...],
        "model": "gpt2-small",
        "features": [100, 200, 300, ...]
    }
    
    Response:
    {
        "results": [
            {"text": "text1", "activatingFeatures": [100, 200]},
            {"text": "text2", "activatingFeatures": [100]},
            ...
        ],
        "featureStats": {
            "100": {"activationCount": 2, "texts": [0, 1]},
            "200": {"activationCount": 1, "texts": [0]},
            ...
        }
    }
    """
    try:
        data = request.json
        texts = data.get('texts', [])
        model_id = data.get('model', 'gpt2-small')
        requested_features = data.get('features', [])
        threshold = data.get('threshold', 0.1)  # Activation threshold
        
        if not texts:
            return jsonify({"error": "No texts provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        
        results = []
        feature_stats = {str(f): {"activationCount": 0, "texts": [], "maxActivation": 0.0} for f in requested_features}
        
        for text_idx, text in enumerate(texts):
            tokens = model.to_tokens(text)
            
            with torch.no_grad():
                _, cache = model.run_with_cache(tokens)
                hook_name = get_hook_name(model_id)
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
            
            # Check which features activated above threshold
            activating_features = []
            for feat_idx in requested_features:
                if feat_idx < feature_acts.shape[-1]:
                    max_act = feature_acts[0, :, feat_idx].max().item()
                    if max_act > threshold:
                        activating_features.append(feat_idx)
                        feat_key = str(feat_idx)
                        feature_stats[feat_key]["activationCount"] += 1
                        feature_stats[feat_key]["texts"].append(text_idx)
                        feature_stats[feat_key]["maxActivation"] = max(
                            feature_stats[feat_key]["maxActivation"], max_act
                        )
            
            results.append({
                "textIndex": text_idx,
                "activatingFeatures": activating_features
            })
            
            print(f"[SAE Service] Processed text {text_idx + 1}/{len(texts)}: {len(activating_features)} features activated")
        
        return jsonify({
            "results": results,
            "featureStats": feature_stats,
            "model": model_id,
            "textsAnalyzed": len(texts),
            "featuresAnalyzed": len(requested_features)
        })
        
    except Exception as e:
        print(f"[SAE Service] Batch error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/steer', methods=['POST'])
def steer():
    """
    Generate text with a feature ablated (set to 0) or amplified.
    Shows how the feature affects model output.
    
    Request body:
    {
        "text": "Hello world",
        "model": "gpt2-small",
        "feature": 16758,
        "mode": "ablate" | "amplify",
        "strength": 5.0  // For amplify mode, multiplier for feature activation
    }
    
    Response:
    {
        "original": {
            "text": "Hello world",
            "continuation": " is a common greeting...",
            "featureActivation": 0.5
        },
        "modified": {
            "text": "Hello world", 
            "continuation": " is used in programming...",
            "featureActivation": 0.0  // or amplified value
        }
    }
    """
    try:
        data = request.json
        text = data.get('text', '')
        model_id = data.get('model', 'gpt2-small')
        feature_idx = data.get('feature', 0)
        mode = data.get('mode', 'ablate')  # "ablate" or "amplify"
        strength = data.get('strength', 5.0)  # Multiplier for amplify
        max_new_tokens = data.get('maxNewTokens', 20)
        
        if not text:
            return jsonify({"error": "No text provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        
        # Get hook name based on model
        hook_name = get_hook_name(model_id)
        
        # Tokenize
        tokens = model.to_tokens(text)
        
        # Generate original continuation
        with torch.no_grad():
            original_output = model.generate(
                tokens, 
                max_new_tokens=max_new_tokens,
                temperature=0.7,
                do_sample=True
            )
            original_text = model.to_string(original_output[0])
            
            # Get original feature activation
            _, cache = model.run_with_cache(tokens)
            activations = safe_get_cache(cache, hook_name)
            feature_acts = sae.encode(activations)
            original_activation = feature_acts[0, :, feature_idx].max().item()
        
        # Validate feature index
        sae_dim = sae.W_dec.shape[0] if hasattr(sae, 'W_dec') else 24576
        if feature_idx < 0 or feature_idx >= sae_dim:
            return jsonify({"error": f"Feature index {feature_idx} out of range (0-{sae_dim-1})"}), 400
        
        # Create steering hook that preserves reconstruction error
        def steering_hook(activations, hook):
            # Encode to get feature activations
            original_feature_acts = sae.encode(activations)
            
            # Compute reconstruction error (what SAE doesn't capture)
            original_reconstruction = sae.decode(original_feature_acts)
            reconstruction_error = activations - original_reconstruction
            
            # Modify the target feature
            modified_feature_acts = original_feature_acts.clone()
            if mode == "ablate":
                # Set target feature to zero
                modified_feature_acts[:, :, feature_idx] = 0.0
            else:  # amplify
                # Multiply target feature activation
                modified_feature_acts[:, :, feature_idx] = modified_feature_acts[:, :, feature_idx] * strength
            
            # Decode modified features and ADD BACK the reconstruction error
            modified_reconstruction = sae.decode(modified_feature_acts)
            return modified_reconstruction + reconstruction_error
        
        # Generate with steering
        with torch.no_grad():
            model.reset_hooks()
            model.add_hook(hook_name, steering_hook)
            
            steered_output = model.generate(
                tokens,
                max_new_tokens=max_new_tokens,
                temperature=0.7,
                do_sample=True
            )
            steered_text = model.to_string(steered_output[0])
            
            # Get modified feature activation (should be near 0 for ablate)
            _, cache = model.run_with_cache(tokens)
            activations = safe_get_cache(cache, hook_name)
            modified_acts = steering_hook(activations, None)
            modified_feature_acts = sae.encode(modified_acts)
            modified_activation = modified_feature_acts[0, :, feature_idx].max().item()
            
            model.reset_hooks()
        
        # Extract just the continuation
        original_continuation = original_text[len(model.to_string(tokens[0])):]
        steered_continuation = steered_text[len(model.to_string(tokens[0])):]
        
        return jsonify({
            "original": {
                "text": text,
                "continuation": original_continuation,
                "featureActivation": original_activation
            },
            "modified": {
                "text": text,
                "continuation": steered_continuation,
                "featureActivation": modified_activation,
                "mode": mode,
                "strength": strength if mode == "amplify" else 0
            },
            "feature": feature_idx,
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Steer error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-ablation', methods=['POST'])
def feature_ablation():
    """
    Measure causal impact of a specific SAE feature by ablating it.
    Zeros out one feature's contribution and measures output change.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompt": "The capital of France is",
        "feature": 16758,
        "layer": 6  (optional, uses model default)
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompt = data.get('prompt', '')
        feature_idx = data.get('feature', 0)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        tokens = model.to_tokens(prompt)
        
        with torch.no_grad():
            baseline_logits = model(tokens)
            baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
            
            _, cache = model.run_with_cache(tokens)
            activations = safe_get_cache(cache, hook_name)
            feature_acts = sae.encode(activations)
            feature_activation = feature_acts[0, :, feature_idx].max().item()
            
            def ablation_hook(acts, hook):
                original_feature_acts = sae.encode(acts)
                original_reconstruction = sae.decode(original_feature_acts)
                reconstruction_error = acts - original_reconstruction
                modified_feature_acts = original_feature_acts.clone()
                modified_feature_acts[:, :, feature_idx] = 0.0
                modified_reconstruction = sae.decode(modified_feature_acts)
                return modified_reconstruction + reconstruction_error
            
            model.reset_hooks()
            ablated_logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[(hook_name, ablation_hook)],
                reset_hooks_end=True
            )
            ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
        
        top_k = 10
        top_baseline_idx = torch.topk(baseline_probs, k=top_k).indices
        top_ablated_idx = torch.topk(ablated_probs, k=top_k).indices
        combined_idx = torch.unique(torch.cat([top_baseline_idx, top_ablated_idx]))
        logit_diff = (baseline_probs[combined_idx] - ablated_probs[combined_idx]).abs().sum().item()
        
        baseline_token = model.to_single_str_token(baseline_probs.argmax().item())
        ablated_token = model.to_single_str_token(ablated_probs.argmax().item())
        
        return jsonify({
            "feature": feature_idx,
            "featureActivation": feature_activation,
            "logitDifference": logit_diff,
            "baselineTopToken": baseline_token.replace('\u0120', ' '),
            "ablatedTopToken": ablated_token.replace('\u0120', ' '),
            "tokenChanged": baseline_token != ablated_token,
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Feature ablation error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-ablation-batch', methods=['POST'])
def feature_ablation_batch():
    """
    Measure causal impact of a specific SAE feature across multiple prompts.
    Returns per-prompt and averaged results.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompts": ["prompt1", "prompt2", ...],
        "feature": 16758
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompts = data.get('prompts', [])
        feature_idx = data.get('feature', 0)
        
        if not prompts:
            return jsonify({"error": "No prompts provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        
        results = []
        for prompt in prompts:
            tokens = model.to_tokens(prompt)
            
            with torch.no_grad():
                baseline_logits = model(tokens)
                baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
                
                _, cache = model.run_with_cache(tokens)
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
                feature_activation = feature_acts[0, :, feature_idx].max().item()
                
                def ablation_hook(acts, hook):
                    original_feature_acts = sae.encode(acts)
                    original_reconstruction = sae.decode(original_feature_acts)
                    reconstruction_error = acts - original_reconstruction
                    modified = original_feature_acts.clone()
                    modified[:, :, feature_idx] = 0.0
                    return sae.decode(modified) + reconstruction_error
                
                model.reset_hooks()
                ablated_logits = model.run_with_hooks(
                    tokens,
                    fwd_hooks=[(hook_name, ablation_hook)],
                    reset_hooks_end=True
                )
                ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
            
            del cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            top_k = 10
            top_b = torch.topk(baseline_probs, k=top_k).indices
            top_a = torch.topk(ablated_probs, k=top_k).indices
            combined = torch.unique(torch.cat([top_b, top_a]))
            logit_diff = (baseline_probs[combined] - ablated_probs[combined]).abs().sum().item()
            
            baseline_token = model.to_single_str_token(baseline_probs.argmax().item())
            ablated_token = model.to_single_str_token(ablated_probs.argmax().item())
            
            results.append({
                "prompt": prompt,
                "featureActivation": feature_activation,
                "logitDifference": logit_diff,
                "baselineTopToken": baseline_token.replace('\u0120', ' '),
                "ablatedTopToken": ablated_token.replace('\u0120', ' '),
                "tokenChanged": baseline_token != ablated_token
            })
        
        avg_logit_diff = sum(r["logitDifference"] for r in results) / len(results)
        avg_activation = sum(r["featureActivation"] for r in results) / len(results)
        tokens_changed = sum(1 for r in results if r["tokenChanged"])
        
        return jsonify({
            "feature": feature_idx,
            "results": results,
            "averageLogitDiff": avg_logit_diff,
            "averageActivation": avg_activation,
            "tokensChanged": tokens_changed,
            "totalPrompts": len(prompts),
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Feature ablation batch error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-ablation-sweep', methods=['POST'])
def feature_ablation_sweep():
    """
    Ablate MULTIPLE features across ALL prompts in one call.
    Returns per-feature causal impact with KL divergence.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompts": ["prompt1", "prompt2", ...],
        "features": [6471, 7393, 10543, ...]
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompts = data.get('prompts', [])
        feature_indices = data.get('features', [])
        
        if not prompts:
            return jsonify({"error": "No prompts provided"}), 400
        if not feature_indices:
            return jsonify({"error": "No features provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        
        sweep_results = {}
        
        for fi_idx, feature_idx in enumerate(feature_indices):
            print(f"[SAE Service] Ablation sweep: feature {feature_idx} ({fi_idx+1}/{len(feature_indices)})")
            
            per_prompt = []
            for prompt in prompts:
                tokens = model.to_tokens(prompt)
                
                with torch.no_grad():
                    baseline_logits = model(tokens)
                    baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
                    
                    _, cache = model.run_with_cache(tokens)
                    activations = safe_get_cache(cache, hook_name)
                    feature_acts = sae.encode(activations)
                    feature_activation = feature_acts[0, :, feature_idx].max().item()
                    
                    current_feature = feature_idx
                    def ablation_hook(acts, hook, feat_idx=current_feature):
                        orig_acts = sae.encode(acts)
                        orig_recon = sae.decode(orig_acts)
                        recon_error = acts - orig_recon
                        modified = orig_acts.clone()
                        modified[:, :, feat_idx] = 0.0
                        return sae.decode(modified) + recon_error
                    
                    model.reset_hooks()
                    ablated_logits = model.run_with_hooks(
                        tokens,
                        fwd_hooks=[(hook_name, ablation_hook)],
                        reset_hooks_end=True
                    )
                    ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
                
                del cache
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

                eps = 1e-10
                kl_div = (baseline_probs * torch.log((baseline_probs + eps) / (ablated_probs + eps))).sum().item()
                
                top_k = 10
                top_b = torch.topk(baseline_probs, k=top_k).indices
                top_a = torch.topk(ablated_probs, k=top_k).indices
                combined = torch.unique(torch.cat([top_b, top_a]))
                logit_diff = (baseline_probs[combined] - ablated_probs[combined]).abs().sum().item()
                
                baseline_token = model.to_single_str_token(baseline_probs.argmax().item())
                ablated_token = model.to_single_str_token(ablated_probs.argmax().item())
                
                per_prompt.append({
                    "prompt": prompt,
                    "featureActivation": round(feature_activation, 6),
                    "logitDifference": round(logit_diff, 6),
                    "klDivergence": round(kl_div, 6),
                    "baselineTopToken": baseline_token.replace('\u0120', ' '),
                    "ablatedTopToken": ablated_token.replace('\u0120', ' '),
                    "tokenChanged": baseline_token != ablated_token
                })
            
            avg_logit_diff = sum(r["logitDifference"] for r in per_prompt) / len(per_prompt)
            avg_kl = sum(r["klDivergence"] for r in per_prompt) / len(per_prompt)
            avg_activation = sum(r["featureActivation"] for r in per_prompt) / len(per_prompt)
            tokens_changed = sum(1 for r in per_prompt if r["tokenChanged"])
            
            sweep_results[str(feature_idx)] = {
                "feature": feature_idx,
                "averageLogitDiff": round(avg_logit_diff, 6),
                "averageKLDivergence": round(avg_kl, 6),
                "averageActivation": round(avg_activation, 6),
                "tokensChanged": tokens_changed,
                "totalPrompts": len(prompts),
                "perPrompt": per_prompt
            }
        
        return jsonify({
            "success": True,
            "results": sweep_results,
            "featuresTotal": len(feature_indices),
            "promptsTotal": len(prompts),
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Feature ablation sweep error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/patch', methods=['POST'])
def activation_patch():
    """Single-layer activation patching between clean and corrupted prompts."""
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        clean_prompt = data.get('cleanPrompt', '')
        corrupted_prompt = data.get('corruptedPrompt', '')
        patch_layer = data.get('patchLayer', 6)
        patch_component = data.get('patchComponent', 'residual')
        patch_direction = data.get('patchDirection', 'noising')

        if not clean_prompt or not corrupted_prompt:
            return jsonify({"error": "Both clean and corrupted prompts are required"}), 400

        with _model_lock:
            if model_id in ACTIVATION_ONLY_MODELS:
                model = _load_model_only(model_id)
            else:
                model, sae = get_model_and_sae(model_id)

            clean_tokens = model.to_tokens(clean_prompt)
            corrupted_tokens = model.to_tokens(corrupted_prompt)

            n_layers = model.cfg.n_layers
            if patch_layer < 0 or patch_layer >= n_layers:
                return jsonify({"error": f"Layer {patch_layer} out of range (0-{n_layers-1})"}), 400

            if patch_component == "attention":
                hook_name = f"blocks.{patch_layer}.attn.hook_result"
            elif patch_component == "mlp":
                hook_name = f"blocks.{patch_layer}.hook_mlp_out"
            else:
                hook_name = f"blocks.{patch_layer}.hook_resid_post"

            is_attn = patch_component == "attention"

            with torch.no_grad():
                _, clean_cache = model.run_with_cache(clean_tokens)
                _, corrupted_cache = model.run_with_cache(corrupted_tokens)
                clean_activations = safe_get_cache(clean_cache, hook_name)
                corrupted_activations = safe_get_cache(corrupted_cache, hook_name)

                if patch_direction == "noising":
                    run_tokens = clean_tokens
                    patch_in_activations = corrupted_activations
                else:
                    run_tokens = corrupted_tokens
                    patch_in_activations = clean_activations

                baseline_logits = model(run_tokens)
                baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
                baseline_top_tokens = torch.topk(baseline_probs, k=10)

                def patch_hook(activations, hook):
                    new_activations = activations.clone()
                    last_pos = min(activations.shape[1], patch_in_activations.shape[1]) - 1
                    if is_attn and activations.dim() == 4 and patch_in_activations.dim() == 4:
                        new_activations[:, last_pos, :, :] = patch_in_activations[:, last_pos, :, :].clone()
                    else:
                        new_activations[:, last_pos, :] = patch_in_activations[:, last_pos, :].clone()
                    return new_activations

                model.reset_hooks()
                patched_logits = model.run_with_hooks(
                    run_tokens,
                    fwd_hooks=[(hook_name, patch_hook)],
                    reset_hooks_end=True
                )

                patched_probs = torch.softmax(patched_logits[0, -1, :], dim=-1)
                patched_top_tokens = torch.topk(patched_probs, k=10)

                del clean_cache, corrupted_cache
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

            with torch.no_grad():
                baseline_gen = model.generate(run_tokens, max_new_tokens=20, temperature=0.7)
                baseline_output = model.to_string(baseline_gen[0])

                def gen_patch_hook(activations, hook):
                    min_len = min(activations.shape[1], patch_in_activations.shape[1])
                    if is_attn and activations.dim() == 4 and patch_in_activations.dim() == 4:
                        activations[:, :min_len, :, :] = patch_in_activations[:, :min_len, :, :]
                    else:
                        activations[:, :min_len, :] = patch_in_activations[:, :min_len, :]
                    return activations

                model.add_hook(hook_name, gen_patch_hook)
                try:
                    patched_gen = model.generate(run_tokens, max_new_tokens=20, temperature=0.7)
                    patched_output = model.to_string(patched_gen[0])
                except Exception as gen_err:
                    print(f"[SAE Service] Patched generation failed (non-fatal): {gen_err}")
                    patched_output = "(generation failed)"
                finally:
                    model.reset_hooks()

            top_k = 10
            top_baseline_idx = torch.topk(baseline_probs, k=top_k).indices
            top_patch_idx = torch.topk(patched_probs, k=top_k).indices
            combined_idx = torch.unique(torch.cat([top_baseline_idx, top_patch_idx]))
            logit_diff = (baseline_probs[combined_idx] - patched_probs[combined_idx]).abs().sum().item()

            token_changes = []
            for i in range(min(10, len(baseline_top_tokens.indices))):
                token_idx = baseline_top_tokens.indices[i].item()
                token_str = model.to_single_str_token(token_idx)
                baseline_prob = baseline_probs[token_idx].item()
                patch_prob = patched_probs[token_idx].item()
                token_changes.append({
                    "token": token_str.replace('\u0120', ' ').replace('\u010a', '\n'),
                    "baselineProb": round(baseline_prob, 4),
                    "patchedProb": round(patch_prob, 4),
                    "change": round(patch_prob - baseline_prob, 4)
                })

            return jsonify({
                "success": True,
                "baselineOutput": baseline_output,
                "patchedOutput": patched_output,
                "logitDifference": round(logit_diff, 4),
                "tokenChanges": token_changes,
                "patchLayer": patch_layer,
                "patchComponent": patch_component,
                "patchDirection": patch_direction
            })

    except Exception as e:
        if _model is not None:
            _model.reset_hooks()
        print(f"[SAE Service] Patch error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/sweep', methods=['POST'])
def layer_sweep():
    """
    Sweep across all layers to find which layer has strongest patching effect.
    
    Request body:
    {
        "model": "gpt2-small",
        "cleanPrompt": "I am an AI assistant",
        "corruptedPrompt": "A calculator is a device",
        "patchComponent": "residual",
        "patchDirection": "noising",
        "numLayers": 12
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        clean_prompt = data.get('cleanPrompt', '')
        corrupted_prompt = data.get('corruptedPrompt', '')
        patch_component = data.get('patchComponent', 'residual')
        patch_direction = data.get('patchDirection', 'noising')
        num_layers = data.get('numLayers', 12)
        skip_generate = data.get('skipGenerate', False)
        
        if not clean_prompt or not corrupted_prompt:
            return jsonify({"error": "Both clean and corrupted prompts are required"}), 400

        with _model_lock:
            if model_id in ACTIVATION_ONLY_MODELS:
                model = _load_model_only(model_id)
            else:
                model, sae = get_model_and_sae(model_id)
            actual_layers = min(num_layers, model.cfg.n_layers)

            clean_tokens = model.to_tokens(clean_prompt)
            corrupted_tokens = model.to_tokens(corrupted_prompt)

            primary_hook_names = []
            for layer in range(actual_layers):
                if patch_component == "attention":
                    primary_hook_names.append(f"blocks.{layer}.attn.hook_result")
                elif patch_component == "mlp":
                    primary_hook_names.append(f"blocks.{layer}.hook_mlp_out")
                else:
                    primary_hook_names.append(f"blocks.{layer}.hook_resid_post")

            results = []

            with torch.no_grad():
                clean_logits, clean_cache = model.run_with_cache(clean_tokens)
                corrupted_logits, corrupted_cache = model.run_with_cache(corrupted_tokens)

                if patch_direction == "noising":
                    run_tokens = clean_tokens
                    patch_cache = corrupted_cache
                    baseline_probs = torch.softmax(clean_logits[0, -1, :], dim=-1)
                else:
                    run_tokens = corrupted_tokens
                    patch_cache = clean_cache
                    baseline_probs = torch.softmax(corrupted_logits[0, -1, :], dim=-1)

                baseline_output = None
                if not skip_generate:
                    baseline_gen = model.generate(run_tokens, max_new_tokens=15, temperature=0.7)
                    baseline_output = model.to_string(baseline_gen[0])

                for layer in range(actual_layers):
                    hook_name = primary_hook_names[layer]

                    try:
                        patch_in_acts = safe_get_cache(patch_cache, hook_name)
                    except KeyError as ke:
                        print(f"[SAE Service] Skipping layer {layer}: {ke}")
                        continue

                    resolved_hook = hook_name if hook_name in patch_cache else None
                    if resolved_hook is None:
                        for candidate in [hook_name, hook_name.replace("hook_result", "hook_z"), hook_name.replace("hook_result", "hook_attn_out")]:
                            if candidate in patch_cache:
                                resolved_hook = candidate
                                break
                        if resolved_hook is None:
                            resolved_hook = hook_name

                    def make_patch_hook(patch_activations, layer_num, is_attn=False):
                        def patch_hook(activations, hook):
                            new_activations = activations.clone()
                            last_pos = min(activations.shape[1], patch_activations.shape[1]) - 1
                            if is_attn and activations.dim() == 4 and patch_activations.dim() == 4:
                                new_activations[:, last_pos, :, :] = patch_activations[:, last_pos, :, :].clone()
                            else:
                                new_activations[:, last_pos, :] = patch_activations[:, last_pos, :].clone()
                            return new_activations
                        return patch_hook

                    is_attn = patch_component == "attention"

                    model.reset_hooks()
                    try:
                        patched_logits = model.run_with_hooks(
                            run_tokens,
                            fwd_hooks=[(resolved_hook, make_patch_hook(patch_in_acts.clone(), layer, is_attn))],
                            reset_hooks_end=True
                        )
                    except Exception as hook_err:
                        print(f"[SAE Service] Layer {layer} patching failed: {hook_err}")
                        model.reset_hooks()
                        continue

                    patched_probs = torch.softmax(patched_logits[0, -1, :], dim=-1)

                    top_k = 10
                    top_baseline_indices = torch.topk(baseline_probs, k=top_k).indices
                    top_patch_indices = torch.topk(patched_probs, k=top_k).indices
                    combined_indices = torch.unique(torch.cat([top_baseline_indices, top_patch_indices]))
                    logit_diff = (baseline_probs[combined_indices] - patched_probs[combined_indices]).abs().sum().item()

                    token_details = []
                    top_indices = torch.topk(baseline_probs, k=5).indices
                    for idx in top_indices:
                        token_str = model.to_single_str_token(idx.item())
                        baseline_p = baseline_probs[idx].item()
                        patch_p = patched_probs[idx].item()
                        token_details.append({
                            "token": token_str.replace('\u0120', ' ').replace('\u010a', '\n'),
                            "position": 0,
                            "baselineProb": round(baseline_p, 4),
                            "patchedProb": round(patch_p, 4),
                            "change": round(patch_p - baseline_p, 4)
                        })

                    patched_output = None
                    if not skip_generate:
                        try:
                            model.add_hook(resolved_hook, make_patch_hook(patch_in_acts, layer, is_attn))
                            patched_gen = model.generate(run_tokens, max_new_tokens=15, temperature=0.7)
                            patched_output = model.to_string(patched_gen[0])
                        except Exception as gen_err:
                            print(f"[SAE Service] Layer {layer} generate failed (non-fatal): {gen_err}")
                            patched_output = "(generation failed)"
                        finally:
                            model.reset_hooks()
                    else:
                        model.reset_hooks()

                    results.append({
                        "layer": layer,
                        "logitDifference": round(logit_diff, 4),
                        "component": patch_component,
                        "tokenDetails": token_details,
                        "baselineOutput": baseline_output,
                        "patchedOutput": patched_output
                    })

                del clean_cache, corrupted_cache, patch_cache

            return jsonify({
                "success": True,
                "results": results,
                "patchDirection": patch_direction,
                "model": model_id
            })

    except Exception as e:
        if _model is not None:
            _model.reset_hooks()
        print(f"[SAE Service] Sweep error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-layer-activations', methods=['POST'])
def feature_layer_activations():
    """
    Get per-feature activation projections across all layers.
    
    Projects residual stream activations at each layer onto each feature's
    SAE decoder direction. This shows where in the network each feature's
    representation emerges and strengthens.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompt": "What are your main limitations as an AI model?",
        "featureIndices": [6471, 7393, 10543],
        "promptGroup": "A"  // optional, for labeling
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompt = data.get('prompt', '')
        feature_indices = data.get('featureIndices', [])
        
        if not prompt:
            return jsonify({"error": "Prompt is required"}), 400
        if not feature_indices:
            return jsonify({"error": "Feature indices are required"}), 400
        
        model, sae = get_model_and_sae(model_id)
        
        num_layers = model.cfg.n_layers
        tokens = model.to_tokens(prompt)
        
        with torch.no_grad():
            _, cache = model.run_with_cache(tokens)
            
            decoder_weights = sae.W_dec.detach()
            
            results = {}
            for feat_idx in feature_indices:
                if feat_idx >= decoder_weights.shape[0]:
                    continue
                    
                feat_direction = decoder_weights[feat_idx]
                feat_direction_norm = feat_direction / (feat_direction.norm() + 1e-8)
                
                layer_activations = {}
                for l in range(num_layers):
                    hook_name = f"blocks.{l}.hook_resid_post"
                    if hook_name not in cache:
                        hook_name = f"blocks.{l}.hook_resid_pre"
                    if hook_name not in cache:
                        layer_activations[l] = 0.0
                        continue
                    
                    resid = safe_get_cache(cache, hook_name)
                    last_token_resid = resid[0, -1, :]
                    
                    projection = torch.dot(last_token_resid, feat_direction_norm).item()
                    layer_activations[l] = round(max(0, projection), 4)
                
                results[str(feat_idx)] = layer_activations
        
        return jsonify({
            "success": True,
            "results": results,
            "numLayers": num_layers,
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Feature layer activations error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/feature-layer-activations-batch', methods=['POST'])
def feature_layer_activations_batch():
    """
    Batch version: average feature projections across multiple prompts.
    This gives more robust per-feature, per-layer activation profiles.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompts": ["prompt1", "prompt2", ...],
        "featureIndices": [6471, 7393, 10543]
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompts = data.get('prompts', [])
        feature_indices = data.get('featureIndices', [])
        
        if not prompts:
            return jsonify({"error": "Prompts are required"}), 400
        if not feature_indices:
            return jsonify({"error": "Feature indices are required"}), 400
        
        model, sae = get_model_and_sae(model_id)
        
        num_layers = model.cfg.n_layers
        decoder_weights = sae.W_dec.detach()
        
        accumulated = {str(fi): {l: 0.0 for l in range(num_layers)} for fi in feature_indices if fi < decoder_weights.shape[0]}
        valid_features = [fi for fi in feature_indices if fi < decoder_weights.shape[0]]
        
        feat_directions = {}
        for fi in valid_features:
            d = decoder_weights[fi]
            feat_directions[fi] = d / (d.norm() + 1e-8)
        
        with torch.no_grad():
            for prompt in prompts:
                tokens = model.to_tokens(prompt)
                _, cache = model.run_with_cache(tokens)
                
                for fi in valid_features:
                    for l in range(num_layers):
                        hook_name = f"blocks.{l}.hook_resid_post"
                        if hook_name not in cache:
                            hook_name = f"blocks.{l}.hook_resid_pre"
                        if hook_name not in cache:
                            continue
                        
                        resid = safe_get_cache(cache, hook_name)
                        last_token_resid = resid[0, -1, :]
                        projection = torch.dot(last_token_resid, feat_directions[fi]).item()
                        accumulated[str(fi)][l] += max(0, projection)
        
        num_prompts = len(prompts)
        results = {}
        for fi_str, layers in accumulated.items():
            results[fi_str] = {l: round(v / num_prompts, 4) for l, v in layers.items()}
        
        return jsonify({
            "success": True,
            "results": results,
            "numLayers": num_layers,
            "numPrompts": num_prompts,
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Batch feature layer activations error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Chain of Thought Monitoring Endpoints (Fidelity)
# ============================================

def parse_cot_steps(text: str) -> list:
    """Parse numbered reasoning steps from text"""
    import re
    steps = []
    # Match patterns like "1.", "Step 1:", "1)", etc.
    pattern = r'(?:^|\n)\s*(?:Step\s*)?(\d+)[.:)\]]\s*(.+?)(?=(?:\n\s*(?:Step\s*)?\d+[.:)\]]|\Z))'
    matches = re.findall(pattern, text, re.DOTALL | re.IGNORECASE)
    for num, content in matches:
        steps.append({
            "stepNumber": int(num),
            "content": content.strip()
        })
    # If no numbered steps found, treat each sentence as a step
    if not steps:
        sentences = [s.strip() for s in text.split('.') if s.strip()]
        for i, s in enumerate(sentences[:10]):  # Limit to 10
            steps.append({"stepNumber": i + 1, "content": s})
    return steps


@app.route('/cot/analyze', methods=['POST'])
def cot_analyze():
    """
    Analyze chain of thought reasoning for faithfulness.
    
    Request body:
    {
        "prompt": "Solve this step by step: What is 15 + 27?",
        "model": "gpt2-small",
        "layers": [4, 6, 8, 10]  // Which layers to analyze
    }
    
    Response:
    {
        "generatedText": "Step 1: Add 15 and 27...",
        "steps": [
            {
                "stepNumber": 1,
                "content": "Add 15 and 27",
                "tokenSpan": [5, 10],
                "faithfulness": {
                    "activationAlignment": 0.85,
                    "causalImpact": 0.72,
                    "score": 0.78,
                    "verdict": "faithful"
                },
                "layerActivations": {...}
            }
        ],
        "overallFaithfulness": 0.75
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gpt2-small')
        layers = data.get('layers', [4, 6, 8, 10])
        max_tokens = min(data.get('maxTokens', 40), 50)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        n_layers = model.cfg.n_layers
        layers = [l for l in layers if l < n_layers]
        if not layers:
            layers = [n_layers // 3, n_layers // 2, n_layers * 2 // 3, n_layers - 1]
            layers = [l for l in layers if l < n_layers]
        
        hook_names = [f"blocks.{l}.hook_resid_post" for l in layers]
        names_filter = lambda name: name in hook_names
        
        tokens = model.to_tokens(prompt)
        prompt_len = tokens.shape[1]
        
        with torch.no_grad():
            output = model.generate(tokens, max_new_tokens=max_tokens, temperature=0.7)
            generated_text = model.to_string(output[0])
            
            _, cache = model.run_with_cache(output, names_filter=names_filter)
            
            layer_activations = {}
            for layer in layers:
                hook_name = f"blocks.{layer}.hook_resid_post"
                acts = safe_get_cache(cache, hook_name)[0, prompt_len:, :]
                layer_activations[layer] = {
                    "mean": acts.mean().item(),
                    "norm": acts.norm(dim=-1).mean().item(),
                    "shape": list(acts.shape)
                }
            
            final_emb = cache[f"blocks.{layers[-1]}.hook_resid_post"][0, -1, :]
        
        response_text = generated_text[len(model.to_string(tokens[0])):]
        steps = parse_cot_steps(response_text)
        
        analyzed_steps = []
        analysis_layer = layers[-1]
        analysis_hook = f"blocks.{analysis_layer}.hook_resid_post"
        analysis_filter = lambda name: name == analysis_hook
        
        with torch.no_grad():
            baseline_logits = model(output)
            baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
        
        full_text = model.to_string(output[0])
        
        for step in steps:
            step_text = step["content"]
            
            with torch.no_grad():
                step_tokens = model.to_tokens(step_text)
                _, step_cache = model.run_with_cache(step_tokens, names_filter=analysis_filter)
                
                step_emb = step_cache[analysis_hook][0, -1, :]
                
                cos_sim = torch.nn.functional.cosine_similarity(
                    step_emb.unsqueeze(0), final_emb.unsqueeze(0)
                ).item()
            
            causal_impact = 0.0
            try:
                step_start = full_text.find(step_text)
                if step_start >= 0:
                    ablated_text = full_text[:step_start] + full_text[step_start + len(step_text):]
                    ablated_tokens = model.to_tokens(ablated_text)
                    with torch.no_grad():
                        ablated_logits = model(ablated_tokens)
                        ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
                        top_k = 10
                        top_baseline = torch.topk(baseline_probs, k=top_k).indices
                        top_ablated = torch.topk(ablated_probs, k=top_k).indices
                        combined = torch.unique(torch.cat([top_baseline, top_ablated]))
                        causal_impact = (baseline_probs[combined] - ablated_probs[combined]).abs().sum().item()
            except Exception as ablation_err:
                print(f"[CoT Service] Ablation failed for step {step['stepNumber']}: {ablation_err}")
                causal_impact = 0.0
            
            faithfulness_score = (abs(cos_sim) * 0.4 + min(causal_impact, 1.0) * 0.6)
            
            verdict = "faithful" if faithfulness_score > 0.5 else "suspicious" if faithfulness_score > 0.25 else "unfaithful"
            
            analyzed_steps.append({
                "stepNumber": step["stepNumber"],
                "content": step["content"],
                "faithfulness": {
                    "activationAlignment": round(abs(cos_sim), 3),
                    "causalImpact": round(causal_impact, 4),
                    "score": round(faithfulness_score, 3),
                    "verdict": verdict
                }
            })
        
        del cache
        
        if analyzed_steps:
            overall = sum(s["faithfulness"]["score"] for s in analyzed_steps) / len(analyzed_steps)
        else:
            overall = 0.0
        
        return jsonify({
            "success": True,
            "prompt": prompt,
            "generatedText": response_text,
            "steps": analyzed_steps,
            "layerActivations": layer_activations,
            "overallFaithfulness": round(overall, 3),
            "model": model_id
        })
        
    except Exception as e:
        print(f"[CoT Service] Analyze error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/cot/ablate', methods=['POST'])
def cot_ablate():
    """
    Ablate a specific reasoning step to measure its causal impact.
    
    Request body:
    {
        "prompt": "Solve step by step: 15 + 27",
        "generatedText": "Step 1: ... Step 2: ...",
        "stepToAblate": 1,
        "model": "gpt2-small"
    }
    
    Response:
    {
        "baselineOutput": "42",
        "ablatedOutput": "38",
        "logitDifference": 0.35,
        "stepCausalImpact": 0.72
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        step_num = data.get('stepToAblate', 1)
        model_id = data.get('model', 'gpt2-small')
        layer = data.get('layer', 6)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        model, sae = get_model_and_sae(model_id, layer=layer)
        
        # Generate baseline
        tokens = model.to_tokens(prompt)
        
        with torch.no_grad():
            # Baseline generation
            baseline_output = model.generate(tokens, max_new_tokens=50, temperature=0.7)
            baseline_text = model.to_string(baseline_output[0])
            
            # Get baseline activations
            _, cache = model.run_with_cache(baseline_output)
            baseline_acts = cache[f"blocks.{layer}.hook_resid_post"].clone()
            
            # Create ablation hook that zeros out middle portion of activations
            # (representing the "step" being ablated)
            prompt_len = tokens.shape[1]
            total_len = baseline_output.shape[1]
            step_start = prompt_len + (step_num - 1) * 5  # Approximate
            step_end = min(step_start + 10, total_len)
            
            def ablation_hook(activations, hook):
                new_acts = activations.clone()
                if step_start < new_acts.shape[1]:
                    end = min(step_end, new_acts.shape[1])
                    new_acts[:, step_start:end, :] = 0.0
                return new_acts
            
            # Generate with ablation
            model.reset_hooks()
            model.add_hook(f"blocks.{layer}.hook_resid_post", ablation_hook)
            ablated_output = model.generate(tokens, max_new_tokens=50, temperature=0.7)
            ablated_text = model.to_string(ablated_output[0])
            model.reset_hooks()
            
            # Compute logit difference
            baseline_logits = model(baseline_output)
            
            model.add_hook(f"blocks.{layer}.hook_resid_post", ablation_hook)
            ablated_logits = model(tokens)
            model.reset_hooks()
            
            baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
            ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
            
            top_k = 10
            top_idx = torch.topk(baseline_probs, k=top_k).indices
            logit_diff = (baseline_probs[top_idx] - ablated_probs[top_idx]).abs().sum().item()
        
        # Causal impact: how much did ablating change the output?
        causal_impact = min(1.0, logit_diff / 2.0)  # Normalize
        
        return jsonify({
            "success": True,
            "baselineOutput": baseline_text[len(model.to_string(tokens[0])):],
            "ablatedOutput": ablated_text[len(model.to_string(tokens[0])):],
            "logitDifference": round(logit_diff, 4),
            "stepCausalImpact": round(causal_impact, 3),
            "stepAblated": step_num,
            "layer": layer
        })
        
    except Exception as e:
        print(f"[CoT Service] Ablate error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Attention Analysis Endpoints
# ============================================

@app.route('/attention/analyze', methods=['POST'])
def attention_analyze():
    """
    Analyze attention patterns across all heads.
    
    Request body:
    {
        "prompt": "The capital of France is Paris",
        "model": "gpt2-small",
        "topK": 10
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gpt2-small')
        top_k = data.get('topK', 10)
        
        if not prompt:
            return jsonify({"error": "Prompt required"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[Attention] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        tokens = _model.to_tokens(prompt)
        token_strs = _model.to_str_tokens(prompt)
        
        with torch.no_grad():
            _, cache = _model.run_with_cache(tokens)
        
        n_layers = _model.cfg.n_layers
        n_heads = _model.cfg.n_heads
        
        heads_data = []
        top_heads = []
        
        for layer in range(n_layers):
            attn_pattern = cache[f"blocks.{layer}.attn.hook_pattern"][0]  # [n_heads, seq, seq]
            
            for head in range(n_heads):
                pattern = attn_pattern[head].detach().float().cpu().numpy().tolist()
                
                # Calculate entropy of attention distribution
                attn_probs = attn_pattern[head].detach().float().cpu().numpy()
                entropy = -np.sum(attn_probs * np.log(attn_probs + 1e-10), axis=-1).mean()
                
                # Max attention value
                max_attn = float(attn_pattern[head].max())
                
                # Classify head type based on pattern
                pat_len = len(pattern)
                head_type = "unknown"
                if pat_len > 1:
                    # Check for previous token attention (diagonal -1)
                    prev_score = sum(pattern[i][i-1] if i > 0 else 0 for i in range(pat_len)) / pat_len
                    # Check for first token attention
                    first_score = sum(pattern[i][0] for i in range(pat_len)) / pat_len
                    
                    if prev_score > 0.3:
                        head_type = "previous"
                    elif first_score > 0.4:
                        head_type = "global"
                    elif entropy < 1.0:
                        head_type = "local"
                    elif max_attn > 0.5 and pat_len > 5:
                        head_type = "induction"
                
                heads_data.append({
                    "layer": layer,
                    "head": head,
                    "pattern": pattern,
                    "entropy": round(float(entropy), 3),
                    "maxAttention": round(max_attn, 3),
                    "headType": head_type
                })
                
                # Score for ranking - high max attention + low entropy = interesting
                score = max_attn - entropy * 0.1
                top_heads.append({
                    "layer": layer,
                    "head": head,
                    "score": round(score, 3),
                    "type": head_type
                })
        
        # Sort and take top K
        top_heads.sort(key=lambda x: x["score"], reverse=True)
        top_heads = top_heads[:top_k]
        
        return jsonify(sanitize_for_json({
            "success": True,
            "tokens": token_strs,
            "heads": heads_data,
            "topHeads": top_heads
        }))
        
    except Exception as e:
        print(f"[Attention] Analyze error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Logit Lens Endpoints
# ============================================

@app.route('/logit-lens/analyze', methods=['POST'])
def logit_lens_analyze():
    """
    Project intermediate layer activations to vocabulary space.
    
    Request body:
    {
        "prompt": "The Eiffel Tower is located in",
        "model": "gpt2-small",
        "topK": 5
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gpt2-small')
        top_k = data.get('topK', 5)
        
        if not prompt:
            return jsonify({"error": "Prompt required"}), 400
        
        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, sae = get_model_and_sae(model_id)
        
        tokens = model.to_tokens(prompt)
        token_strs = model.to_str_tokens(prompt)
        
        n_layers = model.cfg.n_layers
        resid_filter = lambda name: name.endswith("hook_resid_post")
        
        with torch.no_grad():
            _, cache = model.run_with_cache(tokens, names_filter=resid_filter)
        layers_data = []
        
        W_U = model.W_U
        b_U = model.b_U if hasattr(model, 'b_U') and model.b_U is not None else None
        
        for layer in range(n_layers):
            resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :]
            
            if hasattr(model, 'ln_final'):
                resid_normed = model.ln_final(resid.unsqueeze(0)).squeeze(0)
            else:
                resid_normed = resid
            
            logits = (resid_normed.float() @ W_U.float())
            if b_U is not None:
                logits = logits + b_U.float()
            probs = torch.softmax(logits, dim=-1)
            
            top_probs, top_indices = probs.topk(top_k)
            
            top_tokens = []
            for i in range(top_k):
                tok_id = top_indices[i].item()
                tok_str = model.tokenizer.decode([tok_id])
                prob_val = top_probs[i].item()
                logit_val = logits[tok_id].item()
                top_tokens.append({
                    "token": tok_str,
                    "prob": round(prob_val, 4) if not np.isnan(prob_val) else 0.0,
                    "logit": round(logit_val, 3) if not np.isnan(logit_val) else 0.0
                })
            
            entropy = -torch.sum(probs * torch.log(probs + 1e-10)).item()
            if np.isnan(entropy) or np.isinf(entropy):
                entropy = 0.0
            
            layers_data.append({
                "layer": layer,
                "topTokens": top_tokens,
                "entropy": round(entropy, 3)
            })
        
        final_resid = cache[f"blocks.{n_layers - 1}.hook_resid_post"][0, -1, :]
        if hasattr(model, 'ln_final'):
            final_normed = model.ln_final(final_resid.unsqueeze(0)).squeeze(0)
        else:
            final_normed = final_resid
        final_logits = final_normed.float() @ W_U.float()
        if b_U is not None:
            final_logits = final_logits + b_U.float()
        final_probs = torch.softmax(final_logits, dim=-1)
        final_token_idx = final_probs.argmax().item()
        final_token = model.tokenizer.decode([final_token_idx])
        
        top_probs, top_indices = final_probs.topk(top_k)
        unembed_tokens = []
        for i in range(top_k):
            tok_id = top_indices[i].item()
            tok_str = model.tokenizer.decode([tok_id])
            prob_val = top_probs[i].item()
            logit_val = final_logits[tok_id].item()
            unembed_tokens.append({
                "token": tok_str,
                "prob": round(prob_val, 4) if not np.isnan(prob_val) else 0.0,
                "logit": round(logit_val, 3) if not np.isnan(logit_val) else 0.0
            })
        
        del cache
        
        return jsonify({
            "success": True,
            "prompt": prompt,
            "tokens": token_strs,
            "targetPosition": len(token_strs) - 1,
            "layers": layers_data,
            "unembedding": unembed_tokens,
            "finalPrediction": final_token
        })
        
    except Exception as e:
        print(f"[LogitLens] Analyze error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


_tuned_lens_probes = {}

TUNED_LENS_CALIBRATION_TEXTS = [
    "The capital of France is Paris, which is known for the Eiffel Tower.",
    "Machine learning models can be trained on large datasets to make predictions.",
    "The quick brown fox jumps over the lazy dog near the river bank.",
    "In 1969, Neil Armstrong became the first person to walk on the moon.",
    "Water boils at 100 degrees Celsius under standard atmospheric pressure.",
    "Shakespeare wrote many plays including Hamlet, Macbeth, and Romeo and Juliet.",
    "The mitochondria is often called the powerhouse of the cell in biology.",
    "Python is a popular programming language used for data science and web development.",
    "The stock market experienced significant volatility during the financial crisis.",
    "Photosynthesis converts sunlight into chemical energy that plants use to grow.",
    "Albert Einstein developed the theory of relativity in the early twentieth century.",
    "The Pacific Ocean is the largest and deepest ocean on Earth.",
    "Democracy is a system of government where citizens vote to elect their leaders.",
    "Antibiotics are medications used to treat bacterial infections in humans.",
    "The Great Wall of China stretches thousands of miles across northern China.",
    "Neural networks consist of layers of interconnected nodes that process information.",
]


def _train_tuned_lens_probes(model, base_model_id):
    """Train affine probes per layer: map raw resid_post to final-layer resid_post (pre-ln_final)."""
    global _tuned_lens_probes
    
    if base_model_id in _tuned_lens_probes:
        return _tuned_lens_probes[base_model_id]
    
    print(f"[TunedLens] Training probes for {base_model_id}...")
    n_layers = model.cfg.n_layers
    d_model = model.cfg.d_model
    
    all_resids = {layer: [] for layer in range(n_layers)}
    all_targets = []
    
    with torch.no_grad():
        for text in TUNED_LENS_CALIBRATION_TEXTS:
            tokens = model.to_tokens(text)
            _, cache = model.run_with_cache(tokens)
            
            seq_len = tokens.shape[1]
            for pos in range(max(0, seq_len - 8), seq_len):
                final_resid = cache[f"blocks.{n_layers - 1}.hook_resid_post"][0, pos, :]
                all_targets.append(final_resid)
                
                for layer in range(n_layers):
                    resid = cache[f"blocks.{layer}.hook_resid_post"][0, pos, :]
                    all_resids[layer].append(resid)
    
    targets = torch.stack(all_targets)
    
    probes = {}
    for layer in range(n_layers):
        X = torch.stack(all_resids[layer])
        Y = targets
        
        X_bias = torch.cat([X, torch.ones(X.shape[0], 1)], dim=1)
        
        solution = torch.linalg.lstsq(X_bias, Y).solution
        
        W = solution[:d_model, :]
        b = solution[d_model, :]
        
        probes[layer] = {"W": W, "b": b}
        
        with torch.no_grad():
            pred = X @ W + b
            mse = ((pred - Y) ** 2).mean().item()
        print(f"  Layer {layer}: MSE = {mse:.6f}")
    
    _tuned_lens_probes[base_model_id] = probes
    print(f"[TunedLens] Probes trained for {base_model_id}")
    return probes


@app.route('/logit-lens/tuned-analyze', methods=['POST'])
def tuned_lens_analyze():
    """
    Tuned Lens: uses learned per-layer affine probes for cleaner intermediate projections.
    
    Request body:
    {
        "prompt": "The Eiffel Tower is located in",
        "model": "gpt2-small",
        "topK": 5
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gpt2-small')
        top_k = data.get('topK', 5)
        
        if not prompt:
            return jsonify({"error": "Prompt required"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or _current_model_id is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[TunedLens] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        probes = _train_tuned_lens_probes(_model, base_model)
        
        tokens = _model.to_tokens(prompt)
        token_strs = _model.to_str_tokens(prompt)
        
        with torch.no_grad():
            _, cache = _model.run_with_cache(tokens)
        
        n_layers = _model.cfg.n_layers
        W_U = _model.W_U
        b_U = _model.b_U if hasattr(_model, 'b_U') and _model.b_U is not None else None
        
        layers_data = []
        
        for layer in range(n_layers):
            resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :]
            
            probe = probes[layer]
            tuned_resid = resid @ probe["W"] + probe["b"]
            
            if hasattr(_model, 'ln_final'):
                tuned_resid = _model.ln_final(tuned_resid.unsqueeze(0)).squeeze(0)
            
            logits = tuned_resid @ W_U
            if b_U is not None:
                logits = logits + b_U
            probs = torch.softmax(logits, dim=-1)
            
            top_probs, top_indices = probs.topk(top_k)
            
            top_tokens = []
            for i in range(top_k):
                tok_id = top_indices[i].item()
                tok_str = _model.tokenizer.decode([tok_id])
                top_tokens.append({
                    "token": tok_str,
                    "prob": round(top_probs[i].item(), 4),
                    "logit": round(logits[tok_id].item(), 3)
                })
            
            entropy = -torch.sum(probs * torch.log(probs + 1e-10)).item()
            
            layers_data.append({
                "layer": layer,
                "topTokens": top_tokens,
                "entropy": round(entropy, 3)
            })
        
        final_resid = cache[f"blocks.{n_layers - 1}.hook_resid_post"][0, -1, :]
        if hasattr(_model, 'ln_final'):
            final_normed = _model.ln_final(final_resid.unsqueeze(0)).squeeze(0)
        else:
            final_normed = final_resid
        final_logits = final_normed @ W_U
        if b_U is not None:
            final_logits = final_logits + b_U
        final_probs = torch.softmax(final_logits, dim=-1)
        final_token_idx = final_probs.argmax().item()
        final_token = _model.tokenizer.decode([final_token_idx])
        
        top_probs, top_indices = final_probs.topk(top_k)
        unembed_tokens = []
        for i in range(top_k):
            tok_id = top_indices[i].item()
            tok_str = _model.tokenizer.decode([tok_id])
            unembed_tokens.append({
                "token": tok_str,
                "prob": round(top_probs[i].item(), 4),
                "logit": round(final_logits[tok_id].item(), 3)
            })
        
        return jsonify({
            "success": True,
            "prompt": prompt,
            "tokens": token_strs,
            "targetPosition": len(token_strs) - 1,
            "layers": layers_data,
            "unembedding": unembed_tokens,
            "finalPrediction": final_token,
            "mode": "tuned"
        })
        
    except Exception as e:
        print(f"[TunedLens] Analyze error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Probing Classifier Endpoints
# ============================================

@app.route('/probe/train', methods=['POST'])
def probe_train():
    """
    Train linear probes on activations at each layer.
    
    Request body:
    {
        "task": "sentiment",
        "examples": [{"text": "I love this", "label": "positive"}, ...],
        "model": "gpt2-small",
        "testSplit": 0.3
    }
    """
    try:
        data = request.json
        task = data.get('task', 'custom')
        examples = data.get('examples', [])
        model_id = data.get('model', 'gpt2-small')
        test_split = data.get('testSplit', 0.3)
        if not isinstance(test_split, (int, float)) or test_split <= 0 or test_split >= 1:
            test_split = 0.3
        
        if not examples or len(examples) < 20:
            unique_lbls = set(ex.get("label", "") for ex in (examples or []))
            test_split_val = request.json.get("test_split", 0.3)
            est_test = max(1, round(len(examples or []) * test_split_val))
            return jsonify({
                "error": f"Insufficient sample size: {len(examples or [])} examples provided, minimum 20 required.",
                "message": f"With {len(examples or [])} examples and {int(test_split_val*100)}% test split, only ~{est_test} test samples would be used — too few for reliable results. Provide at least 20 examples (10+ per label).",
                "minimum": 20,
                "provided": len(examples or []),
                "labels": list(unique_lbls),
                "estimatedTestSamples": est_test
            }), 400
        
        unique_labels = set(ex.get("label", "") for ex in examples)
        if len(unique_labels) < 2:
            return jsonify({"error": f"Need at least 2 different labels for classification, but all examples have the same label: '{next(iter(unique_labels))}'. Provide examples with contrasting labels (e.g. 'true'/'false', 'positive'/'negative')."}), 400
        
        label_counts = {}
        for lbl in unique_labels:
            label_counts[lbl] = sum(1 for ex in examples if ex.get("label") == lbl)
            if label_counts[lbl] < 2:
                return jsonify({"error": f"Label '{lbl}' has only {label_counts[lbl]} example(s). Each label needs at least 2 examples for train/test splitting."}), 400
        
        min_class_count = min(label_counts.values())
        n_test_needed = max(len(unique_labels), 2)
        total = len(examples)
        test_count = int(total * test_split)
        if test_count < n_test_needed:
            test_split = min(0.5, n_test_needed / total)
            test_count = int(total * test_split)
        if min_class_count < 2 or test_count < 1 or (total - test_count) < 1:
            return jsonify({"error": f"Not enough examples ({total}) for {len(unique_labels)} classes with test_split={test_split}. Provide more examples (at least {n_test_needed * 3} total recommended)."}), 400
        
        k_folds = data.get('kFolds', None)
        if k_folds is not None:
            try:
                k_folds = int(k_folds)
                if k_folds < 2 or k_folds > 10:
                    k_folds = None
                elif min_class_count < k_folds:
                    k_folds = min(min_class_count, 3)
                    if k_folds < 2:
                        k_folds = None
            except (ValueError, TypeError):
                k_folds = None
        
        from transformer_lens import HookedTransformer
        from sklearn.linear_model import LogisticRegression
        from sklearn.model_selection import train_test_split, StratifiedKFold
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        if _model is None or get_base_model_id(_current_model_id) != base_model:
            if _model is not None:
                _cleanup_model_vram()
            print(f"[Probe] Loading model: {base_model}")
            _model = _load_hooked_transformer(base_model)
            _saes = {}
            _current_model_id = model_id
        
        # Get unique labels
        labels = list(set(ex["label"] for ex in examples))
        label_to_idx = {l: i for i, l in enumerate(labels)}
        
        # Collect activations for each example at each layer
        n_layers = _model.cfg.n_layers
        layer_activations = {layer: [] for layer in range(n_layers)}
        y_labels = []
        texts = []
        
        for ex in examples:
            text = ex["text"]
            label = ex["label"]
            
            tokens = _model.to_tokens(text)
            with torch.no_grad():
                _, cache = _model.run_with_cache(tokens)
            
            for layer in range(n_layers):
                act = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :].detach().float().cpu().numpy()
                layer_activations[layer].append(act)
            
            del cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            y_labels.append(label_to_idx[label])
            texts.append(text)
        
        y_labels = np.array(y_labels)
        
        indices = np.arange(len(examples))
        layers_data = []
        best_layer = 0
        best_accuracy = 0
        
        if k_folds and k_folds >= 2:
            skf = StratifiedKFold(n_splits=k_folds, shuffle=True, random_state=42)
            
            for layer in range(n_layers):
                X = np.array(layer_activations[layer])
                fold_accuracies = []
                fold_details = []
                all_predictions = []
                
                for fold_idx, (train_idx, test_idx) in enumerate(skf.split(X, y_labels)):
                    X_train, X_test = X[train_idx], X[test_idx]
                    y_train, y_test = y_labels[train_idx], y_labels[test_idx]
                    
                    clf = LogisticRegression(max_iter=1000, random_state=42)
                    clf.fit(X_train, y_train)
                    y_pred = clf.predict(X_test)
                    acc = float((y_pred == y_test).mean())
                    fold_accuracies.append(acc)
                    fold_details.append({"fold": fold_idx + 1, "accuracy": round(acc, 3), "nTrain": len(train_idx), "nTest": len(test_idx)})
                    
                    for i, idx in enumerate(test_idx):
                        all_predictions.append({
                            "input": texts[idx],
                            "predicted": labels[y_pred[i]],
                            "actual": labels[y_test[i]],
                            "correct": bool(y_pred[i] == y_test[i]),
                            "fold": fold_idx + 1,
                        })
                
                mean_acc = float(np.mean(fold_accuracies))
                std_acc = float(np.std(fold_accuracies))
                
                layers_data.append({
                    "layer": layer,
                    "accuracy": round(mean_acc, 3),
                    "accuracyStd": round(std_acc, 3),
                    "folds": fold_details,
                    "kFolds": k_folds,
                    "predictions": all_predictions,
                })
                
                if mean_acc > best_accuracy:
                    best_accuracy = mean_acc
                    best_layer = layer
            
            return jsonify({
                "success": True,
                "task": task,
                "layers": layers_data,
                "bestLayer": best_layer,
                "bestAccuracy": round(float(best_accuracy), 3),
                "crossValidation": {"kFolds": k_folds, "method": "stratified_k_fold"},
            })
        
        else:
            train_idx, test_idx = train_test_split(indices, test_size=test_split, random_state=42, stratify=y_labels)
            
            for layer in range(n_layers):
                X = np.array(layer_activations[layer])
                X_train, X_test = X[train_idx], X[test_idx]
                y_train, y_test = y_labels[train_idx], y_labels[test_idx]
                
                clf = LogisticRegression(max_iter=1000, random_state=42)
                clf.fit(X_train, y_train)
                
                y_pred = clf.predict(X_test)
                accuracy = (y_pred == y_test).mean()
                
                predictions = []
                for i, idx in enumerate(test_idx):
                    predictions.append({
                        "input": texts[idx],
                        "predicted": labels[y_pred[i]],
                        "actual": labels[y_test[i]],
                        "correct": bool(y_pred[i] == y_test[i])
                    })
                
                layers_data.append({
                    "layer": layer,
                    "accuracy": round(float(accuracy), 3),
                    "predictions": predictions
                })
                
                if accuracy > best_accuracy:
                    best_accuracy = accuracy
                    best_layer = layer
            
            return jsonify({
                "success": True,
                "task": task,
                "layers": layers_data,
                "bestLayer": best_layer,
                "bestAccuracy": round(float(best_accuracy), 3)
            })
        
    except Exception as e:
        print(f"[Probe] Train error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# SAE Feature Comparison
# ============================================

@app.route('/sae/compare', methods=['POST'])
def sae_compare():
    """
    Compare feature activations across two prompt groups.
    
    Request body:
    {
        "groupA": ["prompt 1", "prompt 2", ...],
        "groupB": ["prompt 3", "prompt 4", ...],
        "model": "gpt2-small",
        "layer": 6,
        "topK": 20
    }
    
    Response:
    {
        "success": true,
        "totalFeatures": 24576,
        "topFeatures": [
            {
                "featureIndex": 1234,
                "groupARate": 0.8,
                "groupBRate": 0.2,
                "differential": 0.6,
                "groupAPrompts": [...],
                "groupBPrompts": [...]
            }
        ],
        "dominantGroup": "A"
    }
    """
    try:
        data = request.json
        group_a = data.get('groupA', [])
        group_b = data.get('groupB', [])
        model_id = data.get('model', 'gpt2-small')
        layer = data.get('layer', 6)
        top_k = data.get('topK', 20)
        robust_config = data.get('robustExtraction', None)
        feature_ids = data.get('featureIds', None)
        
        if not group_a or not group_b:
            return jsonify({"error": "Both prompt groups required"}), 400
        
        model, sae = get_model_and_sae(model_id, layer=layer)
        n_features = sae.cfg.d_sae
        hook_name = get_hook_name(model_id, layer=layer)
        
        # Collect activations for each group
        def get_group_activations(prompts):
            all_activations = []
            for prompt in prompts:
                tokens = model.to_tokens(prompt)
                with torch.no_grad():
                    _, cache = model.run_with_cache(tokens)
                    acts = safe_get_cache(cache, hook_name)[0, -1, :]
                    sae_acts = sae.encode(acts)
                    all_activations.append(sae_acts)
                del cache
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            return torch.stack(all_activations)
        
        group_a_acts = get_group_activations(group_a)  # [n_prompts_a, n_features]
        group_b_acts = get_group_activations(group_b)  # [n_prompts_b, n_features]
        
        # Compute activation rates (fraction of prompts where feature activates above threshold)
        threshold = 0.1
        group_a_rates = (group_a_acts > threshold).float().mean(dim=0)  # [n_features]
        group_b_rates = (group_b_acts > threshold).float().mean(dim=0)  # [n_features]
        
        # Count features that activated in any prompt (either group)
        combined_acts = torch.cat([group_a_acts, group_b_acts], dim=0)  # [n_all_prompts, n_features]
        activated_features = int((combined_acts.max(dim=0).values > threshold).sum().item())
        
        # Differential score: A rate minus B rate
        differential = group_a_rates - group_b_rates
        
        # Get top features - either from specified feature set or all features
        if feature_ids and len(feature_ids) > 0:
            # Filter to only the specified features, then get top_k from those
            valid_ids = [fid for fid in feature_ids if 0 <= fid < n_features]
            if not valid_ids:
                return jsonify({"error": "No valid feature IDs in specified range"}), 400
            
            # Create tensor of differential values for specified features
            subset_diffs = torch.tensor([differential[fid].item() for fid in valid_ids])
            # Get top_k indices within the subset
            k = min(top_k, len(valid_ids))
            top_subset_indices = subset_diffs.abs().topk(k).indices.tolist()
            # Map back to actual feature indices
            top_indices = [valid_ids[i] for i in top_subset_indices]
            scanned_features = len(valid_ids)
        else:
            # Scan all features
            top_indices = differential.abs().topk(top_k).indices.tolist()
            scanned_features = n_features
        
        # Statistical significance helper functions
        import numpy as np
        
        def bootstrap_ci(a_binary, b_binary, n_bootstrap=500, alpha=0.05):
            """Compute bootstrap 95% CI for difference in activation rates"""
            diffs = []
            n_a, n_b = len(a_binary), len(b_binary)
            for _ in range(n_bootstrap):
                a_sample = np.random.choice(a_binary, size=n_a, replace=True)
                b_sample = np.random.choice(b_binary, size=n_b, replace=True)
                diffs.append(a_sample.mean() - b_sample.mean())
            diffs = np.sort(diffs)
            low_idx = int(alpha / 2 * n_bootstrap)
            high_idx = int((1 - alpha / 2) * n_bootstrap)
            return float(diffs[low_idx]), float(diffs[high_idx])
        
        def permutation_pvalue(a_binary, b_binary, observed_diff, n_perm=500):
            """Compute permutation test p-value"""
            pooled = np.concatenate([a_binary, b_binary])
            n_a = len(a_binary)
            extreme_count = 0
            for _ in range(n_perm):
                np.random.shuffle(pooled)
                perm_a = pooled[:n_a]
                perm_b = pooled[n_a:]
                perm_diff = abs(perm_a.mean() - perm_b.mean())
                if perm_diff >= abs(observed_diff):
                    extreme_count += 1
            return extreme_count / n_perm
        
        def compute_robustness(feature_idx, group_a_acts_tensor, group_b_acts_tensor, 
                               threshold, n_bootstrap=100, holdout_ratio=0.2, min_stability=0.7):
            """Compute robustness metrics for a feature via bootstrap and cross-validation"""
            n_a = group_a_acts_tensor.shape[0]
            n_b = group_b_acts_tensor.shape[0]
            
            a_acts = group_a_acts_tensor[:, feature_idx].numpy()
            b_acts = group_b_acts_tensor[:, feature_idx].numpy()
            
            # Bootstrap stability: how often does this feature remain in top differential set?
            stability_count = 0
            original_diff = abs((a_acts > threshold).mean() - (b_acts > threshold).mean())
            
            for _ in range(n_bootstrap):
                a_sample = np.random.choice(a_acts, size=n_a, replace=True)
                b_sample = np.random.choice(b_acts, size=n_b, replace=True)
                boot_diff = abs((a_sample > threshold).mean() - (b_sample > threshold).mean())
                # Feature is "stable" if bootstrap diff is at least 50% of original
                if boot_diff >= original_diff * 0.5:
                    stability_count += 1
            
            stability = stability_count / n_bootstrap
            
            # Cross-validation: holdout some prompts and test if feature still differentiates
            holdout_a = max(1, int(n_a * holdout_ratio))
            holdout_b = max(1, int(n_b * holdout_ratio))
            
            cv_scores = []
            for _ in range(min(50, n_bootstrap)):
                # Random split
                a_indices = np.random.permutation(n_a)
                b_indices = np.random.permutation(n_b)
                
                train_a = a_acts[a_indices[holdout_a:]]
                test_a = a_acts[a_indices[:holdout_a]]
                train_b = b_acts[b_indices[holdout_b:]]
                test_b = b_acts[b_indices[:holdout_b]]
                
                # Train differential
                train_diff = (train_a > threshold).mean() - (train_b > threshold).mean()
                
                # Test if direction holds
                test_diff = (test_a > threshold).mean() - (test_b > threshold).mean()
                
                # Score 1 if same sign, 0 otherwise
                if train_diff * test_diff > 0 or (abs(train_diff) < 0.05 and abs(test_diff) < 0.05):
                    cv_scores.append(1)
                else:
                    cv_scores.append(0)
            
            cv_score = np.mean(cv_scores) if cv_scores else 0.5
            bootstrap_consistent = stability >= min_stability and cv_score >= 0.6
            
            return {
                "stability": round(stability, 3),
                "crossValidationScore": round(cv_score, 3),
                "bootstrapConsistent": bootstrap_consistent
            }
        
        # Build response with statistical significance
        top_features = []
        for idx in top_indices:
            # Get binary activation arrays for this feature
            a_binary = (group_a_acts[:, idx] > threshold).float().numpy()
            b_binary = (group_b_acts[:, idx] > threshold).float().numpy()
            obs_diff = differential[idx].item()
            
            # Compute stats only for meaningful differences with enough data
            p_value = None
            ci_low, ci_high = None, None
            total_prompts = len(a_binary) + len(b_binary)
            if abs(obs_diff) > 0.1 and total_prompts >= 6:
                p_value = round(permutation_pvalue(a_binary, b_binary, obs_diff), 3)
                ci_low, ci_high = bootstrap_ci(a_binary, b_binary)
                ci_low, ci_high = round(ci_low, 3), round(ci_high, 3)
            
            # Compute robustness metrics if enabled
            robustness = None
            if robust_config and robust_config.get('enabled'):
                robustness = compute_robustness(
                    idx, group_a_acts, group_b_acts, threshold,
                    n_bootstrap=robust_config.get('bootstrapSamples', 100),
                    holdout_ratio=robust_config.get('holdoutRatio', 0.2),
                    min_stability=robust_config.get('minStability', 0.7)
                )
            
            feature_data = {
                "featureIndex": idx,
                "groupARate": round(group_a_rates[idx].item(), 3),
                "groupBRate": round(group_b_rates[idx].item(), 3),
                "differential": round(obs_diff, 3),
                "pValue": p_value,
                "confidenceInterval": {"low": ci_low, "high": ci_high} if ci_low is not None else None,
                "robustness": robustness,
                "groupAPrompts": [
                    {"prompt": p, "activation": round(group_a_acts[i, idx].item(), 3)}
                    for i, p in enumerate(group_a)
                ],
                "groupBPrompts": [
                    {"prompt": p, "activation": round(group_b_acts[i, idx].item(), 3)}
                    for i, p in enumerate(group_b)
                ]
            }
            top_features.append(feature_data)
        
        # Sort by absolute differential (highest first)
        top_features.sort(key=lambda x: abs(x["differential"]), reverse=True)
        
        # Determine dominant group
        mean_diff = differential.mean().item()
        if mean_diff > 0.05:
            dominant = "A"
        elif mean_diff < -0.05:
            dominant = "B"
        else:
            dominant = "neutral"
        
        return jsonify({
            "success": True,
            "totalFeatures": n_features,
            "scannedFeatures": scanned_features,
            "activatedFeatures": activated_features,
            "topFeatures": top_features,
            "dominantGroup": dominant,
            "model": model_id,
            "layer": layer
        })
        
    except Exception as e:
        print(f"[Backlight] Compare error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Transcoder Circuit Discovery Endpoints (Circuits)
# ============================================

# Lazy loading for transcoders
_transcoder = None
_transcoder_model_id = None


def get_transcoder(model_id: str = "gpt2-small", layer: int = 6):
    """Load transcoder (MLP SAE) for circuit discovery"""
    global _transcoder, _transcoder_model_id, _model, _current_model_id
    
    cache_key = f"{model_id}_layer{layer}"
    if _transcoder is not None and _transcoder_model_id == cache_key:
        return _transcoder, _model
    
    print(f"[Circuits] Loading transcoder for: {model_id} layer {layer}")
    
    from transformer_lens import HookedTransformer
    
    base_model = get_base_model_id(model_id)
    with _model_swap_lock:
        if _model is None or get_base_model_id(_current_model_id) != base_model:
            if _model is not None:
                _cleanup_model_vram()
            _model = _load_hooked_transformer(base_model)
            _current_model_id = model_id
    
    try:
        from sae_lens import SAE
        if "gpt2" in model_id:
            _transcoder, _, _ = SAE.from_pretrained(
                release="gpt2-small-mlp-out-v5-32k",
                sae_id=f"blocks.{layer}.hook_mlp_out",
                device=_get_device()
            )
        else:
            mlp_sae_id = f"layer_{layer}/width_16k/average_l0_{GEMMA_MLP_L0}"
            print(f"[Circuits] Loading MLP SAE: gemma-scope-2b-pt-mlp / {mlp_sae_id}")
            _transcoder, _, _ = SAE.from_pretrained(
                release="gemma-scope-2b-pt-mlp",
                sae_id=mlp_sae_id,
                device=_get_device()
            )
        _transcoder_model_id = cache_key
        print(f"[Circuits] Loaded MLP SAE as transcoder proxy for layer {layer}")
    except Exception as e:
        print(f"[Circuits] Could not load transcoder: {e}")
        _transcoder = None
    
    return _transcoder, _model


@app.route('/transcoder/discover', methods=['POST'])
def transcoder_discover():
    """
    Discover feature circuits from MLP input to output.
    
    Request body:
    {
        "prompt": "The capital of France is",
        "model": "gemma-2-2b",
        "layer": 12,
        "topK": 32,
        "threshold": 0.1
    }
    
    Response:
    {
        "nodes": [
            {"id": "in_123", "type": "input", "feature": 123, "activation": 0.5},
            {"id": "out_456", "type": "output", "feature": 456, "activation": 0.8}
        ],
        "edges": [
            {"source": "in_123", "target": "out_456", "weight": 0.3}
        ],
        "tokens": ["The", " capital", ...],
        "targetToken": "Paris"
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gemma-2-2b')
        layer = data.get('layer', 12)
        top_k = data.get('topK', 32)
        threshold = data.get('threshold', 0.1)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        transcoder, model = get_transcoder(model_id, layer)
        
        if transcoder is None:
            return jsonify({"error": "Transcoder not available for this model"}), 400
        
        tokens = model.to_tokens(prompt)
        str_tokens = [t.replace('Ġ', ' ').replace('Ċ', '\n') for t in model.to_str_tokens(prompt)]
        
        with torch.no_grad():
            # Get MLP input and output
            _, cache = model.run_with_cache(tokens)
            
            mlp_in = cache[f"blocks.{layer}.hook_resid_pre"][0, -1, :]  # Last token
            mlp_out = cache[f"blocks.{layer}.hook_mlp_out"][0, -1, :]
            
            # Encode through transcoder
            input_features = transcoder.encode(mlp_in.unsqueeze(0).unsqueeze(0))
            output_features = transcoder.encode(mlp_out.unsqueeze(0).unsqueeze(0))
            
            input_acts = input_features[0, 0, :]
            output_acts = output_features[0, 0, :]
            
            # Get top-k input features
            top_in_vals, top_in_idx = torch.topk(input_acts, k=min(top_k, len(input_acts)))
            # Get top-k output features
            top_out_vals, top_out_idx = torch.topk(output_acts, k=min(top_k, len(output_acts)))
        
        # Build circuit graph
        nodes = []
        edges = []
        
        # Input feature nodes
        for i, (idx, val) in enumerate(zip(top_in_idx.tolist(), top_in_vals.tolist())):
            if val > threshold:
                nodes.append({
                    "id": f"in_{idx}",
                    "type": "input",
                    "feature": idx,
                    "activation": round(val, 4),
                    "layer": layer
                })
        
        # Output feature nodes
        for i, (idx, val) in enumerate(zip(top_out_idx.tolist(), top_out_vals.tolist())):
            if val > threshold:
                nodes.append({
                    "id": f"out_{idx}",
                    "type": "output",
                    "feature": idx,
                    "activation": round(val, 4),
                    "layer": layer
                })
        
        # Compute edges based on decoder weight contributions
        # Edge weight = input_activation * W_dec contribution to output
        with torch.no_grad():
            W_dec = transcoder.W_dec  # [d_sae, d_model]
            
            for in_node in [n for n in nodes if n["type"] == "input"]:
                in_idx = in_node["feature"]
                in_act = in_node["activation"]
                
                for out_node in [n for n in nodes if n["type"] == "output"]:
                    out_idx = out_node["feature"]
                    
                    # Approximate edge weight via decoder correlation
                    # This is simplified - real implementation would use transcoder decoder
                    weight = abs(in_act * output_acts[out_idx].item()) / 10.0
                    
                    if weight > 0.01:  # Only significant edges
                        edges.append({
                            "source": in_node["id"],
                            "target": out_node["id"],
                            "weight": round(weight, 4)
                        })
        
        # Get model's next token prediction
        with torch.no_grad():
            logits = model(tokens)
            probs = torch.softmax(logits[0, -1, :], dim=-1)
            top_token_idx = probs.argmax().item()
            target_token = model.to_single_str_token(top_token_idx).replace('Ġ', ' ')
        
        return jsonify({
            "success": True,
            "nodes": nodes,
            "edges": edges,
            "tokens": str_tokens,
            "targetToken": target_token,
            "layer": layer,
            "model": model_id
        })
        
    except Exception as e:
        print(f"[Circuits] Discover error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# Cache for multi-layer SAEs
_multilayer_saes = {}

@app.route('/transcoder/discover-multilayer', methods=['POST'])
def transcoder_discover_multilayer():
    """
    Discover feature activations across multiple layers to trace circuit evolution.
    Supports both GPT-2 Small and Gemma 2 2B.
    
    Request body:
    {
        "prompt": "The capital of France is",
        "model": "gemma-2-2b",
        "layers": [6, 12, 18, 24],
        "topK": 10,
        "threshold": 0.1
    }
    """
    global _multilayer_saes, _model, _current_model_id
    
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gemma-2-2b')
        layers = data.get('layers', [6, 12, 18, 24])
        top_k = data.get('topK', 10)
        threshold = data.get('threshold', 0.1)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        is_gemma = "gemma" in model_id
        base_model = get_base_model_id(model_id)
        
        from transformer_lens import HookedTransformer
        from sae_lens import SAE
        
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[Circuits] Loading {base_model} for multi-layer analysis...")
                _model = _load_hooked_transformer(base_model)
                _current_model_id = model_id
        
        for layer in layers:
            cache_key = f"{base_model}_mlp_{layer}"
            if cache_key not in _multilayer_saes:
                try:
                    print(f"[Circuits] Loading MLP SAE for {base_model} layer {layer}...")
                    if is_gemma:
                        sae, _, _ = SAE.from_pretrained(
                            release="gemma-scope-2b-pt-mlp",
                            sae_id=f"layer_{layer}/width_16k/average_l0_{GEMMA_MLP_L0}",
                            device=_get_device()
                        )
                    else:
                        sae, _, _ = SAE.from_pretrained(
                            release="gpt2-small-mlp-out-v5-32k",
                            sae_id=f"blocks.{layer}.hook_mlp_out",
                            device=_get_device()
                        )
                    _multilayer_saes[cache_key] = sae
                    print(f"[Circuits] Loaded SAE for {base_model} layer {layer}")
                except Exception as e:
                    print(f"[Circuits] Could not load SAE for layer {layer}: {e}")
                    _multilayer_saes[cache_key] = None
        
        tokens = _model.to_tokens(prompt)
        str_tokens = _model.to_str_tokens(prompt)
        if not is_gemma:
            str_tokens = [t.replace('Ġ', ' ').replace('Ċ', '\n') for t in str_tokens]
        else:
            str_tokens = list(str_tokens)
        
        with torch.no_grad():
            _, cache = _model.run_with_cache(tokens)
            
            logits = _model(tokens)
            probs = torch.softmax(logits[0, -1, :], dim=-1)
            top_token_idx = probs.argmax().item()
            target_token = _model.to_single_str_token(top_token_idx)
            if not is_gemma:
                target_token = target_token.replace('Ġ', ' ')
        
        layer_results = []
        all_feature_activations = {}
        
        with torch.no_grad():
            for layer in layers:
                cache_key = f"{base_model}_mlp_{layer}"
                sae = _multilayer_saes.get(cache_key)
                
                if sae is None:
                    layer_results.append({
                        "layer": layer,
                        "features": [],
                        "error": "SAE not available for this layer"
                    })
                    continue
                
                hook_name = f"blocks.{layer}.hook_mlp_out"
                try:
                    mlp_out = cache[hook_name][0, -1, :]
                except KeyError:
                    resid_hook = f"blocks.{layer}.hook_resid_post"
                    mlp_out = cache[resid_hook][0, -1, :]
                
                features = sae.encode(mlp_out.unsqueeze(0).unsqueeze(0))
                acts = features[0, 0, :]
                
                top_vals, top_idx = torch.topk(acts, k=min(top_k, len(acts)))
                
                layer_features = []
                for idx, val in zip(top_idx.tolist(), top_vals.tolist()):
                    if val > threshold:
                        layer_features.append({
                            "feature": idx,
                            "activation": round(val, 4)
                        })
                        
                        if idx not in all_feature_activations:
                            all_feature_activations[idx] = {}
                        all_feature_activations[idx][layer] = val
                
                layer_results.append({
                    "layer": layer,
                    "features": layer_features
                })
        
        persistence = []
        for feature_id, layer_acts in all_feature_activations.items():
            if len(layer_acts) >= 2:
                sorted_layers = sorted(layer_acts.keys())
                peak_layer = max(layer_acts, key=layer_acts.get)
                persistence.append({
                    "feature": feature_id,
                    "firstLayer": sorted_layers[0],
                    "lastLayer": sorted_layers[-1],
                    "peakLayer": peak_layer,
                    "peakActivation": round(layer_acts[peak_layer], 4),
                    "layerCount": len(layer_acts),
                    "trajectory": {str(l): round(a, 4) for l, a in sorted(layer_acts.items())}
                })
        
        persistence.sort(key=lambda x: (-x["layerCount"], -x["peakActivation"]))
        
        return jsonify({
            "success": True,
            "layers": layer_results,
            "persistence": persistence[:20],
            "tokens": str_tokens,
            "targetToken": target_token,
            "model": model_id,
            "requestedLayers": layers
        })
        
    except Exception as e:
        print(f"[Circuits] Multi-layer discover error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/transcoder/ablate', methods=['POST'])
def transcoder_ablate():
    """
    Ablate specific transcoder features to verify circuit causality.
    
    Request body:
    {
        "prompt": "The capital of France is",
        "model": "gemma-2-2b",
        "layer": 12,
        "featuresToAblate": [123, 456]
    }
    
    Response:
    {
        "baselineOutput": "Paris",
        "ablatedOutput": "London",
        "logitDifference": 0.45,
        "featuresCausalImpact": [
            {"feature": 123, "impact": 0.3},
            {"feature": 456, "impact": 0.15}
        ]
    }
    """
    try:
        data = request.json
        prompt = data.get('prompt', '')
        model_id = data.get('model', 'gemma-2-2b')
        layer = data.get('layer', 12)
        features_to_ablate = data.get('featuresToAblate', [])
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        
        transcoder, model = get_transcoder(model_id)
        
        if transcoder is None:
            return jsonify({"error": "Transcoder not available"}), 400
        
        tokens = model.to_tokens(prompt)
        
        with torch.no_grad():
            # Baseline
            baseline_logits = model(tokens)
            baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
            baseline_token = model.to_single_str_token(baseline_probs.argmax().item())
            
            baseline_gen = model.generate(tokens, max_new_tokens=5, temperature=0.1)
            baseline_output = model.to_string(baseline_gen[0])[len(model.to_string(tokens[0])):]
            
            # Create ablation hook
            def ablation_hook(activations, hook):
                # Encode, ablate features, decode
                features = transcoder.encode(activations)
                for feat_idx in features_to_ablate:
                    if feat_idx < features.shape[-1]:
                        features[:, :, feat_idx] = 0.0
                return transcoder.decode(features)
            
            # Ablated
            model.reset_hooks()
            model.add_hook(f"blocks.{layer}.hook_mlp_out", ablation_hook)
            
            ablated_logits = model(tokens)
            ablated_probs = torch.softmax(ablated_logits[0, -1, :], dim=-1)
            ablated_token = model.to_single_str_token(ablated_probs.argmax().item())
            
            ablated_gen = model.generate(tokens, max_new_tokens=5, temperature=0.1)
            ablated_output = model.to_string(ablated_gen[0])[len(model.to_string(tokens[0])):]
            
            model.reset_hooks()
            
            # Compute logit difference
            top_k = 10
            top_idx = torch.topk(baseline_probs, k=top_k).indices
            logit_diff = (baseline_probs[top_idx] - ablated_probs[top_idx]).abs().sum().item()
        
        # Compute per-feature causal impact
        feature_impacts = []
        for feat in features_to_ablate:
            # Simplified: distribute total impact equally
            impact = logit_diff / len(features_to_ablate) if features_to_ablate else 0
            feature_impacts.append({
                "feature": feat,
                "impact": round(impact, 4)
            })
        
        return jsonify({
            "success": True,
            "baselineOutput": baseline_output,
            "ablatedOutput": ablated_output,
            "baselineToken": baseline_token.replace('Ġ', ' '),
            "ablatedToken": ablated_token.replace('Ġ', ' '),
            "logitDifference": round(logit_diff, 4),
            "featuresCausalImpact": feature_impacts,
            "layer": layer
        })
        
    except Exception as e:
        print(f"[Circuits] Ablate error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/domain-profile', methods=['POST'])
def domain_profile():
    """
    Profile which SAE features are domain-specific vs shared vs irrelevant.
    Runs domain and baseline texts through the model, encodes with SAE,
    and classifies features by activation frequency.
    
    Request body:
    {
        "model": "gpt2-small",
        "domain_texts": ["text1", "text2", ...],
        "baseline_texts": ["text1", "text2", ...],
        "activation_threshold": 0.5
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        domain_texts = data.get('domain_texts', [])
        baseline_texts = data.get('baseline_texts', [])
        threshold = data.get('activation_threshold', 0.5)
        
        if not domain_texts:
            return jsonify({"error": "No domain texts provided"}), 400
        if not baseline_texts:
            return jsonify({"error": "No baseline texts provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        
        sae_dim = sae.W_dec.shape[0] if hasattr(sae, 'W_dec') else 24576
        
        domain_activation_counts = np.zeros(sae_dim)
        baseline_activation_counts = np.zeros(sae_dim)
        domain_activation_strengths = np.zeros(sae_dim)
        baseline_activation_strengths = np.zeros(sae_dim)
        
        print(f"[SAE Service] Domain profiling: {len(domain_texts)} domain texts, {len(baseline_texts)} baseline texts")
        
        with torch.no_grad():
            for i, text in enumerate(domain_texts):
                print(f"[SAE Service] Profiling domain text {i+1}/{len(domain_texts)}")
                tokens = model.to_tokens(text)
                _, cache = model.run_with_cache(tokens)
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
                max_acts = feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
                domain_activation_counts += (max_acts > threshold).astype(float)
                domain_activation_strengths += max_acts
            
            for i, text in enumerate(baseline_texts):
                print(f"[SAE Service] Profiling baseline text {i+1}/{len(baseline_texts)}")
                tokens = model.to_tokens(text)
                _, cache = model.run_with_cache(tokens)
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
                max_acts = feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
                baseline_activation_counts += (max_acts > threshold).astype(float)
                baseline_activation_strengths += max_acts
        
        domain_rates = domain_activation_counts / len(domain_texts)
        baseline_rates = baseline_activation_counts / len(baseline_texts)
        domain_avg_strength = domain_activation_strengths / len(domain_texts)
        baseline_avg_strength = baseline_activation_strengths / len(baseline_texts)
        
        domain_critical = []
        shared = []
        irrelevant = []
        feature_details = {}
        
        for idx in range(sae_dim):
            dr = float(domain_rates[idx])
            br = float(baseline_rates[idx])
            ds = float(domain_avg_strength[idx])
            bs = float(baseline_avg_strength[idx])
            
            if dr >= 0.3 and (dr - br) >= 0.2:
                classification = "domain_critical"
                domain_critical.append(idx)
            elif dr >= 0.2 or br >= 0.2:
                classification = "shared"
                shared.append(idx)
            else:
                classification = "irrelevant"
                irrelevant.append(idx)
            
            if dr > 0.01 or br > 0.01:
                feature_details[str(idx)] = {
                    "domainRate": round(dr, 4),
                    "baselineRate": round(br, 4),
                    "domainStrength": round(ds, 4),
                    "baselineStrength": round(bs, 4),
                    "classification": classification
                }
        
        top_domain = sorted(domain_critical, key=lambda x: float(domain_rates[x]), reverse=True)[:50]
        top_shared = sorted(shared, key=lambda x: max(float(domain_rates[x]), float(baseline_rates[x])), reverse=True)[:30]
        
        return jsonify({
            "success": True,
            "totalFeatures": sae_dim,
            "domainCritical": len(domain_critical),
            "shared": len(shared),
            "irrelevant": len(irrelevant),
            "topDomainFeatures": top_domain,
            "topSharedFeatures": top_shared,
            "featureDetails": feature_details,
            "model": model_id,
            "domainTextsCount": len(domain_texts),
            "baselineTextsCount": len(baseline_texts)
        })
        
    except Exception as e:
        print(f"[SAE Service] Domain profile error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/pruned-inference', methods=['POST'])
def pruned_inference():
    """
    Run inference with only retained features active, zeroing all others.
    Compares full model output with pruned model output.
    
    Request body:
    {
        "model": "gpt2-small",
        "prompt": "The patient presented with",
        "retained_features": [100, 200, 300, ...],
        "max_new_tokens": 30
    }
    """
    try:
        data = request.json
        model_id = data.get('model', 'gpt2-small')
        prompt = data.get('prompt', '')
        retained_features = set(data.get('retained_features', []))
        max_new_tokens = data.get('max_new_tokens', 30)
        
        if not prompt:
            return jsonify({"error": "No prompt provided"}), 400
        if not retained_features:
            return jsonify({"error": "No retained features provided"}), 400
        
        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        sae_dim = sae.W_dec.shape[0] if hasattr(sae, 'W_dec') else 24576
        
        tokens = model.to_tokens(prompt)
        
        with torch.no_grad():
            baseline_logits = model(tokens)
            baseline_probs = torch.softmax(baseline_logits[0, -1, :], dim=-1)
            baseline_top5 = torch.topk(baseline_probs, k=5)
            baseline_tokens = [model.to_single_str_token(t.item()).replace('\u0120', ' ') for t in baseline_top5.indices]
            baseline_probs_list = [round(p.item(), 4) for p in baseline_top5.values]
            
            original_output = model.generate(
                tokens, max_new_tokens=max_new_tokens, temperature=0.1
            )
            original_text = model.to_string(original_output[0])
            original_continuation = original_text[len(model.to_string(tokens[0])):]
        
        retained_set = retained_features
        
        mask_tensor = torch.zeros(sae_dim, device=_get_device())
        for f in retained_set:
            if f < sae_dim:
                mask_tensor[f] = 1.0
        
        def pruning_hook(activations, hook):
            orig_acts = sae.encode(activations)
            orig_recon = sae.decode(orig_acts)
            recon_error = activations - orig_recon
            pruned_acts = orig_acts * mask_tensor
            return sae.decode(pruned_acts) + recon_error
        
        try:
            with torch.no_grad():
                model.reset_hooks()
                pruned_logits = model.run_with_hooks(
                    tokens,
                    fwd_hooks=[(hook_name, pruning_hook)],
                    reset_hooks_end=True
                )
                pruned_probs = torch.softmax(pruned_logits[0, -1, :], dim=-1)
                pruned_top5 = torch.topk(pruned_probs, k=5)
                pruned_tokens = [model.to_single_str_token(t.item()).replace('\u0120', ' ') for t in pruned_top5.indices]
                pruned_probs_list = [round(p.item(), 4) for p in pruned_top5.values]
                
                eps = 1e-10
                kl_div = (baseline_probs * torch.log((baseline_probs + eps) / (pruned_probs + eps))).sum().item()
                
                model.reset_hooks()
                model.add_hook(hook_name, pruning_hook)
                pruned_output = model.generate(
                    tokens, max_new_tokens=max_new_tokens, temperature=0.1
                )
                pruned_text = model.to_string(pruned_output[0])
                pruned_continuation = pruned_text[len(model.to_string(tokens[0])):]
        finally:
            model.reset_hooks()
        
        fidelity = max(0, 1.0 - kl_div) * 100
        
        return jsonify({
            "success": True,
            "original": {
                "continuation": original_continuation,
                "topTokens": baseline_tokens,
                "topProbs": baseline_probs_list
            },
            "pruned": {
                "continuation": pruned_continuation,
                "topTokens": pruned_tokens,
                "topProbs": pruned_probs_list
            },
            "klDivergence": round(kl_div, 6),
            "fidelity": round(fidelity, 2),
            "totalFeatures": sae_dim,
            "retainedFeatures": len(retained_features),
            "prunedFeatures": sae_dim - len(retained_features),
            "reductionPercent": round((1 - len(retained_features) / sae_dim) * 100, 1),
            "model": model_id
        })
        
    except Exception as e:
        print(f"[SAE Service] Pruned inference error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/tokenize-compare', methods=['POST'])
def tokenize_compare():
    """
    Compare how different tokenizers break down the same text.
    Loads only tokenizers (not full models) for speed.
    
    Request body:
    {
        "text": "Hello world こんにちは世界",
        "tokenizers": ["gpt2", "cl100k_base", "llama3", "gemma"]
    }
    """
    try:
        data = request.json
        text = data.get('text', '')
        requested = data.get('tokenizers', ['gpt2', 'cl100k_base'])
        
        if not text:
            return jsonify({"error": "No text provided"}), 400
        
        def tiktoken_decode_tokens(enc, token_ids):
            """Decode tiktoken tokens, showing hex for incomplete UTF-8 byte fragments.
            Keeps 1:1 correspondence with token_ids (no grouping)."""
            str_tokens = []
            byte_lengths = []
            for t in token_ids:
                raw = enc.decode_single_token_bytes(t)
                byte_lengths.append(len(raw))
                try:
                    str_tokens.append(raw.decode('utf-8'))
                except UnicodeDecodeError:
                    str_tokens.append("<0x" + raw.hex().upper() + ">")
            return str_tokens, byte_lengths
        
        results = {}
        
        for tok_id in requested:
            try:
                if tok_id == 'gpt2':
                    import tiktoken
                    enc = tiktoken.get_encoding("gpt2")
                    token_ids = enc.encode(text, allowed_special="all")
                    str_tokens, byte_lengths = tiktoken_decode_tokens(enc, token_ids)
                    results[tok_id] = {
                        "name": "GPT-2",
                        "vocabSize": enc.n_vocab,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "BPE (Byte-Pair Encoding)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'cl100k_base':
                    import tiktoken
                    enc = tiktoken.get_encoding("cl100k_base")
                    token_ids = enc.encode(text, allowed_special="all")
                    str_tokens, byte_lengths = tiktoken_decode_tokens(enc, token_ids)
                    results[tok_id] = {
                        "name": "GPT-4 (cl100k)",
                        "vocabSize": enc.n_vocab,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "BPE (Byte-Pair Encoding)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'o200k_base':
                    import tiktoken
                    enc = tiktoken.get_encoding("o200k_base")
                    token_ids = enc.encode(text, allowed_special="all")
                    str_tokens, byte_lengths = tiktoken_decode_tokens(enc, token_ids)
                    results[tok_id] = {
                        "name": "GPT-4o (o200k)",
                        "vocabSize": enc.n_vocab,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "BPE (Byte-Pair Encoding)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'llama3':
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        "philschmid/meta-llama-3-tokenizer",
                        use_fast=True,
                        token=os.environ.get("HF_TOKEN", None)
                    )
                    token_ids = tokenizer.encode(text, add_special_tokens=False)
                    str_tokens = [tokenizer.decode([t]) for t in token_ids]
                    byte_lengths = [len(s.encode('utf-8')) for s in str_tokens]
                    results[tok_id] = {
                        "name": "Llama 3",
                        "vocabSize": tokenizer.vocab_size,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "BPE (SentencePiece)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'gemma':
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        "pcuenq/gemma-tokenizer",
                        use_fast=True,
                        token=os.environ.get("HF_TOKEN", None)
                    )
                    token_ids = tokenizer.encode(text, add_special_tokens=False)
                    str_tokens = [tokenizer.decode([t]) for t in token_ids]
                    byte_lengths = [len(s.encode('utf-8')) for s in str_tokens]
                    results[tok_id] = {
                        "name": "Gemma",
                        "vocabSize": tokenizer.vocab_size,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "SentencePiece (Unigram)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'rinna':
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        "rinna/japanese-gpt2-medium",
                        use_fast=False,
                        token=os.environ.get("HF_TOKEN", None)
                    )
                    token_ids = tokenizer.encode(text, add_special_tokens=False)
                    str_tokens = [tokenizer.decode([t]) for t in token_ids]
                    byte_lengths = [len(s.encode('utf-8')) for s in str_tokens]
                    results[tok_id] = {
                        "name": "Rinna",
                        "vocabSize": tokenizer.vocab_size,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "SentencePiece (Unigram)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'calm2':
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        "cyberagent/calm2-7b",
                        use_fast=True,
                        token=os.environ.get("HF_TOKEN", None)
                    )
                    token_ids = tokenizer.encode(text, add_special_tokens=False)
                    str_tokens = [tokenizer.decode([t]) for t in token_ids]
                    byte_lengths = [len(s.encode('utf-8')) for s in str_tokens]
                    results[tok_id] = {
                        "name": "CyberAgent CALM2",
                        "vocabSize": tokenizer.vocab_size,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "BPE (SentencePiece)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                elif tok_id == 'stablelm_jp':
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        "stabilityai/japanese-stablelm-base-gamma-7b",
                        use_fast=True,
                        token=os.environ.get("HF_TOKEN", None)
                    )
                    token_ids = tokenizer.encode(text, add_special_tokens=False)
                    str_tokens = [tokenizer.decode([t]) for t in token_ids]
                    byte_lengths = [len(s.encode('utf-8')) for s in str_tokens]
                    results[tok_id] = {
                        "name": "StableLM JP",
                        "vocabSize": tokenizer.vocab_size,
                        "tokenCount": len(token_ids),
                        "tokens": str_tokens,
                        "tokenIds": token_ids,
                        "byteLengths": byte_lengths,
                        "algorithm": "SentencePiece (BPE+Unigram)",
                        "charsPerToken": round(len(text) / max(len(token_ids), 1), 2)
                    }
                    
                else:
                    results[tok_id] = {"error": f"Unknown tokenizer: {tok_id}"}
                    
            except Exception as tok_err:
                print(f"[SAE Service] Tokenizer {tok_id} error: {str(tok_err)}")
                results[tok_id] = {"error": str(tok_err)}
        
        return jsonify({
            "success": True,
            "text": text,
            "textLength": len(text),
            "textBytes": len(text.encode('utf-8')),
            "results": results
        })
        
    except Exception as e:
        print(f"[SAE Service] Tokenize compare error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


def extract_all_layer_activations(model, texts, position="last"):
    """
    Shared utility: extract residual stream activations at every layer for a list of texts.
    Returns dict with per-text, per-layer activations as numpy arrays.
    
    position: "last" = last token, "mean" = mean across all tokens
    """
    n_layers = model.cfg.n_layers
    resid_hooks = [f"blocks.{l}.hook_resid_post" for l in range(n_layers)]
    all_activations = []
    all_token_strs = []
    
    for text in texts:
        tokens = model.to_tokens(text)
        token_strs = model.to_str_tokens(text)
        
        with torch.no_grad():
            logits, cache = model.run_with_cache(tokens, names_filter=resid_hooks)
        
        layer_acts = {}
        for layer in range(n_layers):
            hook_key = f"blocks.{layer}.hook_resid_post"
            if hook_key not in cache:
                hook_key = f"blocks.{layer}.hook_resid_pre"
            if hook_key not in cache:
                available = [k for k in cache.keys() if f"blocks.{layer}" in k and "resid" in k]
                if available:
                    hook_key = available[0]
                else:
                    print(f"[extract_all_layer_activations] No resid hook found for layer {layer}, available keys: {[k for k in cache.keys() if f'blocks.{layer}' in k]}")
                    continue
            
            resid = cache[hook_key][0]
            if position == "last":
                act = resid[-1, :].float().cpu().numpy()
            elif position == "mean":
                act = resid.mean(dim=0).float().cpu().numpy()
            else:
                act = resid[-1, :].float().cpu().numpy()
            layer_acts[layer] = act
        
        all_activations.append(layer_acts)
        all_token_strs.append(token_strs)
    
    return all_activations, all_token_strs


def extract_final_logits(model, text):
    """Get final-layer logits and top-k predictions for a text."""
    tokens = model.to_tokens(text)
    with torch.no_grad():
        logits = model(tokens)
    
    last_logits = logits[0, -1, :].float()
    probs = torch.softmax(last_logits, dim=-1)
    entropy = -torch.sum(probs * torch.log(probs + 1e-10)).item()
    
    top_probs, top_indices = probs.topk(10)
    top_tokens = []
    for i in range(10):
        tok_id = top_indices[i].item()
        tok_str = model.tokenizer.decode([tok_id])
        top_tokens.append({
            "token": tok_str,
            "prob": round(top_probs[i].item(), 4),
            "id": tok_id
        })
    
    return {
        "topTokens": top_tokens,
        "entropy": round(entropy, 3),
        "confidence": round(top_probs[0].item(), 4)
    }


@app.route('/latent-language/analyze', methods=['POST'])
def latent_language_analyze():
    """
    Detect at which layers a model recognizes input language.
    Trains per-layer linear probes on Japanese vs English activations.
    """
    try:
        data = request.json
        japanese_texts = data.get('japaneseTexts', [])
        english_texts = data.get('englishTexts', [])
        model_id = data.get('model', 'gpt2-small')
        translation_pairs = data.get('translationPairs', [])
        
        if len(japanese_texts) < 2 or len(english_texts) < 2:
            return jsonify({"error": "Need at least 2 texts in each language"}), 400
        
        from transformer_lens import HookedTransformer
        from sklearn.linear_model import LogisticRegression
        from sklearn.model_selection import cross_val_score
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[LatentLang] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        all_texts = english_texts + japanese_texts
        labels = [0] * len(english_texts) + [1] * len(japanese_texts)
        
        print(f"[LatentLang] Extracting activations for {len(all_texts)} texts...")
        activations, _ = extract_all_layer_activations(_model, all_texts)
        
        n_layers = _model.cfg.n_layers
        y = np.array(labels)
        
        layers_data = []
        for layer in range(n_layers):
            X = np.array([act[layer] for act in activations])
            
            min_class_size = min(len(english_texts), len(japanese_texts))
            if min_class_size >= 3:
                cv_folds = min(5, min_class_size)
                clf = LogisticRegression(max_iter=1000, random_state=42)
                scores = cross_val_score(clf, X, y, cv=cv_folds, scoring='accuracy')
                accuracy = float(scores.mean())
            else:
                clf = LogisticRegression(max_iter=1000, random_state=42)
                clf.fit(X, y)
                accuracy = float(clf.score(X, y))
            
            clf_full = LogisticRegression(max_iter=1000, random_state=42)
            clf_full.fit(X, y)
            
            raw_weights = clf_full.coef_[0]
            weight_magnitudes = np.abs(raw_weights)
            top_neuron_indices = weight_magnitudes.argsort()[-10:][::-1].tolist()
            top_neuron_weights = [round(float(raw_weights[i]), 4) for i in top_neuron_indices]
            
            layers_data.append({
                "layer": layer,
                "accuracy": round(accuracy, 3),
                "topNeurons": top_neuron_indices,
                "topNeuronWeights": top_neuron_weights
            })
        
        similarity_data = []
        if translation_pairs:
            for pair in translation_pairs:
                en_text = pair.get("english", "")
                ja_text = pair.get("japanese", "")
                if not en_text or not ja_text:
                    continue
                
                pair_acts, _ = extract_all_layer_activations(_model, [en_text, ja_text])
                en_acts = pair_acts[0]
                ja_acts = pair_acts[1]
                
                pair_similarities = []
                for layer in range(n_layers):
                    en_vec = en_acts[layer]
                    ja_vec = ja_acts[layer]
                    cos_sim = float(np.dot(en_vec, ja_vec) / (np.linalg.norm(en_vec) * np.linalg.norm(ja_vec) + 1e-10))
                    pair_similarities.append({
                        "layer": layer,
                        "similarity": round(cos_sim, 4)
                    })
                
                similarity_data.append({
                    "english": en_text,
                    "japanese": ja_text,
                    "layers": pair_similarities
                })
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nEnglish": len(english_texts),
            "nJapanese": len(japanese_texts),
            "probeResults": layers_data,
            "translationSimilarity": similarity_data
        })
        
    except Exception as e:
        print(f"[LatentLang] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/cross-lingual/analyze', methods=['POST'])
def cross_lingual_analyze():
    """
    Compare activation patterns between parallel Japanese/English facts.
    Measures cosine similarity at each layer to find where representations converge.
    """
    try:
        data = request.json
        pairs = data.get('pairs', [])
        model_id = data.get('model', 'gpt2-small')
        
        if not pairs or len(pairs) < 1:
            return jsonify({"error": "Need at least 1 translation pair"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[CrossLingual] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        n_layers = _model.cfg.n_layers
        pair_results = []
        
        for pair in pairs:
            en_text = pair.get("english", "")
            ja_text = pair.get("japanese", "")
            label = pair.get("label", "")
            domain = pair.get("domain", "general")
            
            if not en_text or not ja_text:
                continue
            
            print(f"[CrossLingual] Analyzing pair: {label}")
            acts, token_strs = extract_all_layer_activations(_model, [en_text, ja_text])
            en_acts = acts[0]
            ja_acts = acts[1]
            
            en_logits = extract_final_logits(_model, en_text)
            ja_logits = extract_final_logits(_model, ja_text)
            
            layer_data = []
            for layer in range(n_layers):
                en_vec = en_acts[layer]
                ja_vec = ja_acts[layer]
                
                cos_sim = float(np.dot(en_vec, ja_vec) / (np.linalg.norm(en_vec) * np.linalg.norm(ja_vec) + 1e-10))
                
                l2_dist = float(np.linalg.norm(en_vec - ja_vec))
                
                layer_data.append({
                    "layer": layer,
                    "cosineSimilarity": round(cos_sim, 4),
                    "l2Distance": round(l2_dist, 2)
                })
            
            convergence_layer = -1
            max_sim = 0
            for ld in layer_data:
                if ld["cosineSimilarity"] > max_sim:
                    max_sim = ld["cosineSimilarity"]
                    convergence_layer = ld["layer"]
            
            pair_results.append({
                "english": en_text,
                "japanese": ja_text,
                "label": label,
                "domain": domain,
                "layers": layer_data,
                "convergenceLayer": convergence_layer,
                "maxSimilarity": round(max_sim, 4),
                "englishPrediction": en_logits,
                "japanesePrediction": ja_logits,
                "englishTokens": len(token_strs[0]),
                "japaneseTokens": len(token_strs[1])
            })
        
        avg_similarities = []
        for layer in range(n_layers):
            sims = [p["layers"][layer]["cosineSimilarity"] for p in pair_results if len(p["layers"]) > layer]
            avg_sim = float(np.mean(sims)) if sims else 0
            avg_similarities.append({
                "layer": layer,
                "avgSimilarity": round(avg_sim, 4)
            })
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nPairs": len(pair_results),
            "pairs": pair_results,
            "averageSimilarity": avg_similarities
        })
        
    except Exception as e:
        print(f"[CrossLingual] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/japanese-bench/analyze', methods=['POST'])
def japanese_bench_analyze():
    """
    Run domain-specific prompts in both Japanese and English,
    compare model confidence, entropy, and top predictions.
    """
    try:
        data = request.json
        prompts = data.get('prompts', [])
        model_id = data.get('model', 'gpt2-small')
        
        if not prompts:
            return jsonify({"error": "Need at least 1 prompt pair"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[JapaneseBench] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        n_layers = _model.cfg.n_layers
        W_U = _model.W_U
        b_U = _model.b_U if hasattr(_model, 'b_U') and _model.b_U is not None else None
        
        results = []
        
        for prompt_pair in prompts:
            en_text = prompt_pair.get("english", "")
            ja_text = prompt_pair.get("japanese", "")
            domain = prompt_pair.get("domain", "general")
            label = prompt_pair.get("label", "")
            
            if not en_text or not ja_text:
                continue
            
            print(f"[JapaneseBench] Analyzing: {label} ({domain})")
            
            en_logits = extract_final_logits(_model, en_text)
            ja_logits = extract_final_logits(_model, ja_text)
            
            acts, token_strs = extract_all_layer_activations(_model, [en_text, ja_text])
            
            en_lens = []
            ja_lens = []
            for layer in range(n_layers):
                for lang_idx, lens_list in [(0, en_lens), (1, ja_lens)]:
                    if layer not in acts[lang_idx]:
                        continue
                    resid = torch.tensor(acts[lang_idx][layer]).float()
                    if hasattr(_model, 'ln_final'):
                        resid_normed = _model.ln_final(resid.unsqueeze(0)).squeeze(0)
                    else:
                        resid_normed = resid
                    logits = resid_normed.float() @ W_U.float()
                    if b_U is not None:
                        logits = logits + b_U.float()
                    probs = torch.softmax(logits, dim=-1)
                    
                    top_p, top_i = probs.topk(5)
                    top_toks = []
                    for k in range(5):
                        tok_id = top_i[k].item()
                        tok_str = _model.tokenizer.decode([tok_id])
                        top_toks.append({
                            "token": tok_str,
                            "prob": round(top_p[k].item(), 4)
                        })
                    
                    entropy = -torch.sum(probs * torch.log(probs + 1e-10)).item()
                    lens_list.append({
                        "layer": layer,
                        "topTokens": top_toks,
                        "entropy": round(entropy, 3),
                        "confidence": round(top_p[0].item(), 4)
                    })
            
            confidence_gap = en_logits["confidence"] - ja_logits["confidence"]
            entropy_gap = ja_logits["entropy"] - en_logits["entropy"]
            
            results.append({
                "english": en_text,
                "japanese": ja_text,
                "domain": domain,
                "label": label,
                "englishPrediction": en_logits,
                "japanesePrediction": ja_logits,
                "englishLens": en_lens,
                "japaneseLens": ja_lens,
                "confidenceGap": round(confidence_gap, 4),
                "entropyGap": round(entropy_gap, 3),
                "englishTokenCount": len(token_strs[0]),
                "japaneseTokenCount": len(token_strs[1])
            })
        
        domain_summaries = {}
        for r in results:
            d = r["domain"]
            if d not in domain_summaries:
                domain_summaries[d] = {"confidenceGaps": [], "entropyGaps": [], "count": 0}
            domain_summaries[d]["confidenceGaps"].append(r["confidenceGap"])
            domain_summaries[d]["entropyGaps"].append(r["entropyGap"])
            domain_summaries[d]["count"] += 1
        
        for d in domain_summaries:
            domain_summaries[d]["avgConfidenceGap"] = round(float(np.mean(domain_summaries[d]["confidenceGaps"])), 4)
            domain_summaries[d]["avgEntropyGap"] = round(float(np.mean(domain_summaries[d]["entropyGaps"])), 3)
            del domain_summaries[d]["confidenceGaps"]
            del domain_summaries[d]["entropyGaps"]
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nPrompts": len(results),
            "results": results,
            "domainSummary": domain_summaries
        })
        
    except Exception as e:
        print(f"[JapaneseBench] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/concept-transfer/analyze', methods=['POST'])
def concept_transfer_analyze():
    """
    Causal cross-lingual knowledge transfer experiment.
    Tests whether patching English activations into a Japanese forward pass
    changes the model's predictions — revealing at which layers concepts
    are language-agnostic vs language-specific.
    """
    try:
        data = request.json
        pairs = data.get('pairs', [])
        model_id = data.get('model', 'gpt2-small')
        
        if not pairs or len(pairs) < 1:
            return jsonify({"error": "Need at least 1 prompt pair"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[ConceptTransfer] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        model = _model
        n_layers = model.cfg.n_layers
        pair_results = []
        
        for pair in pairs:
            en_text = pair.get("english", "")
            ja_text = pair.get("japanese", "")
            target = pair.get("target", "")
            label = pair.get("label", en_text[:40])
            
            if not en_text or not ja_text:
                continue
            
            print(f"[ConceptTransfer] Analyzing: {label}")
            
            en_tokens = model.to_tokens(en_text)
            ja_tokens = model.to_tokens(ja_text)
            
            with torch.no_grad():
                en_logits, en_cache = model.run_with_cache(en_tokens)
                ja_logits, ja_cache = model.run_with_cache(ja_tokens)
            
            en_last_logits = en_logits[0, -1, :]
            ja_last_logits = ja_logits[0, -1, :]
            en_probs = torch.softmax(en_last_logits, dim=-1)
            ja_probs = torch.softmax(ja_last_logits, dim=-1)
            
            target_token_id = None
            if target:
                target_ids = model.to_tokens(target, prepend_bos=False)[0]
                target_token_id = target_ids[0].item()
            else:
                target_token_id = en_probs.argmax().item()
            
            target_str = model.tokenizer.decode([target_token_id])
            en_target_prob = en_probs[target_token_id].item()
            ja_target_prob = ja_probs[target_token_id].item()
            
            en_top5 = []
            top_p, top_i = en_probs.topk(5)
            for k in range(5):
                en_top5.append({"token": model.tokenizer.decode([top_i[k].item()]), "prob": round(top_p[k].item(), 4)})
            
            ja_top5 = []
            top_p, top_i = ja_probs.topk(5)
            for k in range(5):
                ja_top5.append({"token": model.tokenizer.decode([top_i[k].item()]), "prob": round(top_p[k].item(), 4)})
            
            layer_results = []
            for layer in range(n_layers):
                en_resid = en_cache[f"blocks.{layer}.hook_resid_post"]
                
                def patch_hook(activation, hook, source_act=en_resid):
                    patched = activation.clone()
                    patched[0, -1, :] = source_act[0, -1, :]
                    return patched
                
                hook_name = f"blocks.{layer}.hook_resid_post"
                model.reset_hooks()
                with torch.no_grad():
                    patched_logits = model.run_with_hooks(
                        ja_tokens,
                        fwd_hooks=[(hook_name, patch_hook)],
                        reset_hooks_end=True
                    )
                
                patched_probs = torch.softmax(patched_logits[0, -1, :], dim=-1)
                patched_target_prob = patched_probs[target_token_id].item()
                
                transfer_effect = patched_target_prob - ja_target_prob
                
                patched_top3 = []
                top_p, top_i = patched_probs.topk(3)
                for k in range(3):
                    patched_top3.append({"token": model.tokenizer.decode([top_i[k].item()]), "prob": round(top_p[k].item(), 4)})
                
                en_vec = en_cache[f"blocks.{layer}.hook_resid_post"][0, -1, :].detach().float().cpu().numpy()
                ja_vec = ja_cache[f"blocks.{layer}.hook_resid_post"][0, -1, :].detach().float().cpu().numpy()
                cos_sim = float(np.dot(en_vec, ja_vec) / (np.linalg.norm(en_vec) * np.linalg.norm(ja_vec) + 1e-10))
                
                layer_results.append({
                    "layer": layer,
                    "transferEffect": round(transfer_effect, 6),
                    "patchedTargetProb": round(patched_target_prob, 6),
                    "cosineSimilarity": round(cos_sim, 4),
                    "patchedTop3": patched_top3
                })
            
            peak_layer = max(layer_results, key=lambda x: x["transferEffect"])
            
            del en_cache, ja_cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            pair_results.append({
                "english": en_text,
                "japanese": ja_text,
                "target": target_str.strip(),
                "targetTokenId": target_token_id,
                "label": label,
                "englishTargetProb": round(en_target_prob, 6),
                "japaneseTargetProb": round(ja_target_prob, 6),
                "englishTop5": en_top5,
                "japaneseTop5": ja_top5,
                "layers": layer_results,
                "peakTransferLayer": peak_layer["layer"],
                "peakTransferEffect": round(peak_layer["transferEffect"], 6),
                "englishTokenCount": en_tokens.shape[1],
                "japaneseTokenCount": ja_tokens.shape[1]
            })
        
        avg_transfer = []
        for layer in range(n_layers):
            effects = [p["layers"][layer]["transferEffect"] for p in pair_results if len(p["layers"]) > layer]
            sims = [p["layers"][layer]["cosineSimilarity"] for p in pair_results if len(p["layers"]) > layer]
            avg_transfer.append({
                "layer": layer,
                "avgTransferEffect": round(float(np.mean(effects)), 6) if effects else 0,
                "avgCosineSimilarity": round(float(np.mean(sims)), 4) if sims else 0
            })
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nPairs": len(pair_results),
            "pairs": pair_results,
            "averageTransfer": avg_transfer
        })
        
    except Exception as e:
        print(f"[ConceptTransfer] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/token-trajectory/analyze', methods=['POST'])
def token_trajectory_analyze():
    """
    Token Trajectory: compare how English vs Japanese representations
    evolve through the model's layers. Shows where fragmented Japanese
    byte-tokens converge into coherent representations (or don't)
    relative to a single English token.
    """
    try:
        data = request.json
        concepts = data.get('concepts', [])
        model_id = data.get('model', 'gpt2-small')
        aggregation = data.get('aggregation', 'mean')
        
        if not concepts or len(concepts) < 1:
            return jsonify({"error": "Need at least 1 concept pair"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[TokenTrajectory] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        model = _model
        n_layers = model.cfg.n_layers
        W_U = model.W_U
        concept_results = []
        
        for concept in concepts:
            en_text = concept.get("english", "")
            ja_text = concept.get("japanese", "")
            label = concept.get("label", en_text[:30])
            
            if not en_text or not ja_text:
                continue
            
            print(f"[TokenTrajectory] Analyzing: {label} ({en_text} / {ja_text})")
            
            en_tokens = model.to_tokens(en_text)
            ja_tokens = model.to_tokens(ja_text)
            
            en_token_count = en_tokens.shape[1]
            ja_token_count = ja_tokens.shape[1]
            
            en_token_strs = [model.tokenizer.decode([en_tokens[0, t].item()]) for t in range(en_token_count)]
            ja_token_strs = [model.tokenizer.decode([ja_tokens[0, t].item()]) for t in range(ja_token_count)]
            
            resid_hooks = [f"blocks.{l}.hook_resid_post" for l in range(n_layers)]
            with torch.no_grad():
                _, en_cache = model.run_with_cache(en_tokens, names_filter=resid_hooks)
                _, ja_cache = model.run_with_cache(ja_tokens, names_filter=resid_hooks)
            
            layer_results = []
            for layer in range(n_layers):
                en_resid = en_cache[f"blocks.{layer}.hook_resid_post"]
                ja_resid = ja_cache[f"blocks.{layer}.hook_resid_post"]
                
                en_vec = en_resid[0, -1, :]
                
                if aggregation == "last":
                    ja_vec = ja_resid[0, -1, :]
                else:
                    ja_content_start = 1
                    ja_vec = ja_resid[0, ja_content_start:, :].mean(dim=0)
                
                cos_sim = torch.nn.functional.cosine_similarity(
                    en_vec.unsqueeze(0), ja_vec.unsqueeze(0)
                ).item()
                
                en_logits = en_vec @ W_U
                en_probs = torch.softmax(en_logits, dim=-1)
                en_top_p, en_top_i = en_probs.topk(5)
                en_top5 = [{"token": model.tokenizer.decode([en_top_i[k].item()]), "prob": round(en_top_p[k].item(), 4)} for k in range(5)]
                
                ja_logits = ja_vec @ W_U
                ja_probs = torch.softmax(ja_logits, dim=-1)
                ja_top_p, ja_top_i = ja_probs.topk(5)
                ja_top5 = [{"token": model.tokenizer.decode([ja_top_i[k].item()]), "prob": round(ja_top_p[k].item(), 4)} for k in range(5)]
                
                shared_tokens = set()
                en_top20_i = en_probs.topk(20).indices.tolist()
                ja_top20_i = ja_probs.topk(20).indices.tolist()
                overlap_count = len(set(en_top20_i) & set(ja_top20_i))
                
                layer_results.append({
                    "layer": layer,
                    "cosineSimilarity": round(cos_sim, 4),
                    "englishTop5": en_top5,
                    "japaneseTop5": ja_top5,
                    "vocabOverlap20": overlap_count
                })
            
            first_sim = layer_results[0]["cosineSimilarity"]
            last_sim = layer_results[-1]["cosineSimilarity"]
            peak_sim = max(layer_results, key=lambda x: x["cosineSimilarity"])
            min_sim = min(layer_results, key=lambda x: x["cosineSimilarity"])
            
            convergence = last_sim - first_sim
            
            del en_cache, ja_cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            concept_results.append({
                "english": en_text,
                "japanese": ja_text,
                "label": label,
                "englishTokens": en_token_strs,
                "japaneseTokens": ja_token_strs,
                "englishTokenCount": en_token_count,
                "japaneseTokenCount": ja_token_count,
                "layers": layer_results,
                "convergence": round(convergence, 4),
                "peakSimilarityLayer": peak_sim["layer"],
                "peakSimilarity": round(peak_sim["cosineSimilarity"], 4),
                "minSimilarityLayer": min_sim["layer"],
                "minSimilarity": round(min_sim["cosineSimilarity"], 4)
            })
        
        avg_similarity = []
        for layer in range(n_layers):
            sims = [c["layers"][layer]["cosineSimilarity"] for c in concept_results if len(c["layers"]) > layer]
            overlaps = [c["layers"][layer]["vocabOverlap20"] for c in concept_results if len(c["layers"]) > layer]
            avg_similarity.append({
                "layer": layer,
                "avgSimilarity": round(float(np.mean(sims)), 4) if sims else 0,
                "avgVocabOverlap": round(float(np.mean(overlaps)), 1) if overlaps else 0
            })
        
        avg_convergence = round(float(np.mean([c["convergence"] for c in concept_results])), 4) if concept_results else 0
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nConcepts": len(concept_results),
            "aggregation": aggregation,
            "concepts": concept_results,
            "averageSimilarity": avg_similarity,
            "averageConvergence": avg_convergence
        })
        
    except Exception as e:
        print(f"[TokenTrajectory] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Feature Universality Endpoint
# ============================================

@app.route('/feature-universality/analyze', methods=['POST'])
def feature_universality_analyze():
    """
    Cross-lingual Feature Universality: find SAE features that activate
    for both English and Japanese versions of the same concept.
    Universal features = language-agnostic concept encoding.
    """
    try:
        data = request.json
        pairs = data.get('pairs', [])
        model_id = data.get('model', 'gpt2-small')
        top_k = data.get('topK', 20)
        
        if not pairs or len(pairs) < 1:
            return jsonify({"error": "Need at least 1 concept pair"}), 400
        
        from transformer_lens import HookedTransformer
        from sae_lens import SAE
        
        model, sae = get_model_and_sae(model_id)
        if sae is None:
            return jsonify({"error": f"SAE not available for model {model_id}"}), 503
        hook_name = get_hook_name(model_id)
        n_features = sae.cfg.d_sae
        
        pair_results = []
        en_activation_matrix = []
        jp_activation_matrix = []
        
        for pair in pairs:
            en_text = pair.get('english', '')
            jp_text = pair.get('japanese', '')
            label = pair.get('label', '')
            
            if not en_text or not jp_text:
                continue
            
            en_tokens = model.to_tokens(en_text)
            jp_tokens = model.to_tokens(jp_text)
            en_token_strs = model.to_str_tokens(en_text)
            jp_token_strs = model.to_str_tokens(jp_text)
            
            with torch.no_grad():
                _, en_cache = model.run_with_cache(en_tokens, names_filter=[hook_name])
                _, jp_cache = model.run_with_cache(jp_tokens, names_filter=[hook_name])
            
            en_acts = safe_get_cache(en_cache, hook_name)
            jp_acts = safe_get_cache(jp_cache, hook_name)
            
            en_feature_acts = sae.encode(en_acts)
            jp_feature_acts = sae.encode(jp_acts)
            
            en_max_acts = en_feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
            jp_max_acts = jp_feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
            
            del en_cache, jp_cache, en_acts, jp_acts, en_feature_acts, jp_feature_acts
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            en_activation_matrix.append(en_max_acts)
            jp_activation_matrix.append(jp_max_acts)
            
            en_top_indices = np.argsort(en_max_acts)[-top_k:][::-1]
            jp_top_indices = np.argsort(jp_max_acts)[-top_k:][::-1]
            
            en_top_set = set(en_top_indices.tolist())
            jp_top_set = set(jp_top_indices.tolist())
            shared_features = en_top_set & jp_top_set
            en_only = en_top_set - jp_top_set
            jp_only = jp_top_set - en_top_set
            
            both_active = (en_max_acts > 0) & (jp_max_acts > 0)
            cosine_sim = 0.0
            en_norm = np.linalg.norm(en_max_acts)
            jp_norm = np.linalg.norm(jp_max_acts)
            if en_norm > 0 and jp_norm > 0:
                cosine_sim = float(np.dot(en_max_acts, jp_max_acts) / (en_norm * jp_norm))
            
            shared_detail = []
            for fi in sorted(shared_features):
                shared_detail.append({
                    "feature": int(fi),
                    "englishActivation": round(float(en_max_acts[fi]), 4),
                    "japaneseActivation": round(float(jp_max_acts[fi]), 4),
                    "ratio": round(float(min(en_max_acts[fi], jp_max_acts[fi]) / max(en_max_acts[fi], jp_max_acts[fi])), 4) if max(en_max_acts[fi], jp_max_acts[fi]) > 0 else 0
                })
            shared_detail.sort(key=lambda x: min(x["englishActivation"], x["japaneseActivation"]), reverse=True)
            
            pair_results.append({
                "english": en_text,
                "japanese": jp_text,
                "label": label,
                "englishTokens": list(en_token_strs),
                "japaneseTokens": list(jp_token_strs),
                "cosineSimilarity": round(cosine_sim, 4),
                "sharedFeatureCount": len(shared_features),
                "englishOnlyCount": len(en_only),
                "japaneseOnlyCount": len(jp_only),
                "totalActiveEnglish": int(np.sum(en_max_acts > 0)),
                "totalActiveJapanese": int(np.sum(jp_max_acts > 0)),
                "sharedFeatures": shared_detail[:top_k],
                "englishOnlyFeatures": [{"feature": int(fi), "activation": round(float(en_max_acts[fi]), 4)} for fi in sorted(en_only, key=lambda x: en_max_acts[x], reverse=True)][:10],
                "japaneseOnlyFeatures": [{"feature": int(fi), "activation": round(float(jp_max_acts[fi]), 4)} for fi in sorted(jp_only, key=lambda x: jp_max_acts[x], reverse=True)][:10]
            })
        
        if len(en_activation_matrix) < 1:
            return jsonify({"error": "No valid pairs processed"}), 400
        
        en_mat = np.array(en_activation_matrix)
        jp_mat = np.array(jp_activation_matrix)
        
        universal_features = []
        n_pairs = len(en_activation_matrix)
        
        for fi in range(n_features):
            en_col = en_mat[:, fi]
            jp_col = jp_mat[:, fi]
            
            en_active = np.sum(en_col > 0)
            jp_active = np.sum(jp_col > 0)
            both_active_count = np.sum((en_col > 0) & (jp_col > 0))
            
            if both_active_count >= max(1, n_pairs * 0.5):
                en_mean = float(np.mean(en_col[en_col > 0])) if en_active > 0 else 0
                jp_mean = float(np.mean(jp_col[jp_col > 0])) if jp_active > 0 else 0
                
                correlation = 0.0
                if n_pairs > 1 and np.std(en_col) > 0 and np.std(jp_col) > 0:
                    correlation = float(np.corrcoef(en_col, jp_col)[0, 1])
                
                universal_features.append({
                    "feature": int(fi),
                    "bothActiveCount": int(both_active_count),
                    "englishMeanActivation": round(en_mean, 4),
                    "japaneseMeanActivation": round(jp_mean, 4),
                    "correlation": round(correlation, 4),
                    "universalityScore": round(float(both_active_count / n_pairs) * (1 + abs(correlation)) / 2, 4)
                })
        
        universal_features.sort(key=lambda x: x["universalityScore"], reverse=True)
        
        en_specific = []
        jp_specific = []
        for fi in range(n_features):
            en_col = en_mat[:, fi]
            jp_col = jp_mat[:, fi]
            en_active = np.sum(en_col > 0)
            jp_active = np.sum(jp_col > 0)
            
            if en_active >= max(1, n_pairs * 0.5) and jp_active == 0:
                en_specific.append({
                    "feature": int(fi),
                    "activeCount": int(en_active),
                    "meanActivation": round(float(np.mean(en_col[en_col > 0])), 4)
                })
            elif jp_active >= max(1, n_pairs * 0.5) and en_active == 0:
                jp_specific.append({
                    "feature": int(fi),
                    "activeCount": int(jp_active),
                    "meanActivation": round(float(np.mean(jp_col[jp_col > 0])), 4)
                })
        
        en_specific.sort(key=lambda x: x["meanActivation"], reverse=True)
        jp_specific.sort(key=lambda x: x["meanActivation"], reverse=True)
        
        avg_cosine = round(float(np.mean([p["cosineSimilarity"] for p in pair_results])), 4)
        avg_shared = round(float(np.mean([p["sharedFeatureCount"] for p in pair_results])), 1)
        
        return jsonify({
            "model": model_id,
            "nFeatures": int(n_features),
            "nPairs": len(pair_results),
            "saeHook": hook_name,
            "pairs": pair_results,
            "universalFeatures": universal_features[:50],
            "englishSpecificFeatures": en_specific[:30],
            "japaneseSpecificFeatures": jp_specific[:30],
            "averageCosineSimilarity": avg_cosine,
            "averageSharedFeatures": avg_shared,
            "totalUniversalFound": len(universal_features),
            "totalEnglishSpecific": len(en_specific),
            "totalJapaneseSpecific": len(jp_specific)
        })
        
    except Exception as e:
        print(f"[FeatureUniversality] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


# ============================================
# Tokenization Tax Endpoint
# ============================================

@app.route('/tokenization-tax/analyze', methods=['POST'])
def tokenization_tax_analyze():
    """
    Tokenization Tax: quantify the representation cost of Japanese
    multi-byte tokenization vs English single-token encoding.
    Measures per-layer prediction confidence, entropy, and how many
    extra layers Japanese needs to reach equivalent representation quality.
    """
    try:
        data = request.json
        pairs = data.get('pairs', [])
        model_id = data.get('model', 'gpt2-small')
        
        if not pairs or len(pairs) < 1:
            return jsonify({"error": "Need at least 1 concept pair"}), 400
        
        from transformer_lens import HookedTransformer
        
        global _model, _saes, _current_model_id
        base_model = get_base_model_id(model_id)
        with _model_swap_lock:
            if _model is None or get_base_model_id(_current_model_id) != base_model:
                if _model is not None:
                    _cleanup_model_vram()
                print(f"[TokenizationTax] Loading model: {base_model}")
                _model = _load_hooked_transformer(base_model)
                _saes = {}
                _current_model_id = model_id
        
        model = _model
        n_layers = model.cfg.n_layers
        
        W_U = model.W_U.detach().float().cpu()
        b_U = model.b_U.detach().float().cpu() if model.b_U is not None else None
        
        pair_results = []
        
        for pair in pairs:
            en_text = pair.get('english', '')
            jp_text = pair.get('japanese', '')
            label = pair.get('label', '')
            
            if not en_text or not jp_text:
                continue
            
            en_tokens = model.to_tokens(en_text)
            jp_tokens = model.to_tokens(jp_text)
            en_token_strs = model.to_str_tokens(en_text)
            jp_token_strs = model.to_str_tokens(jp_text)
            
            resid_hooks = [f"blocks.{l}.hook_resid_post" for l in range(n_layers)]
            with torch.no_grad():
                _, en_cache = model.run_with_cache(en_tokens, names_filter=resid_hooks)
                _, jp_cache = model.run_with_cache(jp_tokens, names_filter=resid_hooks)
            
            en_layers = []
            jp_layers = []
            
            for layer in range(n_layers):
                en_resid = en_cache[f"blocks.{layer}.hook_resid_post"][0, -1].detach().float().cpu()
                jp_resid = jp_cache[f"blocks.{layer}.hook_resid_post"][0, -1].detach().float().cpu()
                
                en_logits = en_resid @ W_U
                jp_logits = jp_resid @ W_U
                if b_U is not None:
                    en_logits = en_logits + b_U
                    jp_logits = jp_logits + b_U
                
                en_probs = torch.softmax(en_logits, dim=-1)
                jp_probs = torch.softmax(jp_logits, dim=-1)
                
                en_entropy = -torch.sum(en_probs * torch.log(en_probs + 1e-10)).item()
                jp_entropy = -torch.sum(jp_probs * torch.log(jp_probs + 1e-10)).item()
                
                en_top1_prob = en_probs.max().item()
                jp_top1_prob = jp_probs.max().item()
                
                en_top5_vals, en_top5_idx = torch.topk(en_probs, 5)
                jp_top5_vals, jp_top5_idx = torch.topk(jp_probs, 5)
                
                en_top5 = [{"token": model.tokenizer.decode([idx.item()]), "prob": round(val.item(), 6)} for val, idx in zip(en_top5_vals, en_top5_idx)]
                jp_top5 = [{"token": model.tokenizer.decode([idx.item()]), "prob": round(val.item(), 6)} for val, idx in zip(jp_top5_vals, jp_top5_idx)]
                
                en_layers.append({
                    "layer": layer,
                    "entropy": round(en_entropy, 4),
                    "top1Prob": round(en_top1_prob, 6),
                    "top5Cumulative": round(en_top5_vals.sum().item(), 6),
                    "topPredictions": en_top5
                })
                
                jp_layers.append({
                    "layer": layer,
                    "entropy": round(jp_entropy, 4),
                    "top1Prob": round(jp_top1_prob, 6),
                    "top5Cumulative": round(jp_top5_vals.sum().item(), 6),
                    "topPredictions": jp_top5
                })
            
            en_final_entropy = en_layers[-1]["entropy"]
            jp_final_entropy = jp_layers[-1]["entropy"]
            
            en_half_confidence_layer = n_layers - 1
            jp_half_confidence_layer = n_layers - 1
            final_en_top1 = en_layers[-1]["top1Prob"]
            final_jp_top1 = jp_layers[-1]["top1Prob"]
            
            for l in en_layers:
                if l["top1Prob"] >= final_en_top1 * 0.5:
                    en_half_confidence_layer = l["layer"]
                    break
            for l in jp_layers:
                if l["top1Prob"] >= final_jp_top1 * 0.5:
                    jp_half_confidence_layer = l["layer"]
                    break
            
            confidence_tax = jp_half_confidence_layer - en_half_confidence_layer
            
            entropy_ratios = []
            for i in range(n_layers):
                if en_layers[i]["entropy"] > 0:
                    entropy_ratios.append(jp_layers[i]["entropy"] / en_layers[i]["entropy"])
                else:
                    entropy_ratios.append(1.0)
            avg_entropy_ratio = float(np.mean(entropy_ratios))
            
            del en_cache, jp_cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            pair_results.append({
                "english": en_text,
                "japanese": jp_text,
                "label": label,
                "englishTokenCount": len(en_token_strs),
                "japaneseTokenCount": len(jp_token_strs),
                "tokenCountRatio": round(len(jp_token_strs) / max(len(en_token_strs), 1), 2),
                "englishTokens": list(en_token_strs),
                "japaneseTokens": list(jp_token_strs),
                "englishLayers": en_layers,
                "japaneseLayers": jp_layers,
                "englishHalfConfidenceLayer": en_half_confidence_layer,
                "japaneseHalfConfidenceLayer": jp_half_confidence_layer,
                "confidenceTax": confidence_tax,
                "averageEntropyRatio": round(avg_entropy_ratio, 4),
                "englishFinalEntropy": round(en_final_entropy, 4),
                "japaneseFinalEntropy": round(jp_final_entropy, 4)
            })
        
        if not pair_results:
            return jsonify({"error": "No valid pairs processed"}), 400
        
        avg_layers = []
        for layer in range(n_layers):
            en_entropies = [p["englishLayers"][layer]["entropy"] for p in pair_results]
            jp_entropies = [p["japaneseLayers"][layer]["entropy"] for p in pair_results]
            en_top1s = [p["englishLayers"][layer]["top1Prob"] for p in pair_results]
            jp_top1s = [p["japaneseLayers"][layer]["top1Prob"] for p in pair_results]
            
            avg_layers.append({
                "layer": layer,
                "avgEnglishEntropy": round(float(np.mean(en_entropies)), 4),
                "avgJapaneseEntropy": round(float(np.mean(jp_entropies)), 4),
                "avgEnglishTop1": round(float(np.mean(en_top1s)), 6),
                "avgJapaneseTop1": round(float(np.mean(jp_top1s)), 6),
                "entropyGap": round(float(np.mean(jp_entropies)) - float(np.mean(en_entropies)), 4),
                "confidenceGap": round(float(np.mean(en_top1s)) - float(np.mean(jp_top1s)), 6)
            })
        
        avg_confidence_tax = round(float(np.mean([p["confidenceTax"] for p in pair_results])), 2)
        avg_token_ratio = round(float(np.mean([p["tokenCountRatio"] for p in pair_results])), 2)
        avg_entropy_ratio_overall = round(float(np.mean([p["averageEntropyRatio"] for p in pair_results])), 4)
        
        return jsonify({
            "model": model_id,
            "nLayers": n_layers,
            "nPairs": len(pair_results),
            "pairs": pair_results,
            "averageLayers": avg_layers,
            "averageConfidenceTax": avg_confidence_tax,
            "averageTokenCountRatio": avg_token_ratio,
            "averageEntropyRatio": avg_entropy_ratio_overall
        })
        
    except Exception as e:
        print(f"[TokenizationTax] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/model-merge/analyze', methods=['POST'])
def model_merge_analyze():
    """
    Analyze what happens to SAE features during model merging.
    Simulates interpolation between two domain activation vectors.
    """
    try:
        data = request.json
        domain_a_texts = data.get('domainA', [])
        domain_b_texts = data.get('domainB', [])
        model_id = data.get('model', 'gpt2-small')
        alpha_steps = data.get('alphaSteps', [0.0, 0.25, 0.5, 0.75, 1.0])

        if not domain_a_texts or not domain_b_texts:
            return jsonify({"error": "Both domainA and domainB texts are required"}), 400

        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        n_features = sae.W_dec.shape[0] if hasattr(sae, 'W_dec') else sae.W_enc.shape[1]

        domain_a_vectors = []
        domain_a_info = []
        domain_b_vectors = []
        domain_b_info = []

        all_items = [(item, 'a') for item in domain_a_texts] + [(item, 'b') for item in domain_b_texts]
        with torch.no_grad():
            for item, domain_tag in all_items:
                tokens = model.to_tokens(item['text'])
                n_tok = tokens.shape[-1]
                _, cache = model.run_with_cache(tokens)
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
                max_pooled = feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
                n_active = int(np.sum(max_pooled > 0.1))
                info = {
                    "text": item['text'],
                    "label": item.get('label', ''),
                    "nTokens": int(n_tok),
                    "nActiveFeatures": n_active
                }
                if domain_tag == 'a':
                    domain_a_vectors.append(max_pooled)
                    domain_a_info.append(info)
                else:
                    domain_b_vectors.append(max_pooled)
                    domain_b_info.append(info)
                del cache, activations, feature_acts
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        domain_a_mean = np.mean(np.stack(domain_a_vectors), axis=0)
        domain_b_mean = np.mean(np.stack(domain_b_vectors), axis=0)

        merged_vectors = []
        for alpha in alpha_steps:
            merged = alpha * domain_a_mean + (1.0 - alpha) * domain_b_mean
            merged_vectors.append(merged)

        merged_matrix = np.stack(merged_vectors)
        per_feature_std = np.std(merged_matrix, axis=0)
        per_feature_max = np.max(merged_matrix, axis=0)
        stability = 1.0 - (per_feature_std / (per_feature_max + 1e-8))

        mean_across_alphas = np.mean(merged_matrix, axis=0)

        stable_mask = (stability > 0.8) & (mean_across_alphas > 0.1)
        stable_indices = np.where(stable_mask)[0]
        stable_sorted = stable_indices[np.argsort(-stability[stable_indices])][:30]

        fragile_mask = (stability < 0.3) & (per_feature_max > 0.1)
        fragile_indices = np.where(fragile_mask)[0]
        fragile_sorted = fragile_indices[np.argsort(stability[fragile_indices])][:30]

        domain_a_exclusive_mask = (domain_a_mean > 0.1) & (domain_b_mean <= 0.1)
        domain_a_excl_indices = np.where(domain_a_exclusive_mask)[0]
        domain_a_excl_sorted = domain_a_excl_indices[np.argsort(-domain_a_mean[domain_a_excl_indices])][:20]

        domain_b_exclusive_mask = (domain_b_mean > 0.1) & (domain_a_mean <= 0.1)
        domain_b_excl_indices = np.where(domain_b_exclusive_mask)[0]
        domain_b_excl_sorted = domain_b_excl_indices[np.argsort(-domain_b_mean[domain_b_excl_indices])][:20]

        active_mask = per_feature_max > 0.05
        if np.any(active_mask):
            merge_compatibility = float(np.mean(stability[active_mask]))
        else:
            merge_compatibility = 0.0

        def cosine_sim(a, b):
            dot = np.dot(a, b)
            norm_a = np.linalg.norm(a)
            norm_b = np.linalg.norm(b)
            if norm_a < 1e-10 or norm_b < 1e-10:
                return 0.0
            return float(dot / (norm_a * norm_b))

        alpha_results = []
        for i, alpha in enumerate(alpha_steps):
            merged = merged_vectors[i]
            cos_a = cosine_sim(merged, domain_a_mean)
            cos_b = cosine_sim(merged, domain_b_mean)
            active_count = int(np.sum(merged > 0.1))
            alpha_results.append({
                "alpha": round(float(alpha), 4),
                "cosineToDomainA": round(cos_a, 4),
                "cosineToDomainB": round(cos_b, 4),
                "activeFeatureCount": active_count
            })

        mid_alpha_idx = len(alpha_steps) // 2
        mid_merged = merged_vectors[mid_alpha_idx]

        stable_features_list = []
        for idx in stable_sorted:
            idx = int(idx)
            stable_features_list.append({
                "feature": idx,
                "stability": round(float(stability[idx]), 4),
                "domainAMean": round(float(domain_a_mean[idx]), 4),
                "domainBMean": round(float(domain_b_mean[idx]), 4),
                "mergedMean": round(float(mid_merged[idx]), 4)
            })

        fragile_features_list = []
        for idx in fragile_sorted:
            idx = int(idx)
            fragile_features_list.append({
                "feature": idx,
                "stability": round(float(stability[idx]), 4),
                "domainAMean": round(float(domain_a_mean[idx]), 4),
                "domainBMean": round(float(domain_b_mean[idx]), 4),
                "variance": round(float(per_feature_std[idx] ** 2), 4)
            })

        domain_a_excl_list = []
        for idx in domain_a_excl_sorted:
            idx = int(idx)
            domain_a_excl_list.append({
                "feature": idx,
                "activation": round(float(domain_a_mean[idx]), 4)
            })

        domain_b_excl_list = []
        for idx in domain_b_excl_sorted:
            idx = int(idx)
            domain_b_excl_list.append({
                "feature": idx,
                "activation": round(float(domain_b_mean[idx]), 4)
            })

        domain_a_active = int(np.sum(domain_a_mean > 0.1))
        domain_b_active = int(np.sum(domain_b_mean > 0.1))

        return jsonify({
            "model": model_id,
            "nFeatures": int(n_features),
            "saeHook": hook_name,
            "nDomainA": len(domain_a_texts),
            "nDomainB": len(domain_b_texts),
            "domainATexts": domain_a_info,
            "domainBTexts": domain_b_info,
            "alphaSteps": alpha_results,
            "stableFeatures": stable_features_list,
            "fragileFeatures": fragile_features_list,
            "domainAExclusive": domain_a_excl_list,
            "domainBExclusive": domain_b_excl_list,
            "mergeCompatibility": round(float(merge_compatibility), 4),
            "totalStable": int(np.sum(stable_mask)),
            "totalFragile": int(np.sum(fragile_mask)),
            "totalDomainAExclusive": int(np.sum(domain_a_exclusive_mask)),
            "totalDomainBExclusive": int(np.sum(domain_b_exclusive_mask)),
            "domainASummary": {
                "totalActive": domain_a_active,
                "meanActivation": round(float(np.mean(domain_a_mean[domain_a_mean > 0.1])) if domain_a_active > 0 else 0.0, 4)
            },
            "domainBSummary": {
                "totalActive": domain_b_active,
                "meanActivation": round(float(np.mean(domain_b_mean[domain_b_mean > 0.1])) if domain_b_active > 0 else 0.0, 4)
            }
        })

    except Exception as e:
        print(f"[ModelMerge] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/financial-circuits/analyze', methods=['POST'])
def financial_circuits_analyze():
    """
    Map which SAE features encode financial concepts across English and Japanese.
    """
    try:
        data = request.json
        financial_pairs = data.get('financialPairs', [])
        baseline_pairs = data.get('baselinePairs', [])
        model_id = data.get('model', 'gpt2-small')
        top_k = data.get('topK', 20)

        if not financial_pairs:
            return jsonify({"error": "financialPairs are required"}), 400
        if not baseline_pairs:
            return jsonify({"error": "baselinePairs are required"}), 400

        model, sae = get_model_and_sae(model_id)
        hook_name = get_hook_name(model_id)
        n_features = sae.W_dec.shape[0] if hasattr(sae, 'W_dec') else sae.W_enc.shape[1]

        def cosine_sim(a, b):
            dot = np.dot(a, b)
            norm_a = np.linalg.norm(a)
            norm_b = np.linalg.norm(b)
            if norm_a < 1e-10 or norm_b < 1e-10:
                return 0.0
            return float(dot / (norm_a * norm_b))

        all_texts = []
        all_tags = []
        for i, pair in enumerate(financial_pairs):
            all_texts.append(pair['english'])
            all_tags.append(('fin_en', i))
            all_texts.append(pair['japanese'])
            all_tags.append(('fin_jp', i))
        for i, pair in enumerate(baseline_pairs):
            all_texts.append(pair['english'])
            all_tags.append(('base_en', i))
            all_texts.append(pair['japanese'])
            all_tags.append(('base_jp', i))

        all_vectors = {}
        all_str_tokens = {}
        with torch.no_grad():
            for text, tag in zip(all_texts, all_tags):
                tokens = model.to_tokens(text)
                str_tokens = model.to_str_tokens(text)
                _, cache = model.run_with_cache(tokens, names_filter=[hook_name])
                activations = safe_get_cache(cache, hook_name)
                feature_acts = sae.encode(activations)
                max_pooled = feature_acts[0].max(dim=0).values.detach().float().cpu().numpy()
                all_vectors[tag] = max_pooled
                all_str_tokens[tag] = [t.replace('\u0120', ' ').replace('\u010a', '\n') for t in str_tokens]
                del cache, activations, feature_acts
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        fin_en_vectors = []
        fin_jp_vectors = []
        financial_pair_results = []

        for i, pair in enumerate(financial_pairs):
            en_vec = all_vectors[('fin_en', i)]
            jp_vec = all_vectors[('fin_jp', i)]
            en_tokens = all_str_tokens[('fin_en', i)]
            jp_tokens = all_str_tokens[('fin_jp', i)]
            fin_en_vectors.append(en_vec)
            fin_jp_vectors.append(jp_vec)

            cos_sim = cosine_sim(en_vec, jp_vec)

            en_top_k_indices = np.argsort(-en_vec)[:top_k]
            jp_top_k_indices = np.argsort(-jp_vec)[:top_k]
            shared_top_k = set(en_top_k_indices.tolist()) & set(jp_top_k_indices.tolist())

            combined_indices = np.unique(np.concatenate([en_top_k_indices, jp_top_k_indices]))
            top_financial = []
            for idx in combined_indices[:top_k]:
                idx = int(idx)
                top_financial.append({
                    "feature": idx,
                    "englishActivation": round(float(en_vec[idx]), 4),
                    "japaneseActivation": round(float(jp_vec[idx]), 4)
                })
            top_financial.sort(key=lambda x: max(x['englishActivation'], x['japaneseActivation']), reverse=True)
            top_financial = top_financial[:top_k]

            financial_pair_results.append({
                "english": pair['english'],
                "japanese": pair['japanese'],
                "label": pair.get('label', ''),
                "englishTokens": en_tokens,
                "japaneseTokens": jp_tokens,
                "cosineSimilarity": round(cos_sim, 4),
                "sharedFeatureCount": len(shared_top_k),
                "topFinancialFeatures": top_financial,
                "totalActiveEnglish": int(np.sum(en_vec > 0.1)),
                "totalActiveJapanese": int(np.sum(jp_vec > 0.1))
            })

        base_en_vectors = []
        base_jp_vectors = []
        baseline_pair_results = []

        for i, pair in enumerate(baseline_pairs):
            en_vec = all_vectors[('base_en', i)]
            jp_vec = all_vectors[('base_jp', i)]
            base_en_vectors.append(en_vec)
            base_jp_vectors.append(jp_vec)

            cos_sim = cosine_sim(en_vec, jp_vec)
            en_top_k_indices = np.argsort(-en_vec)[:top_k]
            jp_top_k_indices = np.argsort(-jp_vec)[:top_k]
            shared_top_k = set(en_top_k_indices.tolist()) & set(jp_top_k_indices.tolist())

            baseline_pair_results.append({
                "english": pair['english'],
                "japanese": pair['japanese'],
                "label": pair.get('label', ''),
                "cosineSimilarity": round(cos_sim, 4),
                "sharedFeatureCount": len(shared_top_k)
            })

        mean_fin_en = np.mean(np.stack(fin_en_vectors), axis=0)
        mean_fin_jp = np.mean(np.stack(fin_jp_vectors), axis=0)
        mean_base_en = np.mean(np.stack(base_en_vectors), axis=0)
        mean_base_jp = np.mean(np.stack(base_jp_vectors), axis=0)

        specificity = np.maximum(mean_fin_en, mean_fin_jp) - np.maximum(mean_base_en, mean_base_jp)

        finance_specific_mask = specificity > 0.1
        finance_specific_indices = np.where(finance_specific_mask)[0]
        finance_sorted = finance_specific_indices[np.argsort(-specificity[finance_specific_indices])][:30]

        finance_specific_list = []
        for idx in finance_sorted:
            idx = int(idx)
            finance_specific_list.append({
                "feature": idx,
                "specificityScore": round(float(specificity[idx]), 4),
                "financialMeanEN": round(float(mean_fin_en[idx]), 4),
                "financialMeanJP": round(float(mean_fin_jp[idx]), 4),
                "baselineMeanEN": round(float(mean_base_en[idx]), 4),
                "baselineMeanJP": round(float(mean_base_jp[idx]), 4)
            })

        universal_mask = finance_specific_mask & (mean_fin_en > 0.1) & (mean_fin_jp > 0.1)
        universal_indices = np.where(universal_mask)[0]
        universal_sorted = universal_indices[np.argsort(-specificity[universal_indices])][:20]

        universal_list = []
        for idx in universal_sorted:
            idx = int(idx)
            universal_list.append({
                "feature": idx,
                "englishActivation": round(float(mean_fin_en[idx]), 4),
                "japaneseActivation": round(float(mean_fin_jp[idx]), 4),
                "specificityScore": round(float(specificity[idx]), 4)
            })

        en_only_mask = finance_specific_mask & (mean_fin_en > 0.1) & (mean_fin_jp < 0.05)
        en_only_indices = np.where(en_only_mask)[0]
        en_only_sorted = en_only_indices[np.argsort(-mean_fin_en[en_only_indices])][:15]

        en_only_list = []
        for idx in en_only_sorted:
            idx = int(idx)
            en_only_list.append({
                "feature": idx,
                "activation": round(float(mean_fin_en[idx]), 4),
                "specificityScore": round(float(specificity[idx]), 4)
            })

        jp_only_mask = finance_specific_mask & (mean_fin_jp > 0.1) & (mean_fin_en < 0.05)
        jp_only_indices = np.where(jp_only_mask)[0]
        jp_only_sorted = jp_only_indices[np.argsort(-mean_fin_jp[jp_only_indices])][:15]

        jp_only_list = []
        for idx in jp_only_sorted:
            idx = int(idx)
            jp_only_list.append({
                "feature": idx,
                "activation": round(float(mean_fin_jp[idx]), 4),
                "specificityScore": round(float(specificity[idx]), 4)
            })

        cross_lingual_overlap = cosine_sim(mean_fin_en, mean_fin_jp)

        avg_fin_cos = round(float(np.mean([r['cosineSimilarity'] for r in financial_pair_results])), 4)
        avg_base_cos = round(float(np.mean([r['cosineSimilarity'] for r in baseline_pair_results])), 4)

        return jsonify({
            "model": model_id,
            "nFeatures": int(n_features),
            "saeHook": hook_name,
            "nFinancialPairs": len(financial_pairs),
            "nBaselinePairs": len(baseline_pairs),
            "financialPairResults": financial_pair_results,
            "baselinePairResults": baseline_pair_results,
            "financeSpecificFeatures": finance_specific_list,
            "universalFinancialFeatures": universal_list,
            "englishOnlyFinancial": en_only_list,
            "japaneseOnlyFinancial": jp_only_list,
            "crossLingualOverlap": round(float(cross_lingual_overlap), 4),
            "avgFinancialCosineSim": avg_fin_cos,
            "avgBaselineCosineSim": avg_base_cos,
            "financialVsBaselineGap": round(avg_fin_cos - avg_base_cos, 4),
            "totalFinanceSpecific": int(np.sum(finance_specific_mask)),
            "totalUniversalFinancial": int(np.sum(universal_mask)),
            "totalEnglishOnlyFinancial": int(np.sum(en_only_mask)),
            "totalJapaneseOnlyFinancial": int(np.sum(jp_only_mask))
        })

    except Exception as e:
        print(f"[FinancialCircuits] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/domain-benchmark/analyze', methods=['POST'])
def domain_benchmark_analyze():
    """
    Analyze model prediction confidence across domains for English vs Japanese text.

    Request body:
    {
        "domains": [
            {
                "domain": "Legal",
                "englishTexts": ["The court ruled...", ...],
                "japaneseTexts": ["裁判所は...", ...]
            },
            ...
        ],
        "model": "gpt2-small"
    }

    Response:
    {
        "model": "gpt2-small",
        "domains": [
            {
                "domain": "Legal",
                "englishAccuracy": 0.85,
                "japaneseAccuracy": 0.62,
                "englishEntropy": 2.1,
                "japaneseEntropy": 3.4,
                "gap": 0.23
            },
            ...
        ],
        "overallEnglishAccuracy": 0.82,
        "overallJapaneseAccuracy": 0.60
    }
    """
    try:
        data = request.json
        domains = data.get('domains', [])
        model_id = data.get('model', data.get('model_id', 'gpt2-small'))

        if not domains:
            return jsonify({"error": "No domains provided"}), 400

        import torch.nn.functional as F

        with _model_lock:
            model, sae = get_model_and_sae(model_id)

            domain_results = []
            all_en_confidences = []
            all_jp_confidences = []

            with torch.no_grad():
                for domain_entry in domains:
                    domain_name = domain_entry.get('domain', 'Unknown')
                    en_texts = domain_entry.get('englishTexts', [])
                    jp_texts = domain_entry.get('japaneseTexts', [])

                    en_confidences = []
                    en_entropies = []
                    for text in en_texts:
                        if not text or not text.strip():
                            continue
                        try:
                            tokens = model.to_tokens(text)
                            logits = model(tokens)
                            final_logits = logits[0, -1, :]
                            probs = F.softmax(final_logits, dim=-1)
                            top1_prob = probs.max().item()
                            log_probs = torch.log(probs + 1e-10)
                            entropy = -(probs * log_probs).sum().item()
                            en_confidences.append(top1_prob)
                            en_entropies.append(entropy)
                        except Exception as e:
                            print(f"[DomainBenchmark] Error processing EN text: {e}")

                    jp_confidences = []
                    jp_entropies = []
                    for text in jp_texts:
                        if not text or not text.strip():
                            continue
                        try:
                            tokens = model.to_tokens(text)
                            logits = model(tokens)
                            final_logits = logits[0, -1, :]
                            probs = F.softmax(final_logits, dim=-1)
                            top1_prob = probs.max().item()
                            log_probs = torch.log(probs + 1e-10)
                            entropy = -(probs * log_probs).sum().item()
                            jp_confidences.append(top1_prob)
                            jp_entropies.append(entropy)
                        except Exception as e:
                            print(f"[DomainBenchmark] Error processing JP text: {e}")

                    en_acc = float(np.mean(en_confidences)) if en_confidences else 0.0
                    jp_acc = float(np.mean(jp_confidences)) if jp_confidences else 0.0
                    en_ent = float(np.mean(en_entropies)) if en_entropies else 0.0
                    jp_ent = float(np.mean(jp_entropies)) if jp_entropies else 0.0

                    all_en_confidences.extend(en_confidences)
                    all_jp_confidences.extend(jp_confidences)

                    domain_results.append({
                        "domain": domain_name,
                        "englishAccuracy": round(en_acc, 4),
                        "japaneseAccuracy": round(jp_acc, 4),
                        "englishEntropy": round(en_ent, 4),
                        "japaneseEntropy": round(jp_ent, 4),
                        "gap": round(en_acc - jp_acc, 4)
                    })

                    print(f"[DomainBenchmark] {domain_name}: EN={en_acc:.4f}, JP={jp_acc:.4f}, gap={en_acc - jp_acc:.4f}")

            overall_en = round(float(np.mean(all_en_confidences)), 4) if all_en_confidences else 0.0
            overall_jp = round(float(np.mean(all_jp_confidences)), 4) if all_jp_confidences else 0.0

            return jsonify({
                "model": model_id,
                "domains": domain_results,
                "overallEnglishAccuracy": overall_en,
                "overallJapaneseAccuracy": overall_jp
            })

    except Exception as e:
        print(f"[DomainBenchmark] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/statistical_test', methods=['POST'])
def statistical_test():
    """Run formal statistical tests on experimental data.
    
    Supports:
    - paired_bootstrap: Bootstrap confidence intervals for paired measurements (e.g., EN vs JP similarities)
    - permutation_test: Permutation test for group differences
    - effect_size: Cohen's d effect size with bootstrap CI
    - correlation: Pearson/Spearman correlation with significance
    - one_sample: One-sample test against a reference value (e.g., is mean similarity > 0.5?)
    - summary_stats: Descriptive statistics with bootstrap CIs for means
    
    All tests use non-parametric methods (no normality assumption required).
    """
    try:
        data = request.json
        if not data:
            return jsonify({"error": "No data provided"}), 400
        
        test_type = data.get("test", "summary_stats")
        n_bootstrap = min(data.get("nBootstrap", 10000), 50000)
        confidence_level = data.get("confidenceLevel", 0.95)
        
        import scipy.stats as stats
        
        def bootstrap_ci(values, n_boot=10000, ci=0.95):
            values = np.array(values, dtype=float)
            values = values[~np.isnan(values)]
            if len(values) < 2:
                return {"mean": float(np.mean(values)) if len(values) > 0 else None, "ci_lower": None, "ci_upper": None, "n": len(values)}
            boot_means = np.array([np.mean(np.random.choice(values, size=len(values), replace=True)) for _ in range(n_boot)])
            alpha = 1 - ci
            ci_lower = float(np.percentile(boot_means, 100 * alpha / 2))
            ci_upper = float(np.percentile(boot_means, 100 * (1 - alpha / 2)))
            return {
                "mean": float(np.mean(values)),
                "std": float(np.std(values, ddof=1)),
                "ci_lower": ci_lower,
                "ci_upper": ci_upper,
                "ci_level": ci,
                "n": len(values),
                "se": float(np.std(boot_means)),
            }
        
        result = {"test": test_type, "confidenceLevel": confidence_level}
        
        if test_type == "summary_stats":
            groups = data.get("groups", {})
            if not groups:
                values = data.get("values", [])
                groups = {"data": values}
            
            group_results = {}
            for name, values in groups.items():
                arr = np.array(values, dtype=float)
                arr = arr[~np.isnan(arr)]
                ci = bootstrap_ci(arr, n_bootstrap, confidence_level)
                group_results[name] = {
                    **ci,
                    "median": float(np.median(arr)) if len(arr) > 0 else None,
                    "min": float(np.min(arr)) if len(arr) > 0 else None,
                    "max": float(np.max(arr)) if len(arr) > 0 else None,
                    "iqr": float(np.percentile(arr, 75) - np.percentile(arr, 25)) if len(arr) >= 4 else None,
                }
            result["groups"] = group_results
        
        elif test_type == "paired_bootstrap":
            group_a = np.array(data.get("groupA", []), dtype=float)
            group_b = np.array(data.get("groupB", []), dtype=float)
            label_a = data.get("labelA", "A")
            label_b = data.get("labelB", "B")
            
            min_len = min(len(group_a), len(group_b))
            if min_len < 2:
                return jsonify({"error": f"Need at least 2 paired observations, got {min_len}"}), 400
            group_a, group_b = group_a[:min_len], group_b[:min_len]
            
            diffs = group_a - group_b
            boot_diffs = np.array([np.mean(np.random.choice(diffs, size=len(diffs), replace=True)) for _ in range(n_bootstrap)])
            alpha = 1 - confidence_level
            
            p_value = float(2 * min(np.mean(boot_diffs <= 0), np.mean(boot_diffs >= 0)))
            p_value = min(p_value, 1.0)
            
            pooled_std = float(np.sqrt((np.var(group_a, ddof=1) + np.var(group_b, ddof=1)) / 2))
            cohens_d = float(np.mean(diffs) / pooled_std) if pooled_std > 0 else 0.0
            
            result.update({
                "labelA": label_a,
                "labelB": label_b,
                "n": min_len,
                "meanDiff": float(np.mean(diffs)),
                "diffCI": {
                    "lower": float(np.percentile(boot_diffs, 100 * alpha / 2)),
                    "upper": float(np.percentile(boot_diffs, 100 * (1 - alpha / 2))),
                },
                "pValue": p_value,
                "significant": p_value < (1 - confidence_level),
                "effectSize": cohens_d,
                "effectMagnitude": "large" if abs(cohens_d) >= 0.8 else "medium" if abs(cohens_d) >= 0.5 else "small" if abs(cohens_d) >= 0.2 else "negligible",
                "summaryA": bootstrap_ci(group_a, n_bootstrap, confidence_level),
                "summaryB": bootstrap_ci(group_b, n_bootstrap, confidence_level),
                "interpretation": f"{'Significant' if p_value < (1 - confidence_level) else 'No significant'} difference between {label_a} and {label_b} (mean diff = {float(np.mean(diffs)):.4f}, p = {p_value:.4f}, Cohen's d = {cohens_d:.3f} [{('large' if abs(cohens_d) >= 0.8 else 'medium' if abs(cohens_d) >= 0.5 else 'small' if abs(cohens_d) >= 0.2 else 'negligible')}], n = {min_len})",
            })
        
        elif test_type == "permutation_test":
            group_a = np.array(data.get("groupA", []), dtype=float)
            group_b = np.array(data.get("groupB", []), dtype=float)
            label_a = data.get("labelA", "A")
            label_b = data.get("labelB", "B")
            n_perms = min(data.get("nPermutations", 10000), 50000)
            
            if len(group_a) < 2 or len(group_b) < 2:
                return jsonify({"error": "Need at least 2 observations per group"}), 400
            
            observed_diff = float(np.mean(group_a) - np.mean(group_b))
            combined = np.concatenate([group_a, group_b])
            na = len(group_a)
            
            perm_diffs = np.zeros(n_perms)
            for i in range(n_perms):
                np.random.shuffle(combined)
                perm_diffs[i] = np.mean(combined[:na]) - np.mean(combined[na:])
            
            p_value = float(np.mean(np.abs(perm_diffs) >= abs(observed_diff)))
            
            pooled_std = float(np.sqrt(((len(group_a) - 1) * np.var(group_a, ddof=1) + (len(group_b) - 1) * np.var(group_b, ddof=1)) / (len(group_a) + len(group_b) - 2)))
            cohens_d = float(observed_diff / pooled_std) if pooled_std > 0 else 0.0
            
            result.update({
                "labelA": label_a,
                "labelB": label_b,
                "nA": len(group_a),
                "nB": len(group_b),
                "observedDiff": observed_diff,
                "pValue": p_value,
                "significant": p_value < (1 - confidence_level),
                "effectSize": cohens_d,
                "effectMagnitude": "large" if abs(cohens_d) >= 0.8 else "medium" if abs(cohens_d) >= 0.5 else "small" if abs(cohens_d) >= 0.2 else "negligible",
                "summaryA": bootstrap_ci(group_a, n_bootstrap, confidence_level),
                "summaryB": bootstrap_ci(group_b, n_bootstrap, confidence_level),
                "interpretation": f"Permutation test: {'Significant' if p_value < (1 - confidence_level) else 'No significant'} difference between {label_a} and {label_b} (observed diff = {observed_diff:.4f}, p = {p_value:.4f}, Cohen's d = {cohens_d:.3f}, nA = {len(group_a)}, nB = {len(group_b)})",
            })
        
        elif test_type == "effect_size":
            group_a = np.array(data.get("groupA", []), dtype=float)
            group_b = np.array(data.get("groupB", []), dtype=float)
            
            if len(group_a) < 2 or len(group_b) < 2:
                return jsonify({"error": "Need at least 2 observations per group"}), 400
            
            pooled_std = float(np.sqrt(((len(group_a) - 1) * np.var(group_a, ddof=1) + (len(group_b) - 1) * np.var(group_b, ddof=1)) / (len(group_a) + len(group_b) - 2)))
            cohens_d = float((np.mean(group_a) - np.mean(group_b)) / pooled_std) if pooled_std > 0 else 0.0
            
            boot_ds = []
            for _ in range(n_bootstrap):
                ba = np.random.choice(group_a, size=len(group_a), replace=True)
                bb = np.random.choice(group_b, size=len(group_b), replace=True)
                ps = float(np.sqrt(((len(ba) - 1) * np.var(ba, ddof=1) + (len(bb) - 1) * np.var(bb, ddof=1)) / (len(ba) + len(bb) - 2)))
                boot_ds.append(float((np.mean(ba) - np.mean(bb)) / ps) if ps > 0 else 0.0)
            boot_ds = np.array(boot_ds)
            alpha = 1 - confidence_level
            
            result.update({
                "cohensD": cohens_d,
                "effectMagnitude": "large" if abs(cohens_d) >= 0.8 else "medium" if abs(cohens_d) >= 0.5 else "small" if abs(cohens_d) >= 0.2 else "negligible",
                "ci": {
                    "lower": float(np.percentile(boot_ds, 100 * alpha / 2)),
                    "upper": float(np.percentile(boot_ds, 100 * (1 - alpha / 2))),
                },
                "nA": len(group_a),
                "nB": len(group_b),
            })
        
        elif test_type == "correlation":
            x = np.array(data.get("x", []), dtype=float)
            y = np.array(data.get("y", []), dtype=float)
            method = data.get("method", "pearson")
            
            min_len = min(len(x), len(y))
            if min_len < 3:
                return jsonify({"error": f"Need at least 3 observations for correlation, got {min_len}"}), 400
            x, y = x[:min_len], y[:min_len]
            
            if method == "spearman":
                r, p = stats.spearmanr(x, y)
            else:
                r, p = stats.pearsonr(x, y)
            
            boot_rs = []
            for _ in range(n_bootstrap):
                idx = np.random.choice(len(x), size=len(x), replace=True)
                if method == "spearman":
                    br, _ = stats.spearmanr(x[idx], y[idx])
                else:
                    br, _ = stats.pearsonr(x[idx], y[idx])
                boot_rs.append(br)
            boot_rs = np.array(boot_rs)
            alpha_val = 1 - confidence_level
            
            result.update({
                "method": method,
                "r": float(r),
                "pValue": float(p),
                "significant": float(p) < (1 - confidence_level),
                "ci": {
                    "lower": float(np.percentile(boot_rs, 100 * alpha_val / 2)),
                    "upper": float(np.percentile(boot_rs, 100 * (1 - alpha_val / 2))),
                },
                "n": min_len,
                "rSquared": float(r ** 2),
                "interpretation": f"{method.capitalize()} r = {float(r):.4f} (p = {float(p):.4f}), {'significant' if float(p) < (1 - confidence_level) else 'not significant'}, R² = {float(r**2):.4f}, n = {min_len}",
            })
        
        elif test_type == "one_sample":
            values = np.array(data.get("values", []), dtype=float)
            reference = data.get("reference", 0.0)
            alternative = data.get("alternative", "two-sided")
            
            values = values[~np.isnan(values)]
            if len(values) < 2:
                return jsonify({"error": "Need at least 2 observations"}), 400
            
            observed_mean = float(np.mean(values))
            ci = bootstrap_ci(values, n_bootstrap, confidence_level)
            
            boot_means = np.array([np.mean(np.random.choice(values, size=len(values), replace=True)) for _ in range(n_bootstrap)])
            centered = boot_means - observed_mean + reference
            
            if alternative == "greater":
                p_value = float(np.mean(centered >= observed_mean))
            elif alternative == "less":
                p_value = float(np.mean(centered <= observed_mean))
            else:
                p_value = float(2 * min(np.mean(centered >= observed_mean), np.mean(centered <= observed_mean)))
            p_value = min(p_value, 1.0)
            
            result.update({
                "observedMean": observed_mean,
                "reference": reference,
                "alternative": alternative,
                "pValue": p_value,
                "significant": p_value < (1 - confidence_level),
                "ci": {"lower": ci["ci_lower"], "upper": ci["ci_upper"]},
                "n": len(values),
                "interpretation": f"One-sample test: mean = {observed_mean:.4f} vs reference = {reference} ({alternative}), p = {p_value:.4f}, {'significant' if p_value < (1 - confidence_level) else 'not significant'}, n = {len(values)}",
            })
        
        else:
            return jsonify({"error": f"Unknown test type: {test_type}. Supported: summary_stats, paired_bootstrap, permutation_test, effect_size, correlation, one_sample"}), 400
        
        print(f"[StatisticalTest] {test_type} completed successfully")
        return jsonify(result)
    
    except Exception as e:
        print(f"[StatisticalTest] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


_specter2_model = None
_specter2_tokenizer = None

def _load_specter2():
    """Lazy-load SPECTER2 model for paper embeddings. ~110MB, runs on CPU."""
    global _specter2_model, _specter2_tokenizer
    if _specter2_model is not None:
        return _specter2_model, _specter2_tokenizer
    
    print("[SPECTER2] Loading allenai/specter2_base...")
    from transformers import AutoTokenizer, AutoModel
    _specter2_tokenizer = AutoTokenizer.from_pretrained("allenai/specter2_base")
    _specter2_model = AutoModel.from_pretrained("allenai/specter2_base")
    _specter2_model.eval()
    print("[SPECTER2] Model loaded successfully")
    return _specter2_model, _specter2_tokenizer

@app.route('/specter2/embed', methods=['POST'])
def specter2_embed():
    """Embed one or more papers using local SPECTER2. 
    Input: { "papers": [{"title": "...", "abstract": "..."}] }
    Returns: { "embeddings": [[...768 floats...], ...] }
    """
    try:
        data = request.json
        papers = data.get("papers", [])
        if not papers:
            return jsonify({"error": "No papers provided"}), 400
        if len(papers) > 32:
            return jsonify({"error": "Max 32 papers per batch"}), 400
        
        model, tokenizer = _load_specter2()
        
        texts = []
        for p in papers:
            title = p.get("title", "")
            abstract = p.get("abstract", "")
            text = title + (tokenizer.sep_token + abstract if abstract else "")
            texts.append(text)
        
        inputs = tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        )
        
        with torch.no_grad():
            outputs = model(**inputs)
            embeddings = outputs.last_hidden_state[:, 0, :]
        
        embedding_list = embeddings.cpu().numpy().tolist()
        
        return jsonify({"embeddings": embedding_list, "count": len(embedding_list), "dim": len(embedding_list[0])})
    
    except Exception as e:
        print(f"[SPECTER2] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/warmup', methods=['POST'])
def warmup():
    """Pre-load a model and run a test inference so subsequent requests are fast."""
    model_id = request.json.get("model", "gpt2-small") if request.json else "gpt2-small"
    try:
        print(f"[SAE Service] Warmup request for {model_id}...")
        model, sae = get_model_and_sae(model_id)
        print(f"[SAE Service] {model_id} loaded, running test inference...")

        import time
        t0 = time.time()
        with torch.no_grad():
            tokens = model.to_tokens("The quick brown fox")
            _, cache = model.run_with_cache(tokens)
        elapsed = time.time() - t0
        print(f"[SAE Service] Test inference complete in {elapsed:.1f}s — fully warmed up.")

        return jsonify({"status": "ready", "model": model_id, "test_inference_ms": int(elapsed * 1000)})
    except Exception as e:
        print(f"[SAE Service] Warmup failed: {e}")
        return jsonify({"status": "error", "error": str(e)}), 500


_openreviewer_model = None
_openreviewer_tokenizer = None

def _load_openreviewer():
    """Load the Llama-OpenReviewer-8B model for peer review generation.
    Clears SAE models from VRAM first since the 8B model needs the full T4."""
    global _openreviewer_model, _openreviewer_tokenizer
    if _openreviewer_model is not None:
        return _openreviewer_model, _openreviewer_tokenizer

    _cleanup_model_vram()

    device = _get_device()
    print(f"[OpenReviewer] Loading maxidl/Llama-OpenReviewer-8B on {device}...")

    from transformers import AutoModelForCausalLM, AutoTokenizer
    hf_token = os.environ.get("HF_TOKEN")

    _openreviewer_tokenizer = AutoTokenizer.from_pretrained(
        "maxidl/Llama-OpenReviewer-8B",
        token=hf_token,
    )
    _openreviewer_model = AutoModelForCausalLM.from_pretrained(
        "maxidl/Llama-OpenReviewer-8B",
        token=hf_token,
        torch_dtype=torch.float16,
        device_map=device,
    )
    print(f"[OpenReviewer] Model loaded on {device}")
    if torch.cuda.is_available():
        free_mb = torch.cuda.mem_get_info()[0] // (1024 * 1024)
        print(f"[OpenReviewer] VRAM free after load: {free_mb} MB")
    return _openreviewer_model, _openreviewer_tokenizer

def _unload_openreviewer():
    """Free OpenReviewer from VRAM so SAE models can load."""
    global _openreviewer_model, _openreviewer_tokenizer
    if _openreviewer_model is not None:
        try:
            _openreviewer_model.cpu()
        except Exception:
            pass
        del _openreviewer_model
        _openreviewer_model = None
    if _openreviewer_tokenizer is not None:
        del _openreviewer_tokenizer
        _openreviewer_tokenizer = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    import gc
    gc.collect()
    print("[OpenReviewer] Model unloaded from VRAM")


@app.route('/openreviewer/chat', methods=['POST'])
def openreviewer_chat():
    """OpenReviewer chat completions endpoint.

    Expects OpenAI-compatible request body:
    {
        "messages": [
            {"role": "system", "content": "..."},
            {"role": "user", "content": "..."}
        ],
        "max_tokens": 2048,
        "temperature": 0.3
    }

    Returns OpenAI-compatible response:
    {
        "choices": [{"message": {"role": "assistant", "content": "..."}}],
        "model": "maxidl/Llama-OpenReviewer-8B"
    }
    """
    try:
        data = request.json or {}
        messages = data.get("messages", [])
        max_tokens = data.get("max_tokens", 2048)
        temperature = data.get("temperature", 0.3)

        if not messages:
            return jsonify({"error": "No messages provided"}), 400

        with _model_swap_lock:
            model, tokenizer = _load_openreviewer()

            if tokenizer.chat_template:
                prompt_text = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            else:
                parts = []
                for m in messages:
                    role = m.get("role", "user")
                    content = m.get("content", "")
                    if role == "system":
                        parts.append(f"<<SYS>>\n{content}\n<</SYS>>\n\n")
                    elif role == "user":
                        parts.append(f"[INST] {content} [/INST]\n")
                    elif role == "assistant":
                        parts.append(f"{content}\n")
                prompt_text = "".join(parts)

            inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)

            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    temperature=max(temperature, 0.01),
                    do_sample=temperature > 0,
                    top_p=0.9,
                    pad_token_id=tokenizer.eos_token_id,
                )

            new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
            response_text = tokenizer.decode(new_tokens, skip_special_tokens=True)

            print(f"[OpenReviewer] Generated {len(new_tokens)} tokens")

        return jsonify({
            "choices": [{"message": {"role": "assistant", "content": response_text}}],
            "model": "maxidl/Llama-OpenReviewer-8B",
        })

    except torch.cuda.OutOfMemoryError:
        _unload_openreviewer()
        return jsonify({"error": "CUDA out of memory loading OpenReviewer. Try again after VRAM is freed."}), 503
    except Exception as e:
        print(f"[OpenReviewer] Error: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500




@app.route('/superposition-analysis/analyze', methods=['POST'])
def superposition_analysis_analyze():
    """
    Profile MLP neurons in a chosen layer to expose polysemanticity / superposition.

    Pipeline:
    1. Forward each prompt with cache; gather `blocks.{L}.mlp.hook_post` -> [N_prompts, T, d_mlp].
    2. Concatenate token-level activations into [N_total, d_mlp]; record (prompt_idx, pos, token_id) per row.
    3. For each neuron compute max, mean, fire_rate (% tokens > 0).
    4. Pick top-K neurons by max activation.
    5. For each chosen neuron, take top-N rows by activation → record (token, context window, activation).
    6. Polysemy score: take top-N tokens' unembedding rows W_U[:,token_id], compute mean pairwise
       cosine similarity; polysemy_score = 1 − mean_cos. Coherence = mean_cos.

    Request:  { model, prompts:[..], layer, topK, examplesPerNeuron }
    Response: { success, device, layer, totalNeurons, totalTokens, meanFireRate, meanPolysemy,
                neurons:[{neuron,maxActivation,meanActivation,fireRate,polysemyScore,
                          semanticCoherence,examples:[{token,activation,context}]}] }
    """
    try:
        data = request.json or {}
        prompts = [p for p in (data.get('prompts') or []) if isinstance(p, str) and p.strip()][:40]
        if len(prompts) < 4:
            return jsonify({"error": "at least 4 prompts required"}), 400
        model_id = data.get('model', 'gpt2-small')
        layer = int(data.get('layer') or 6)
        top_k = max(3, min(20, int(data.get('topK') or 8)))
        examples_per_neuron = max(3, min(10, int(data.get('examplesPerNeuron') or 5)))

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)
        device = next(model.parameters()).device
        layer = max(0, min(model.cfg.n_layers - 1, layer))

        # Collect activations from chosen layer's MLP post-activation
        hook_name = f'blocks.{layer}.mlp.hook_post'
        all_acts = []  # list of [T_i, d_mlp]
        all_token_ids = []  # list of [T_i]
        all_prompt_idx = []  # list of [T_i] indices into prompts
        for pi, p in enumerate(prompts):
            try:
                tokens = model.to_tokens(p).to(device)
                if tokens.shape[1] > 96:
                    tokens = tokens[:, :96]
                with torch.no_grad():
                    _, cache = model.run_with_cache(tokens, names_filter=lambda n: n == hook_name)
                acts = cache[hook_name][0].float()  # [T, d_mlp]
                all_acts.append(acts)
                all_token_ids.append(tokens[0])
                all_prompt_idx.extend([pi] * acts.shape[0])
                del cache
            except Exception as inner:
                print(f"[SuperpositionAnalysis] skip prompt {pi}: {inner}")
        if not all_acts:
            return jsonify({"error": "no usable prompts"}), 400

        acts = torch.cat(all_acts, dim=0)  # [N_total, d_mlp]
        token_ids = torch.cat(all_token_ids, dim=0)  # [N_total]
        prompt_idx = torch.tensor(all_prompt_idx, device=device)
        n_total, d_mlp = acts.shape

        max_per_neuron = acts.max(dim=0).values  # [d_mlp]
        mean_per_neuron = acts.mean(dim=0)  # [d_mlp]
        fire_per_neuron = (acts > 0).float().mean(dim=0)  # [d_mlp]
        mean_fire_rate = fire_per_neuron.mean().item()

        # top-K neurons by max activation
        top_neuron_ids = max_per_neuron.topk(top_k).indices.tolist()

        # Unembedding for cosine: model.W_U is [d_model, vocab] in HookedTransformer; we want columns.
        W_U = model.W_U.float()  # [d_model, vocab]

        neurons_out = []
        polysemy_scores = []
        for nid in top_neuron_ids:
            col = acts[:, nid]  # [N_total]
            top_vals, top_pos = col.topk(min(examples_per_neuron, n_total))
            examples = []
            top_token_ids_for_neuron = []
            for v, p in zip(top_vals.tolist(), top_pos.tolist()):
                pi = prompt_idx[p].item()
                # local position within prompt = p - start_offset
                offset = sum(t.shape[0] for t in all_token_ids[:pi])
                local_pos = p - offset
                tok_seq = all_token_ids[pi]
                tok_id = int(tok_seq[local_pos].item())
                top_token_ids_for_neuron.append(tok_id)
                # build context window with marked token
                lo = max(0, local_pos - 4)
                hi = min(tok_seq.shape[0], local_pos + 5)
                pre = model.to_string(tok_seq[lo:local_pos]) if local_pos > lo else ""
                tok_str = model.to_string(tok_seq[local_pos:local_pos+1])
                post = model.to_string(tok_seq[local_pos+1:hi]) if hi > local_pos+1 else ""
                examples.append({
                    "token": tok_str,
                    "activation": round(float(v), 4),
                    "context": f"{pre}«{tok_str}»{post}",
                })

            # Polysemy: mean pairwise cosine similarity of unembed rows of top tokens.
            # Lower mean cos => tokens are spread out in embedding space => polysemantic.
            tt = torch.tensor(top_token_ids_for_neuron, device=device)
            embs = W_U[:, tt].t()  # [N_top, d_model]
            embs_n = embs / embs.norm(dim=-1, keepdim=True).clamp_min(1e-9)
            sim_matrix = embs_n @ embs_n.t()  # [N_top, N_top]
            n_top = sim_matrix.shape[0]
            if n_top > 1:
                # take upper triangle excluding diagonal
                mask = torch.triu(torch.ones_like(sim_matrix), diagonal=1).bool()
                mean_cos = sim_matrix[mask].mean().item()
            else:
                mean_cos = 1.0
            poly = max(0.0, 1.0 - mean_cos)
            polysemy_scores.append(poly)

            neurons_out.append({
                "neuron": int(nid),
                "maxActivation": round(float(max_per_neuron[nid].item()), 4),
                "meanActivation": round(float(mean_per_neuron[nid].item()), 4),
                "fireRate": round(float(fire_per_neuron[nid].item()), 4),
                "polysemyScore": round(poly, 4),
                "semanticCoherence": round(mean_cos, 4),
                "examples": examples,
            })

        mean_polysemy = sum(polysemy_scores) / max(1, len(polysemy_scores))

        return jsonify({
            "success": True,
            "device": str(device),
            "layer": layer,
            "totalNeurons": int(d_mlp),
            "totalTokens": int(n_total),
            "meanFireRate": round(mean_fire_rate, 4),
            "meanPolysemy": round(mean_polysemy, 4),
            "neurons": neurons_out,
        })
    except Exception as e:
        print(f"[SuperpositionAnalysis] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/data-leakage-scanner/analyze', methods=['POST'])
def data_leakage_scanner_analyze():
    """
    Detect probable training-data memorization for each candidate string.

    Two signals per candidate:
      1. Teacher-forcing perplexity, compared to a fixed-baseline perplexity (computed on a
         neutral reference text on the same model). A candidate ppl significantly below
         baseline indicates the model finds it abnormally predictable.
      2. Exact-completion test: split the candidate at the midpoint (token-wise), feed the
         prefix, greedily generate `len(suffix)` tokens, count exact-token matches against
         the held-out suffix.

    Combined leakage score = 0.5*ppl_norm + 0.5*exact_ratio, where
      ppl_norm = clip(1 - ppl/baseline, 0, 1)  (capped so score in [0,1]).

    Verdicts: >=0.6 likely-memorized, >=0.35 suspicious, else safe.

    Request:  { model, candidates:[..] }
    Response: { success, device, baselinePerplexity,
                candidates:[{candidate,tokens,meanNll,perplexity,prefixTokens,suffixTokens,
                             exactMatchTokens,exactMatchRatio,generatedSuffix,expectedSuffix,
                             leakageScore,verdict}] }
    """
    try:
        data = request.json or {}
        candidates = [c for c in (data.get('candidates') or []) if isinstance(c, str) and c.strip()][:20]
        if not candidates:
            return jsonify({"error": "candidates required"}), 400
        model_id = data.get('model', 'gpt2-small')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)
        device = next(model.parameters()).device

        def compute_nll(text):
            tokens = model.to_tokens(text).to(device)
            if tokens.shape[1] < 2:
                return None, 0
            if tokens.shape[1] > 256:
                tokens = tokens[:, :256]
            with torch.no_grad():
                logits = model(tokens)
            log_probs = torch.log_softmax(logits[0, :-1, :].float(), dim=-1)
            targets = tokens[0, 1:]
            token_lp = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
            return -token_lp.mean().item(), tokens.shape[1] - 1

        # Baseline: fixed neutral reference text. Average a few sentences for stability.
        baseline_refs = [
            "The garden contained several plants of various sizes and colors arranged along the gravel path.",
            "She decided to take the longer route home because she wanted some time to think things over.",
            "The committee will reconvene next week to review the proposed changes to the operating procedures.",
        ]
        baseline_nlls = []
        for ref in baseline_refs:
            n, _ = compute_nll(ref)
            if n is not None:
                baseline_nlls.append(n)
        baseline_mean_nll = sum(baseline_nlls) / max(1, len(baseline_nlls)) if baseline_nlls else 5.0
        baseline_ppl = float(torch.tensor(baseline_mean_nll).exp().item())

        results = []
        for cand in candidates:
            try:
                tokens = model.to_tokens(cand).to(device)
                if tokens.shape[1] < 4:
                    results.append({
                        "candidate": cand, "tokens": int(tokens.shape[1]),
                        "meanNll": 0.0, "perplexity": 0.0,
                        "prefixTokens": 0, "suffixTokens": 0,
                        "exactMatchTokens": 0, "exactMatchRatio": 0.0,
                        "generatedSuffix": "", "expectedSuffix": "",
                        "leakageScore": 0.0, "verdict": "too-short",
                    })
                    continue
                if tokens.shape[1] > 256:
                    tokens = tokens[:, :256]

                # 1) NLL / perplexity on full candidate
                with torch.no_grad():
                    logits_full = model(tokens)
                log_probs = torch.log_softmax(logits_full[0, :-1, :].float(), dim=-1)
                targets = tokens[0, 1:]
                token_lp = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
                mean_nll = -token_lp.mean().item()
                ppl = float(torch.tensor(mean_nll).exp().item())
                del logits_full, log_probs

                # 2) Exact-completion test
                T = tokens.shape[1]
                split = T // 2
                prefix_ids = tokens[:, :split]
                suffix_ids = tokens[0, split:]
                gen_len = suffix_ids.shape[0]
                cur = prefix_ids.clone()
                generated = []
                with torch.no_grad():
                    for _ in range(gen_len):
                        out = model(cur)
                        nxt = out[0, -1, :].argmax().item()
                        generated.append(nxt)
                        cur = torch.cat([cur, torch.tensor([[nxt]], device=device)], dim=1)
                exact_matches = sum(1 for a, b in zip(generated, suffix_ids.tolist()) if a == b)
                exact_ratio = exact_matches / max(1, gen_len)

                expected_suffix = model.to_string(suffix_ids)
                generated_suffix = model.to_string(torch.tensor(generated, device=device))

                # 3) Leakage score
                ppl_norm = max(0.0, min(1.0, 1.0 - ppl / max(baseline_ppl, 1e-9)))
                leakage = 0.5 * ppl_norm + 0.5 * exact_ratio
                if leakage >= 0.6:
                    verdict = "likely-memorized"
                elif leakage >= 0.35:
                    verdict = "suspicious"
                else:
                    verdict = "safe"

                results.append({
                    "candidate": cand,
                    "tokens": int(T),
                    "meanNll": round(mean_nll, 4),
                    "perplexity": round(ppl, 3),
                    "prefixTokens": int(split),
                    "suffixTokens": int(gen_len),
                    "exactMatchTokens": int(exact_matches),
                    "exactMatchRatio": round(exact_ratio, 4),
                    "generatedSuffix": generated_suffix,
                    "expectedSuffix": expected_suffix,
                    "leakageScore": round(leakage, 4),
                    "verdict": verdict,
                })
            except Exception as inner:
                print(f"[DataLeakageScanner] skip candidate: {inner}")
                results.append({
                    "candidate": cand, "tokens": 0, "meanNll": 0.0, "perplexity": 0.0,
                    "prefixTokens": 0, "suffixTokens": 0, "exactMatchTokens": 0, "exactMatchRatio": 0.0,
                    "generatedSuffix": "", "expectedSuffix": "",
                    "leakageScore": 0.0, "verdict": "error",
                })

        return jsonify({
            "success": True,
            "device": str(device),
            "baselinePerplexity": round(baseline_ppl, 3),
            "candidates": results,
        })
    except Exception as e:
        print(f"[DataLeakageScanner] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/multilingual-parity/analyze', methods=['POST'])
def multilingual_parity_analyze():
    """
    Compute per-language mean perplexity and a parity ratio.

    For each (language, prompts) pack:
      - For each prompt, tokenize, forward pass, gather log-prob of each next token (teacher-forcing),
        mean negative log-likelihood = -mean(log P(token_t | tokens_<t)) over t=1..T-1.
      - Perplexity = exp(mean_nll).
    parityRatio = max_lang_ppl / min_lang_ppl. 1.0 == perfect parity.

    Request:  { model, languages: [{language, prompts:[..]}] }
    Response: { success, device, perLanguage:[{language,prompts,meanNll,meanPerplexity,meanTokens}],
                parityRatio, bestLanguage, worstLanguage, totalPrompts }
    """
    try:
        data = request.json or {}
        languages = data.get('languages') or []
        if len(languages) < 2:
            return jsonify({"error": "at least 2 languages required"}), 400
        model_id = data.get('model', 'gpt2-small')
        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)
        device = next(model.parameters()).device

        per_language = []
        total_prompts = 0
        for pack in languages[:12]:
            lang = (pack.get('language') or '').strip()
            prompts = [p for p in (pack.get('prompts') or []) if isinstance(p, str) and p.strip()][:30]
            if not lang or not prompts:
                continue
            nlls = []
            tok_counts = []
            for p in prompts:
                try:
                    tokens = model.to_tokens(p).to(device)
                    if tokens.shape[1] < 2:
                        continue
                    if tokens.shape[1] > 256:
                        tokens = tokens[:, :256]
                    with torch.no_grad():
                        logits = model(tokens)  # [1, T, V]
                    # Teacher-forcing nll on positions 1..T-1
                    log_probs = torch.log_softmax(logits[0, :-1, :].float(), dim=-1)  # [T-1, V]
                    targets = tokens[0, 1:]  # [T-1]
                    token_lp = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)  # [T-1]
                    nll = -token_lp.mean().item()
                    nlls.append(nll)
                    tok_counts.append(tokens.shape[1] - 1)
                    total_prompts += 1
                    del logits, log_probs
                except Exception as inner:
                    print(f"[MultilingualParity] skip {lang!r} prompt: {inner}")
            if not nlls:
                continue
            mean_nll = sum(nlls) / len(nlls)
            mean_tok = sum(tok_counts) / len(tok_counts)
            per_language.append({
                "language": lang,
                "prompts": len(nlls),
                "meanNll": round(mean_nll, 4),
                "meanPerplexity": round(float(torch.tensor(mean_nll).exp().item()), 3),
                "meanTokens": round(mean_tok, 1),
            })

        if len(per_language) < 2:
            return jsonify({"error": "fewer than 2 languages produced valid results"}), 400

        ppls = [(l['language'], l['meanPerplexity']) for l in per_language]
        best = min(ppls, key=lambda x: x[1])
        worst = max(ppls, key=lambda x: x[1])
        parity = worst[1] / max(best[1], 1e-9)

        return jsonify({
            "success": True,
            "device": str(device),
            "perLanguage": per_language,
            "parityRatio": round(parity, 3),
            "bestLanguage": best[0],
            "worstLanguage": worst[0],
            "totalPrompts": total_prompts,
        })
    except Exception as e:
        print(f"[MultilingualParity] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/latency-profiler/analyze', methods=['POST'])
def latency_profiler_analyze():
    """
    Profile end-to-end and per-block forward-pass latency.

    Pipeline:
    1. Tokenize prompt; do 1 warmup pass (always discarded).
    2. For each trial:
       a. Time prefill (forward on full prompt) with cuda.synchronize() if CUDA.
       b. Greedily generate `generatedTokens` extra tokens; time the whole loop.
       c. Hook every `blocks.{l}` to measure per-block time on a single forward pass:
          install pre-/post-hooks recording perf_counter; do one extra forward.
    3. Aggregate: mean ± std end-to-end ms, mean per-layer ms, per-stage breakdown
       (prefill / generation / per-token-mean / overhead).

    Request:  { model, prompt, generatedTokens, trials }
    Response: { success, device, prompt, promptTokens, generatedTokens, trials,
                totalMeanMs, totalStdMs, tokensPerSec,
                perLayer:[{layer, meanMs, stdMs}],
                stages:[{stage, meanMs, pct}] }
    """
    try:
        import time
        data = request.json or {}
        prompt = (data.get('prompt') or '').strip()
        if not prompt:
            return jsonify({"error": "prompt is required"}), 400
        gen_tokens = int(data.get('generatedTokens') or 20)
        gen_tokens = max(1, min(50, gen_tokens))
        trials = int(data.get('trials') or 5)
        trials = max(1, min(15, trials))
        model_id = data.get('model', 'gpt2-small')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)
        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers
        is_cuda = device.type == 'cuda'

        def sync():
            if is_cuda:
                torch.cuda.synchronize()

        tokens = model.to_tokens(prompt).to(device)
        prompt_len = tokens.shape[1]
        if prompt_len > 64:
            return jsonify({"error": f"prompt too long ({prompt_len} tokens); max 64 for latency profiling"}), 400

        # Warmup (discarded).
        with torch.no_grad():
            _ = model(tokens)
        sync()

        prefill_times = []
        gen_times = []
        total_times = []

        for _ in range(trials):
            sync(); t0 = time.perf_counter()
            with torch.no_grad():
                _ = model(tokens)
            sync(); t1 = time.perf_counter()
            prefill_times.append((t1 - t0) * 1000.0)

            sync(); t0 = time.perf_counter()
            cur = tokens.clone()
            with torch.no_grad():
                for _ in range(gen_tokens):
                    out = model(cur)
                    next_id = out[0, -1, :].argmax().item()
                    cur = torch.cat([cur, torch.tensor([[next_id]], device=device)], dim=1)
            sync(); t1 = time.perf_counter()
            gen_times.append((t1 - t0) * 1000.0)
            total_times.append(prefill_times[-1] + gen_times[-1])

        # Per-layer profiling via hooks (single forward pass per trial).
        per_layer_acc = [[] for _ in range(n_layers)]
        for _ in range(trials):
            layer_starts = [0.0] * n_layers
            layer_durations = [0.0] * n_layers
            handles = []
            def make_pre(l_idx):
                def _pre(module, inputs):
                    sync()
                    layer_starts[l_idx] = time.perf_counter()
                return _pre
            def make_post(l_idx):
                def _post(module, inputs, outputs):
                    sync()
                    layer_durations[l_idx] = (time.perf_counter() - layer_starts[l_idx]) * 1000.0
                return _post
            for l in range(n_layers):
                block = model.blocks[l]
                handles.append(block.register_forward_pre_hook(make_pre(l)))
                handles.append(block.register_forward_hook(make_post(l)))
            with torch.no_grad():
                _ = model(tokens)
            for h in handles:
                h.remove()
            for l in range(n_layers):
                per_layer_acc[l].append(layer_durations[l])

        def mean_std(xs):
            if not xs:
                return 0.0, 0.0
            m = sum(xs) / len(xs)
            v = sum((x - m) ** 2 for x in xs) / max(1, len(xs))
            return m, v ** 0.5

        per_layer = []
        for l in range(n_layers):
            m, s = mean_std(per_layer_acc[l])
            per_layer.append({"layer": l, "meanMs": round(m, 4), "stdMs": round(s, 4)})

        total_mean, total_std = mean_std(total_times)
        prefill_mean, _ = mean_std(prefill_times)
        gen_mean, _ = mean_std(gen_times)
        per_token_mean = gen_mean / gen_tokens if gen_tokens > 0 else 0.0
        overhead = max(0.0, total_mean - prefill_mean - gen_mean)
        total_for_pct = max(1e-9, total_mean)
        stages = [
            {"stage": "prefill", "meanMs": round(prefill_mean, 3), "pct": round(prefill_mean / total_for_pct * 100, 2)},
            {"stage": "generation total", "meanMs": round(gen_mean, 3), "pct": round(gen_mean / total_for_pct * 100, 2)},
            {"stage": "per-gen-token avg", "meanMs": round(per_token_mean, 3), "pct": round(per_token_mean / total_for_pct * 100, 2)},
            {"stage": "overhead", "meanMs": round(overhead, 3), "pct": round(overhead / total_for_pct * 100, 2)},
        ]

        tokens_per_sec = (gen_tokens / (gen_mean / 1000.0)) if gen_mean > 0 else 0.0

        return jsonify({
            "success": True,
            "device": str(device),
            "prompt": prompt,
            "promptTokens": prompt_len,
            "generatedTokens": gen_tokens,
            "trials": trials,
            "totalMeanMs": round(total_mean, 3),
            "totalStdMs": round(total_std, 3),
            "tokensPerSec": round(tokens_per_sec, 2),
            "perLayer": per_layer,
            "stages": stages,
        })
    except Exception as e:
        print(f"[LatencyProfiler] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/drift-monitor/analyze', methods=['POST'])
def drift_monitor_analyze():
    """
    Measure representational drift between two prompt distributions (e.g., baseline vs canary).

    Algorithm:
    1. For each prompt in set A and set B, run forward pass with cache; per layer take the
       last-token residual stream (a sentence summary vector). Average within set → mean_a[l], mean_b[l].
    2. Per-layer drift = cosine_distance(mean_a[l], mean_b[l]) and L2 distance.
    3. Output drift = JS divergence between the mean next-token softmax of set A and set B.
    4. Top shifted tokens = top-K tokens by |mean_prob_a - mean_prob_b|.

    Request:  { model, promptsA, promptsB, labelA, labelB }
    Response: { success, labelA, labelB, nA, nB, nLayers,
                perLayerCosineDistance[], perLayerL2Distance[],
                meanCosineDistance, outputJSDivergence,
                topShiftedTokens:[{token, deltaProb, probA, probB}] }
    """
    try:
        data = request.json or {}
        prompts_a = [p for p in (data.get('promptsA') or []) if isinstance(p, str) and p.strip()][:50]
        prompts_b = [p for p in (data.get('promptsB') or []) if isinstance(p, str) and p.strip()][:50]
        if not prompts_a or not prompts_b:
            return jsonify({"error": "Need at least 1 prompt in each set"}), 400
        label_a = (data.get('labelA') or 'A')[:64]
        label_b = (data.get('labelB') or 'B')[:64]
        model_id = data.get('model', 'gpt2-small')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)
        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers
        d_model = model.cfg.d_model

        def encode_set(prompts):
            """Returns (per_layer_mean[L, d_model], mean_probs[vocab], per_prompt_probs[N, vocab])."""
            sums = [torch.zeros(d_model, device=device) for _ in range(n_layers)]
            counts = 0
            prob_sum = None
            per_prompt_probs = []  # store each prompt's softmax for per-prompt JS
            for p in prompts:
                try:
                    tokens = model.to_tokens(p).to(device)
                    if tokens.shape[1] > 128:
                        tokens = tokens[:, :128]
                    with torch.no_grad():
                        logits, cache = model.run_with_cache(tokens)
                    for l in range(n_layers):
                        last_resid = cache[f'blocks.{l}.hook_resid_post'][0, -1, :]
                        sums[l] = sums[l] + last_resid
                    probs = torch.softmax(logits[0, -1, :].float(), dim=-1)
                    per_prompt_probs.append(probs.detach().clone())
                    prob_sum = probs if prob_sum is None else prob_sum + probs
                    counts += 1
                    del logits, cache
                except Exception as inner:
                    print(f"[DriftMonitor] skip prompt: {inner}")
                    continue
            if counts == 0:
                return None, None, None, 0
            means = torch.stack([s / counts for s in sums], dim=0)  # [L, d]
            mean_probs = prob_sum / counts
            return means, mean_probs, per_prompt_probs, counts

        mean_a, probs_a, prompts_probs_a, n_a = encode_set(prompts_a)
        mean_b, probs_b, prompts_probs_b, n_b = encode_set(prompts_b)
        if mean_a is None or mean_b is None:
            return jsonify({"error": "All prompts failed to encode"}), 400

        cos = torch.nn.functional.cosine_similarity(mean_a, mean_b, dim=-1)  # [L]
        cos_dist = (1.0 - cos).cpu().tolist()
        l2 = torch.norm(mean_a - mean_b, dim=-1).cpu().tolist()
        mean_cos = sum(cos_dist) / len(cos_dist)

        eps = 1e-12

        def js_div(p, q):
            m = 0.5 * (p + q)
            kl_pm = (p * (p.clamp_min(eps).log() - m.clamp_min(eps).log())).sum().item()
            kl_qm = (q * (q.clamp_min(eps).log() - m.clamp_min(eps).log())).sum().item()
            return 0.5 * (kl_pm + kl_qm)

        # JS of mean softmaxes — kept for reference but uninformative when
        # averaging many peaked distributions (both means converge to the English
        # unigram marginal regardless of topic, so this is near-zero by design).
        js_of_means = js_div(probs_a, probs_b)

        # MEANINGFUL drift metric: average per-prompt JS divergence to the OTHER
        # set's centroid. Because each prompt's softmax is sharply peaked on a
        # different token, this captures real distributional drift.
        per_prompt_js_a = [js_div(p, probs_b) for p in prompts_probs_a]
        per_prompt_js_b = [js_div(p, probs_a) for p in prompts_probs_b]
        all_js = per_prompt_js_a + per_prompt_js_b
        js = sum(all_js) / len(all_js) if all_js else 0.0

        # Top shifted tokens.
        delta = (probs_a - probs_b).cpu()
        top_k = 12
        abs_delta = delta.abs()
        top_idx = torch.topk(abs_delta, top_k).indices.tolist()
        top_shifted = []
        for tid in top_idx:
            try:
                tok_str = model.to_string(torch.tensor([tid]))
            except Exception:
                tok_str = f"[id={tid}]"
            top_shifted.append({
                "token": tok_str,
                "deltaProb": round(float(delta[tid].item()), 6),
                "probA": round(float(probs_a[tid].item()), 6),
                "probB": round(float(probs_b[tid].item()), 6),
            })

        return jsonify({
            "success": True,
            "labelA": label_a,
            "labelB": label_b,
            "nA": n_a,
            "nB": n_b,
            "nLayers": n_layers,
            "perLayerCosineDistance": [round(x, 6) for x in cos_dist],
            "perLayerL2Distance": [round(x, 4) for x in l2],
            "meanCosineDistance": round(mean_cos, 6),
            "outputJSDivergence": round(float(js), 6),
            "outputJSDivergenceOfMeans": round(float(js_of_means), 6),
            "topShiftedTokens": top_shifted,
        })
    except Exception as e:
        print(f"[DriftMonitor] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/robustness-testing/analyze', methods=['POST'])
def robustness_testing_analyze():
    """
    Measure representational stability under input perturbations.

    Algorithm:
    1. Run clean prompt through HookedTransformer with cache; record per-layer
       residual-stream activations (resid_post for each block).
    2. For k = 1..nSamples: generate a perturbed version of the prompt
       (char-noise / swap-adjacent / delete) at noiseRate, run with cache.
    3. For each layer, compute mean cosine distance between clean and perturbed
       residual streams (averaged over token positions where both are valid).
       Take min(clean_len, pert_len) positions.
    4. Compute per-sample mean layer distance and output-KL between clean and
       perturbed final next-token softmax.
    5. Aggregate to perLayerMeanDistance and meanOverallDistance.

    Request:  { model, prompt, perturbationType, noiseRate, nSamples }
    Response: { success, prompt, perturbationType, noiseRate, nSamples, nLayers,
                perLayerMeanDistance[], meanOverallDistance, meanOutputKL,
                samples[{perturbed, meanLayerDistance, outputKL}] }
    """
    try:
        import random
        import string
        data = request.json or {}
        prompt = (data.get('prompt') or '').strip()
        if not prompt:
            return jsonify({"error": "prompt is required"}), 400
        perturbation_type = data.get('perturbationType', 'char')
        if perturbation_type not in ('char', 'swap', 'delete'):
            perturbation_type = 'char'
        noise_rate = float(data.get('noiseRate') or 0.1)
        noise_rate = max(0.01, min(0.5, noise_rate))
        n_samples = int(data.get('nSamples') or 5)
        n_samples = max(1, min(20, n_samples))
        model_id = data.get('model', 'gpt2-small')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)

        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers

        def perturb(s: str, ptype: str, rate: float) -> str:
            chars = list(s)
            n_changes = max(1, int(len(chars) * rate))
            indices = random.sample(range(len(chars)), min(n_changes, len(chars)))
            if ptype == 'char':
                for i in indices:
                    chars[i] = random.choice(string.ascii_lowercase + ' ')
            elif ptype == 'swap':
                for i in indices:
                    if i + 1 < len(chars):
                        chars[i], chars[i + 1] = chars[i + 1], chars[i]
            elif ptype == 'delete':
                for i in sorted(indices, reverse=True):
                    if len(chars) > 1:
                        del chars[i]
            return ''.join(chars)

        # Clean forward pass with cache for residual streams.
        clean_tokens = model.to_tokens(prompt).to(device)
        if clean_tokens.shape[1] > 128:
            return jsonify({"error": f"prompt too long ({clean_tokens.shape[1]} tokens); max 128 for robustness testing"}), 400

        with torch.no_grad():
            clean_logits, clean_cache = model.run_with_cache(clean_tokens)
        clean_resids = [clean_cache[f'blocks.{l}.hook_resid_post'][0] for l in range(n_layers)]  # each [T, d]
        clean_T = clean_resids[0].shape[0]
        clean_probs = torch.softmax(clean_logits[0, -1, :].float(), dim=-1)

        per_layer_sums = [0.0] * n_layers
        per_layer_counts = [0] * n_layers
        sample_results = []
        all_kls = []

        for s_idx in range(n_samples):
            perturbed = perturb(prompt, perturbation_type, noise_rate)
            try:
                pert_tokens = model.to_tokens(perturbed).to(device)
                with torch.no_grad():
                    pert_logits, pert_cache = model.run_with_cache(pert_tokens)
                pert_T = pert_tokens.shape[1]
                T = min(clean_T, pert_T)

                sample_layer_dists = []
                for l in range(n_layers):
                    pert_resid = pert_cache[f'blocks.{l}.hook_resid_post'][0][:T]
                    clean_resid = clean_resids[l][:T]
                    cos_sim = torch.nn.functional.cosine_similarity(clean_resid, pert_resid, dim=-1)
                    cos_dist = (1.0 - cos_sim).mean().item()
                    per_layer_sums[l] += cos_dist
                    per_layer_counts[l] += 1
                    sample_layer_dists.append(cos_dist)

                pert_probs = torch.softmax(pert_logits[0, -1, :].float(), dim=-1)
                kl = torch.nn.functional.kl_div(
                    pert_probs.clamp_min(1e-12).log(),
                    clean_probs,
                    reduction='sum',
                ).item()
                all_kls.append(kl)
                mean_layer_dist = sum(sample_layer_dists) / len(sample_layer_dists)
                sample_results.append({
                    "perturbed": perturbed,
                    "meanLayerDistance": round(mean_layer_dist, 6),
                    "outputKL": round(float(kl), 6),
                })
                del pert_logits, pert_cache
            except Exception as inner:
                print(f"[RobustnessTesting] sample {s_idx} skipped: {inner}")
                continue

        per_layer_mean = [
            round(per_layer_sums[l] / per_layer_counts[l], 6) if per_layer_counts[l] > 0 else 0.0
            for l in range(n_layers)
        ]
        mean_overall = sum(per_layer_mean) / len(per_layer_mean) if per_layer_mean else 0.0
        mean_kl = sum(all_kls) / len(all_kls) if all_kls else 0.0

        return jsonify({
            "success": True,
            "prompt": prompt,
            "perturbationType": perturbation_type,
            "noiseRate": noise_rate,
            "nSamples": len(sample_results),
            "nLayers": n_layers,
            "perLayerMeanDistance": per_layer_mean,
            "meanOverallDistance": round(mean_overall, 6),
            "meanOutputKL": round(mean_kl, 6),
            "samples": sample_results,
        })
    except Exception as e:
        print(f"[RobustnessTesting] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/bias-probing/analyze', methods=['POST'])
def bias_probing_analyze():
    """
    Behavioral bias probe via paired-prompt next-token logit/prob deltas.

    For each pair {promptA, promptB} and each target token (first sub-token):
      - Compute next-token logit and softmax prob at last position of promptA → (logitA, probA).
      - Same for promptB → (logitB, probB).
      - deltaProb = probA - probB; deltaLogit = logitA - logitB.
    Aggregate per-target across pairs (mean deltaProb, mean |deltaProb|, mean deltaLogit)
    and compute global mean |deltaProb|.

    Request:  { model, pairs:[{promptA,promptB}], targets:[str], labelA, labelB }
    Response: { success, labelA, labelB, total, meanAbsDeltaProb,
                perTarget:[{target, meanDeltaProb, meanAbsDeltaProb, meanDeltaLogit}],
                pairs:[{pairIndex, promptA, promptB, results:[{target,probA,probB,deltaProb,logitA,logitB,deltaLogit}]}] }
    """
    try:
        data = request.json or {}
        pairs_in = data.get('pairs') or []
        targets_in = data.get('targets') or []
        if not isinstance(pairs_in, list) or len(pairs_in) < 1 or not isinstance(targets_in, list) or len(targets_in) < 1:
            return jsonify({"error": "Need at least 1 pair and 1 target"}), 400
        pairs_in = pairs_in[:50]
        targets_in = targets_in[:50]
        label_a = (data.get('labelA') or 'A')[:64]
        label_b = (data.get('labelB') or 'B')[:64]

        model_id = data.get('model', 'gpt2-small')
        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)

        # Resolve target tokens (first sub-token of each).
        target_specs = []
        for t in targets_in:
            try:
                ids = model.to_tokens(t, prepend_bos=False)[0]
                if ids.numel() == 0:
                    continue
                tid = int(ids[0].item())
                target_specs.append({"text": t, "id": tid})
            except Exception:
                continue
        if not target_specs:
            return jsonify({"error": "No targets could be tokenized"}), 400

        # Cache distributions per unique prompt to avoid duplicate forward passes.
        prompt_cache = {}
        def get_dist(prompt: str):
            if prompt in prompt_cache:
                return prompt_cache[prompt]
            tokens = model.to_tokens(prompt)
            with torch.no_grad():
                logits = model(tokens)
            last_logits = logits[0, -1, :]
            probs = torch.softmax(last_logits, dim=-1)
            prompt_cache[prompt] = (last_logits, probs)
            return prompt_cache[prompt]

        pair_results = []
        per_target_acc = {ts["text"]: {"deltaProb": [], "deltaLogit": []} for ts in target_specs}
        all_abs_deltas = []

        for pi, p in enumerate(pairs_in):
            pa = (p.get('promptA') or '').strip()
            pb = (p.get('promptB') or '').strip()
            if not pa or not pb:
                continue
            try:
                logits_a, probs_a = get_dist(pa)
                logits_b, probs_b = get_dist(pb)
            except Exception as inner:
                print(f"[BiasProbing] skip pair {pi}: {inner}")
                continue

            results = []
            for ts in target_specs:
                tid = ts["id"]
                la = float(logits_a[tid].item())
                lb = float(logits_b[tid].item())
                pa_prob = float(probs_a[tid].item())
                pb_prob = float(probs_b[tid].item())
                d_prob = pa_prob - pb_prob
                d_logit = la - lb
                results.append({
                    "target": ts["text"],
                    "probA": round(pa_prob, 6),
                    "probB": round(pb_prob, 6),
                    "deltaProb": round(d_prob, 6),
                    "logitA": round(la, 4),
                    "logitB": round(lb, 4),
                    "deltaLogit": round(d_logit, 4),
                })
                per_target_acc[ts["text"]]["deltaProb"].append(d_prob)
                per_target_acc[ts["text"]]["deltaLogit"].append(d_logit)
                all_abs_deltas.append(abs(d_prob))

            pair_results.append({
                "pairIndex": pi,
                "promptA": pa,
                "promptB": pb,
                "results": results,
            })

        per_target = []
        for ts in target_specs:
            arr = per_target_acc[ts["text"]]["deltaProb"]
            arrl = per_target_acc[ts["text"]]["deltaLogit"]
            if not arr:
                continue
            per_target.append({
                "target": ts["text"],
                "meanDeltaProb": round(sum(arr) / len(arr), 6),
                "meanAbsDeltaProb": round(sum(abs(x) for x in arr) / len(arr), 6),
                "meanDeltaLogit": round(sum(arrl) / len(arrl), 4),
            })

        total = sum(len(p["results"]) for p in pair_results)
        mean_abs = sum(all_abs_deltas) / len(all_abs_deltas) if all_abs_deltas else 0.0

        return jsonify({
            "success": True,
            "labelA": label_a,
            "labelB": label_b,
            "total": total,
            "meanAbsDeltaProb": round(mean_abs, 6),
            "perTarget": per_target,
            "pairs": pair_results,
        })
    except Exception as e:
        print(f"[BiasProbing] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/gradient-attribution/analyze', methods=['POST'])
def gradient_attribution_analyze():
    """
    Integrated-gradients attribution of a target token's logit back to each
    input token's embedding.

    Algorithm:
    1. Tokenize prompt; capture input embeddings via hook_embed.
    2. Resolve target_token_id (user-supplied first token of `target`, else
       argmax of next-token distribution).
    3. For k = 1..steps: alpha = k/steps; replace embeddings with
       (alpha * embeds) using a forward hook; compute target_logit at last
       position; backprop to get d(target_logit)/d(scaled_embeds); accumulate.
    4. avg_grad = sum / steps. Integrated grad attribution = embeds * avg_grad
       (zero baseline). Per-token signed attribution = sum over embed dim.
       Saliency = abs sum.

    Request:  { model, prompt, target?, steps }
    Response: { success, prompt, targetToken, targetProb, topPredicted{token,prob},
                steps, tokens:[{position,token,attribution,saliency}], maxAbsAttribution }
    """
    try:
        data = request.json or {}
        prompt = (data.get('prompt') or '').strip()
        if not prompt:
            return jsonify({"error": "prompt is required"}), 400
        steps = int(data.get('steps') or 20)
        steps = max(1, min(50, steps))
        model_id = data.get('model', 'gpt2-small')
        user_target = data.get('target')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)

        device = next(model.parameters()).device
        model.eval()

        tokens = model.to_tokens(prompt).to(device)
        seq_len = tokens.shape[1]
        if seq_len > 64:
            return jsonify({"error": f"prompt too long ({seq_len} tokens); max 64 for gradient attribution"}), 400
        token_strs = [model.to_string([int(t)]) for t in tokens[0]]

        # Capture clean embeddings
        clean_emb = {}
        def cap_hook(act, hook):
            clean_emb['e'] = act.detach().clone()
            return act
        with torch.no_grad():
            logits_clean = model.run_with_hooks(tokens, fwd_hooks=[('hook_embed', cap_hook)])
        last_logits = logits_clean[0, -1, :]
        probs = torch.softmax(last_logits, dim=-1)
        top_id = int(torch.argmax(probs).item())
        top_prob = float(probs[top_id].item())
        top_str = model.to_string([top_id])

        if user_target is not None and len(user_target) > 0:
            tgt_ids = model.to_tokens(user_target, prepend_bos=False)[0]
            if tgt_ids.numel() == 0:
                return jsonify({"error": "target tokenized to empty sequence"}), 400
            target_id = int(tgt_ids[0].item())
        else:
            target_id = top_id
        target_str = model.to_string([target_id])
        target_prob = float(probs[target_id].item())

        embeds = clean_emb['e']  # [1, T, d_model]

        # Integrated gradients along straight-line path from zero baseline to embeds.
        accum = torch.zeros_like(embeds)
        for k in range(1, steps + 1):
            alpha = k / steps
            # Build a fresh tensor each step that requires grad.
            scaled = (alpha * embeds).clone().detach().requires_grad_(True)
            captured_for_grad = [scaled]
            def replace_hook(act, hook):
                return captured_for_grad[0]
            logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[('hook_embed', replace_hook)],
            )
            target_logit = logits[0, -1, target_id]
            grad = torch.autograd.grad(target_logit, scaled, retain_graph=False, create_graph=False)[0]
            accum = accum + grad.detach()
            del scaled, logits, target_logit, grad

        avg_grad = accum / steps
        ig = embeds * avg_grad  # [1, T, d_model]
        per_token_signed = ig.sum(dim=-1)[0]  # [T]
        per_token_saliency = ig.abs().sum(dim=-1)[0]  # [T]

        per_token_signed_l = per_token_signed.detach().cpu().tolist()
        per_token_saliency_l = per_token_saliency.detach().cpu().tolist()
        max_abs = max((abs(x) for x in per_token_signed_l), default=1.0)

        token_results = []
        for i in range(seq_len):
            token_results.append({
                "position": i,
                "token": token_strs[i],
                "attribution": round(per_token_signed_l[i], 6),
                "saliency": round(per_token_saliency_l[i], 6),
            })

        return jsonify({
            "success": True,
            "prompt": prompt,
            "targetToken": target_str,
            "targetProb": round(target_prob, 6),
            "topPredicted": {"token": top_str, "prob": round(top_prob, 6)},
            "steps": steps,
            "tokens": token_results,
            "maxAbsAttribution": round(float(max_abs), 6),
        })
    except Exception as e:
        print(f"[GradientAttribution] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/confidence-calibration/analyze', methods=['POST'])
def confidence_calibration_analyze():
    """
    Compute Expected Calibration Error (ECE) and reliability bins for a model.

    Algorithm:
    1. For each {prompt, target} pair:
       - Forward pass on prompt; take next-token distribution at last position.
       - target_first_token_id = first token of `target` when tokenized.
       - predictedToken = argmax of distribution; predictedProb = its prob.
       - targetProb = prob assigned to target_first_token_id.
       - correct = (predictedToken_id == target_first_token_id)
    2. Bin examples by predictedProb into nBins equal-width bins on [0,1].
       For each bin: avgConfidence = mean(predictedProb), accuracy = mean(correct).
    3. ECE = sum_bins (n_bin / N) * |avgConfidence - accuracy|

    Request:  { model, pairs: [{prompt, target}], nBins }
    Response: { success, total, ece, meanConfidence, accuracy, bins, examples }
    """
    try:
        data = request.json or {}
        pairs = data.get('pairs') or []
        if not isinstance(pairs, list) or len(pairs) < 2:
            return jsonify({"error": "pairs must be a list of at least 2 {prompt, target} objects"}), 400
        pairs = pairs[:100]

        n_bins = int(data.get('nBins') if data.get('nBins') is not None else 10)
        n_bins = max(2, min(50, n_bins))

        model_id = data.get('model', 'gpt2-small')

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)

        examples = []
        skipped = 0
        for p in pairs:
            prompt = (p.get('prompt') or '').strip()
            target = p.get('target') or ''
            if not prompt or not target:
                skipped += 1
                continue
            try:
                target_ids = model.to_tokens(target, prepend_bos=False)[0]
                if target_ids.numel() == 0:
                    skipped += 1
                    continue
                target_first_id = int(target_ids[0].item())
                target_first_str = model.to_string([target_first_id])

                prompt_tokens = model.to_tokens(prompt)
                with torch.no_grad():
                    logits = model(prompt_tokens)
                last_logits = logits[0, -1, :]
                probs = torch.softmax(last_logits, dim=-1)
                pred_id = int(torch.argmax(probs).item())
                pred_prob = float(probs[pred_id].item())
                target_prob = float(probs[target_first_id].item())
                pred_str = model.to_string([pred_id])

                examples.append({
                    "prompt": prompt,
                    "target": target,
                    "predictedToken": pred_str,
                    "predictedProb": round(pred_prob, 6),
                    "targetProb": round(target_prob, 6),
                    "correct": pred_id == target_first_id,
                })
            except Exception as inner:
                print(f"[ConfidenceCalibration] skipped pair due to error: {inner}")
                skipped += 1
                continue

        n = len(examples)
        if n == 0:
            return jsonify({"error": "No pairs could be scored"}), 400

        bins = []
        ece = 0.0
        for i in range(n_bins):
            low = i / n_bins
            high = (i + 1) / n_bins
            in_bin = [e for e in examples if (e["predictedProb"] > low or i == 0) and e["predictedProb"] <= high]
            count = len(in_bin)
            if count > 0:
                avg_conf = sum(e["predictedProb"] for e in in_bin) / count
                acc = sum(1 for e in in_bin if e["correct"]) / count
                ece += (count / n) * abs(avg_conf - acc)
            else:
                avg_conf = 0.0
                acc = 0.0
            bins.append({
                "binIndex": i,
                "binLow": round(low, 4),
                "binHigh": round(high, 4),
                "count": count,
                "avgConfidence": round(avg_conf, 6),
                "accuracy": round(acc, 6),
            })

        mean_conf = sum(e["predictedProb"] for e in examples) / n
        accuracy = sum(1 for e in examples if e["correct"]) / n

        return jsonify({
            "success": True,
            "total": n,
            "skipped": skipped,
            "ece": round(ece, 6),
            "meanConfidence": round(mean_conf, 6),
            "accuracy": round(accuracy, 6),
            "bins": bins,
            "examples": examples,
        })
    except Exception as e:
        print(f"[ConfidenceCalibration] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/hallucination-detector/analyze', methods=['POST'])
def hallucination_detector_analyze():
    """
    Detect overconfident hallucinations by comparing final-layer output probability
    against per-layer logit-lens support for the chosen token.

    Algorithm:
    1. Greedy-generate up to maxTokens new tokens from the prompt.
    2. For each generated token at position i:
       - p_out = softmax(W_U · ln_final(resid_final[i]))[t_i]
       - For each layer L: p_L = softmax(W_U · ln_final(resid_L[i]))[t_i]
       - support_score = mean(p_L for L in last 50% of layers)
       - divergence = p_out - support_score
       - flagged = (p_out >= confThreshold) and (support_score < supportThreshold)

    Request:
      { model, prompt, maxTokens, confThreshold, supportThreshold }
    Response:
      { success, prompt, generatedTokens, nLayers, analysis: [...], flaggedCount, meanDivergence }
    """
    try:
        data = request.json or {}
        prompt = data.get('prompt', '').strip()
        if not prompt:
            return jsonify({"error": "prompt is required"}), 400

        model_id = data.get('model', 'gpt2-small')
        max_tokens = int(data.get('maxTokens', 8))
        max_tokens = max(1, min(20, max_tokens))
        conf_threshold = float(data.get('confThreshold') if data.get('confThreshold') is not None else 0.6)
        support_threshold = float(data.get('supportThreshold') if data.get('supportThreshold') is not None else 0.3)

        if model_id in ACTIVATION_ONLY_MODELS:
            model = _load_model_only(model_id)
        else:
            model, _ = get_model_and_sae(model_id)

        n_layers = model.cfg.n_layers
        W_U = model.W_U
        b_U = model.b_U if hasattr(model, 'b_U') and model.b_U is not None else None

        prompt_tokens = model.to_tokens(prompt)
        current_tokens = prompt_tokens
        prompt_len = current_tokens.shape[1]

        analysis = []
        generated_strs = []
        resid_filter = lambda name: name.endswith("hook_resid_post")

        for step in range(max_tokens):
            with torch.no_grad():
                _, cache = model.run_with_cache(current_tokens, names_filter=resid_filter)

            # Final layer prediction at the LAST position
            final_resid = cache[f"blocks.{n_layers - 1}.hook_resid_post"][0, -1, :]
            if hasattr(model, 'ln_final'):
                final_normed = model.ln_final(final_resid.unsqueeze(0)).squeeze(0)
            else:
                final_normed = final_resid
            final_logits = final_normed.float() @ W_U.float()
            if b_U is not None:
                final_logits = final_logits + b_U.float()
            final_probs = torch.softmax(final_logits, dim=-1)

            chosen_id = int(final_probs.argmax().item())
            p_out = float(final_probs[chosen_id].item())
            chosen_str = model.tokenizer.decode([chosen_id])

            # Per-layer logit lens prob for the chosen token
            layer_probs = []
            for layer in range(n_layers):
                resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1, :]
                if hasattr(model, 'ln_final'):
                    resid_normed = model.ln_final(resid.unsqueeze(0)).squeeze(0)
                else:
                    resid_normed = resid
                logits = resid_normed.float() @ W_U.float()
                if b_U is not None:
                    logits = logits + b_U.float()
                probs = torch.softmax(logits, dim=-1)
                p_layer = float(probs[chosen_id].item())
                if np.isnan(p_layer) or np.isinf(p_layer):
                    p_layer = 0.0
                layer_probs.append(round(p_layer, 4))

            del cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            # Support score = mean over last half of layers
            half = max(1, n_layers // 2)
            late_probs = layer_probs[-half:]
            support_score = float(sum(late_probs) / len(late_probs))
            divergence = p_out - support_score
            flagged = bool(p_out >= conf_threshold and support_score < support_threshold)

            analysis.append({
                "position": step,
                "token": chosen_str,
                "outputProb": round(p_out, 4),
                "supportScore": round(support_score, 4),
                "divergence": round(divergence, 4),
                "flagged": flagged,
                "layerProbs": layer_probs,
            })
            generated_strs.append(chosen_str)

            # Append chosen token and continue
            new_tok = torch.tensor([[chosen_id]], device=current_tokens.device, dtype=current_tokens.dtype)
            current_tokens = torch.cat([current_tokens, new_tok], dim=1)

        flagged_count = sum(1 for a in analysis if a["flagged"])
        mean_div = float(sum(a["divergence"] for a in analysis) / len(analysis)) if analysis else 0.0

        return jsonify({
            "success": True,
            "prompt": prompt,
            "generatedTokens": generated_strs,
            "nLayers": n_layers,
            "analysis": analysis,
            "flaggedCount": flagged_count,
            "meanDivergence": round(mean_div, 4),
        })
    except Exception as e:
        print(f"[HallucinationDetector] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500
@app.route('/layer-ablation/analyze', methods=['POST'])
def layer_ablation_analyze():
    """
    For a prompt + two candidate tokens (A, B), zero-ablate each layer's
    attention output AND each layer's MLP output (one at a time) and measure
    how much the logit difference between A and B drops.

    This is a coarser-grained version of Direct Logit Attribution — DLA
    decomposes the contribution of every individual head/MLP, while this
    measures the CAUSAL impact of removing each layer's component entirely
    (which captures effects beyond linear projection: downstream feedback,
    second-order effects, etc.)

    Cost: 2 * n_layers + 1 forward passes. On GPT-2 small that's 25 forwards.

    Request:  { prompt, tokenA, tokenB }
    Response: { success, model, device, baseline:{logitA,logitB,logitDiff},
                layers:[{layer, attnAblated:{logitDiff, dropFromBaseline},
                                  mlpAblated:{logitDiff, dropFromBaseline}}] }
    """
    try:
        data = request.json or {}
        prompt = data.get('prompt')
        a = data.get('tokenA')
        b = data.get('tokenB')
        if not prompt or not a or not b:
            return jsonify({"error": "prompt, tokenA, tokenB required"}), 400

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        def first_id(s):
            ids = model.to_tokens(s, prepend_bos=False)[0]
            if ids.numel() == 0:
                raise ValueError(f"empty: {s!r}")
            return int(ids[0].item())

        a_id = first_id(a)
        b_id = first_id(b)

        def zero_hook(activation, hook):
            return torch.zeros_like(activation)

        with torch.no_grad():
            tokens = model.to_tokens(prompt)  # [1, seq] with BOS
            n_layers = int(model.cfg.n_layers)

            # Baseline
            base_logits = model(tokens)[0, -1, :].float().cpu()
            base_a = float(base_logits[a_id].item())
            base_b = float(base_logits[b_id].item())
            base_diff = base_a - base_b

            layer_results = []
            for layer in range(n_layers):
                # Attn ablation
                with model.hooks(fwd_hooks=[(f"blocks.{layer}.hook_attn_out", zero_hook)]):
                    attn_logits = model(tokens)[0, -1, :].float().cpu()
                attn_a = float(attn_logits[a_id].item())
                attn_b = float(attn_logits[b_id].item())
                attn_diff = attn_a - attn_b

                # MLP ablation
                with model.hooks(fwd_hooks=[(f"blocks.{layer}.hook_mlp_out", zero_hook)]):
                    mlp_logits = model(tokens)[0, -1, :].float().cpu()
                mlp_a = float(mlp_logits[a_id].item())
                mlp_b = float(mlp_logits[b_id].item())
                mlp_diff = mlp_a - mlp_b

                layer_results.append({
                    "layer": layer,
                    "attnAblated": {
                        "logitA": round(attn_a, 4),
                        "logitB": round(attn_b, 4),
                        "logitDiff": round(attn_diff, 4),
                        "dropFromBaseline": round(base_diff - attn_diff, 4),
                    },
                    "mlpAblated": {
                        "logitA": round(mlp_a, 4),
                        "logitB": round(mlp_b, 4),
                        "logitDiff": round(mlp_diff, 4),
                        "dropFromBaseline": round(base_diff - mlp_diff, 4),
                    },
                })

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "nLayers": n_layers,
            "tokenA": {"input": a, "used": model.to_string([a_id]), "id": a_id},
            "tokenB": {"input": b, "used": model.to_string([b_id]), "id": b_id},
            "baseline": {
                "logitA": round(base_a, 4),
                "logitB": round(base_b, 4),
                "logitDiff": round(base_diff, 4),
            },
            "layers": layer_results,
        })
    except Exception as e:
        print(f"[LayerAblation] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/logit-difference/analyze', methods=['POST'])
def logit_difference_analyze():
    """
    Compute the logit difference between two candidate next-tokens for a given
    prompt. This is THE foundational metric used in mech-interp literature
    (Wang et al. 2022 IOI, Anthropic transformer-circuits), because:
      - It's a clean scalar per prompt
      - Differences cancel out the unembedding bias
      - Sign tells you which token wins, magnitude tells you by how much

    Returns logits, probabilities, ranks, and the difference for both tokens
    at the FINAL position (i.e., the next-token prediction).

    Request:  { prompt, tokenA, tokenB }   // tokenA/B are literal strings
    Response: { success, model, device, prompt, promptTokens, lastTokenIdx,
                tokenA: {used, id, logit, prob, rank}, tokenB: {...},
                logitDiff, probDiff, top5Predictions:[{token, logit, prob, rank}] }
    """
    try:
        data = request.json or {}
        prompt = data.get('prompt')
        a = data.get('tokenA')
        b = data.get('tokenB')
        if not prompt or not a or not b:
            return jsonify({"error": "prompt, tokenA, tokenB required"}), 400

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        def first_id(s):
            ids = model.to_tokens(s, prepend_bos=False)[0]
            if ids.numel() == 0:
                raise ValueError(f"token tokenized to empty: {s!r}")
            return int(ids[0].item()), ids.numel() > 1, model.to_string([int(ids[0].item())])

        a_id, a_multi, a_used = first_id(a)
        b_id, b_multi, b_used = first_id(b)

        with torch.no_grad():
            tokens = model.to_tokens(prompt)  # [1, seq] with BOS
            logits = model(tokens)  # [1, seq, vocab]
            last_logits = logits[0, -1, :].float().cpu()  # [vocab]
            probs = torch.softmax(last_logits, dim=-1)
            # Ranks: argsort descending then find positions
            sorted_idx = torch.argsort(last_logits, descending=True)
            rank_map = torch.empty_like(sorted_idx)
            rank_map[sorted_idx] = torch.arange(len(sorted_idx))
            # Top-5 alternatives
            top_logits, top_idx = torch.topk(last_logits, 5)

        prompt_token_strs = [model.to_string([int(t.item())]) for t in tokens[0]]
        last_idx = int(tokens.shape[1] - 1)

        a_logit = float(last_logits[a_id].item())
        b_logit = float(last_logits[b_id].item())
        a_prob = float(probs[a_id].item())
        b_prob = float(probs[b_id].item())
        a_rank = int(rank_map[a_id].item()) + 1  # 1-indexed
        b_rank = int(rank_map[b_id].item()) + 1

        top5 = []
        for log_val, i in zip(top_logits.tolist(), top_idx.tolist()):
            r = int(rank_map[i].item()) + 1
            top5.append({
                "token": model.to_string([int(i)]),
                "tokenId": int(i),
                "logit": round(float(log_val), 4),
                "prob": round(float(probs[int(i)].item()), 5),
                "rank": r,
            })

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "prompt": prompt,
            "promptTokens": prompt_token_strs,
            "lastTokenIdx": last_idx,
            "tokenA": {
                "input": a, "used": a_used, "id": a_id, "multi": a_multi,
                "logit": round(a_logit, 4), "prob": round(a_prob, 5), "rank": a_rank,
            },
            "tokenB": {
                "input": b, "used": b_used, "id": b_id, "multi": b_multi,
                "logit": round(b_logit, 4), "prob": round(b_prob, 5), "rank": b_rank,
            },
            "logitDiff": round(a_logit - b_logit, 4),
            "probDiff": round(a_prob - b_prob, 5),
            "winner": "A" if a_logit > b_logit else "B",
            "top5Predictions": top5,
        })
    except Exception as e:
        print(f"[LogitDifference] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/vector-arithmetic/analyze', methods=['POST'])
def vector_arithmetic_analyze():
    """
    Compute v = E(a) - E(b) + E(c) and return the K nearest vocab tokens
    to v by cosine similarity in GPT-2's input embedding space.

    The classic word2vec-style analogy probe: e.g.
      " king" - " man" + " woman" -> ?
      " Paris" - " France" + " Germany" -> ?
      " walking" - " walk" + " run" -> ?

    We exclude the three input tokens themselves from the results so the
    user sees only NEW candidates.

    Request:  { a, b, c, k? }   // strings, e.g. " king", " man", " woman"
    Response: { success, model, device, dModel, vocabSize,
                tokens:{a,b,c} with id+used_str, neighbors:[{token, tokenId, similarity}] }
    """
    try:
        data = request.json or {}
        a = data.get('a'); b = data.get('b'); c = data.get('c')
        if not a or not b or not c:
            return jsonify({"error": "a, b, c required"}), 400
        k = int(data.get('k', 15))
        k = max(1, min(40, k))

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        def first_id(s):
            ids = model.to_tokens(s, prepend_bos=False)[0]
            if ids.numel() == 0:
                raise ValueError(f"token tokenized to empty: {s!r}")
            return int(ids[0].item()), ids.numel() > 1

        a_id, a_multi = first_id(a)
        b_id, b_multi = first_id(b)
        c_id, c_multi = first_id(c)

        with torch.no_grad():
            W_E = model.W_E  # [vocab, d_model]
            v = (W_E[a_id].float() - W_E[b_id].float() + W_E[c_id].float())
            v_norm = v / (v.norm() + 1e-9)
            W_norm = W_E.float() / (W_E.float().norm(dim=-1, keepdim=True) + 1e-9)
            sims = (W_norm @ v_norm).cpu()

            # Pull more than k so we can filter the 3 input tokens out
            top_vals, top_idx = torch.topk(sims, k + 5)

        skip = {a_id, b_id, c_id}
        neighbors = []
        for s, i in zip(top_vals.tolist(), top_idx.tolist()):
            if int(i) in skip:
                continue
            neighbors.append({
                "token": model.to_string([int(i)]),
                "tokenId": int(i),
                "similarity": round(float(s), 5),
            })
            if len(neighbors) >= k:
                break

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "dModel": int(W_E.shape[1]),
            "vocabSize": int(W_E.shape[0]),
            "tokens": {
                "a": {"used": model.to_string([a_id]), "id": a_id, "multi": a_multi},
                "b": {"used": model.to_string([b_id]), "id": b_id, "multi": b_multi},
                "c": {"used": model.to_string([c_id]), "id": c_id, "multi": c_multi},
            },
            "neighbors": neighbors,
        })
    except Exception as e:
        print(f"[VectorArithmetic] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/vocab-neighbors/analyze', methods=['POST'])
def vocab_neighbors_analyze():
    """
    Find the K nearest vocab tokens to an input token in embedding space.

    No forward pass needed — we operate directly on GPT-2's input embedding
    matrix W_E (shape [vocab_size, d_model]). For the input token, we compute
    cosine similarity against every other vocab token and return the top-K.

    Useful for:
      - Exploring what GPT-2 "thinks" is similar (e.g., " king" → " queen", " prince")
      - Finding tokenization quirks (capitalization splits, leading-space variants)
      - Sanity-checking that semantic structure exists in raw embeddings

    Request:  { token, k? }  // token is the literal string, e.g. " king"
    Response: { success, model, device, inputToken, inputTokenId, dModel,
                vocabSize, neighbors:[{token, tokenId, similarity}] }
    """
    try:
        data = request.json or {}
        raw = data.get('token')
        if raw is None or raw == "":
            return jsonify({"error": "token required"}), 400
        k = int(data.get('k', 20))
        k = max(1, min(50, k))

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        # Tokenize. If multi-token, use the first piece and report the actual
        # piece used so the user understands what was queried.
        ids = model.to_tokens(raw, prepend_bos=False)[0]
        if ids.numel() == 0:
            return jsonify({"error": "token tokenized to empty sequence"}), 400
        tok_id = int(ids[0].item())
        used_str = model.to_string([tok_id])

        with torch.no_grad():
            W_E = model.W_E  # [vocab, d_model]
            v = W_E[tok_id].float()
            v_norm = v / (v.norm() + 1e-9)
            # Normalize all rows (vocab is ~50k for GPT-2 — fine on GPU)
            W_norm = W_E.float() / (W_E.float().norm(dim=-1, keepdim=True) + 1e-9)
            sims = (W_norm @ v_norm).cpu()  # [vocab]

            # Get top K+1 (will include the token itself at position 0)
            top_vals, top_idx = torch.topk(sims, k + 1)

        neighbors = []
        for s, i in zip(top_vals.tolist(), top_idx.tolist()):
            if int(i) == tok_id:
                continue  # skip self
            neighbors.append({
                "token": model.to_string([int(i)]),
                "tokenId": int(i),
                "similarity": round(float(s), 5),
            })
            if len(neighbors) >= k:
                break

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "inputToken": used_str,
            "inputTokenId": tok_id,
            "multiTokenInput": ids.numel() > 1,
            "originalPieces": [model.to_string([int(t.item())]) for t in ids] if ids.numel() > 1 else None,
            "dModel": int(W_E.shape[1]),
            "vocabSize": int(W_E.shape[0]),
            "neighbors": neighbors,
        })
    except Exception as e:
        print(f"[VocabNeighbors] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/token-surprisal/analyze', methods=['POST'])
def token_surprisal_analyze():
    """
    Per-token surprisal (negative log probability) under the model.

    For every token at position p (p >= 1), we compute -log P(token_p | tokens<p)
    using the model's logits at position p-1. Lower = the model expected this
    token; higher = the model was surprised by it.

    Also returns the model's top-1 expected token at each position so users
    can see what the model was predicting instead.

    Request:  { prompt }
    Response: { success, model, device, tokens, surprisals (bits),
                expectedTokens, predictedCorrectly, totalLogProb,
                avgSurprisal, perplexity, maxSurprisal }
    """
    try:
        data = request.json or {}
        prompt = (data.get('prompt') or "").strip()
        if not prompt:
            return jsonify({"error": "prompt required"}), 400
        if len(prompt) > 1500:
            prompt = prompt[:1500]

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        toks = model.to_tokens(prompt)
        # Cap for performance and rendering
        if toks.shape[1] > 80:
            toks = toks[:, :80]
        token_strs = [model.to_string([int(t.item())]) for t in toks[0]]
        n = toks.shape[1]

        with torch.no_grad():
            logits = model(toks)  # [1, n, vocab]
            log_probs = torch.log_softmax(logits[0].float(), dim=-1)  # [n, vocab]

        # For position p in 1..n-1, surprisal of toks[p] = -log_probs[p-1, toks[p]]
        # in bits = nats / ln(2)
        ln2 = float(torch.log(torch.tensor(2.0)).item())
        surprisals = [None]  # position 0 has no previous context
        expected_tokens = [None]
        predicted_correctly = [None]
        for p in range(1, n):
            target_id = int(toks[0, p].item())
            nats = -float(log_probs[p - 1, target_id].item())
            bits = nats / ln2
            top1_id = int(log_probs[p - 1].argmax().item())
            surprisals.append(round(bits, 5))
            expected_tokens.append(model.to_string([top1_id]))
            predicted_correctly.append(top1_id == target_id)

        valid_surprisals = [s for s in surprisals if s is not None]
        avg_surprisal = sum(valid_surprisals) / len(valid_surprisals) if valid_surprisals else 0.0
        perplexity = float(2.0 ** avg_surprisal) if avg_surprisal else 0.0
        max_surprisal = max(valid_surprisals) if valid_surprisals else 0.0
        total_log_prob = -sum(valid_surprisals) * ln2  # in nats, total log P(sequence)

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "tokens": token_strs,
            "surprisals": surprisals,
            "expectedTokens": expected_tokens,
            "predictedCorrectly": predicted_correctly,
            "avgSurprisal": round(avg_surprisal, 5),
            "perplexity": round(perplexity, 4),
            "maxSurprisal": round(max_surprisal, 5),
            "totalLogProb": round(total_log_prob, 5),
        })
    except Exception as e:
        print(f"[TokenSurprisal] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/attention-pattern/analyze', methods=['POST'])
def attention_pattern_analyze():
    """
    Visualize the attention pattern of a single (layer, head) on a prompt.

    Returns the post-softmax attention matrix [seq_len, seq_len] with the
    corresponding token labels for both axes, so the frontend can render a
    heatmap of how each query position attends to each key position for the
    chosen head.

    Request:  { prompt, layer, head }
    Response: { success, model, device, layer, head, nLayers, nHeads,
                tokens:[str], pattern:[[float]], maxAttention,
                topAttended:[{from, to, weight}] (top 10 off-diagonal pairs) }
    """
    try:
        data = request.json or {}
        prompt = (data.get('prompt') or "").strip()
        layer = int(data.get('layer', 5))
        head = int(data.get('head', 5))
        if not prompt:
            return jsonify({"error": "prompt required"}), 400
        if len(prompt) > 600:
            prompt = prompt[:600]

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers
        n_heads = model.cfg.n_heads
        layer = max(0, min(layer, n_layers - 1))
        head = max(0, min(head, n_heads - 1))

        toks = model.to_tokens(prompt)
        # Cap sequence length for visualization (40x40 matrix is the max useful render).
        if toks.shape[1] > 40:
            toks = toks[:, :40]
        token_strs = [model.to_string([int(t.item())]) for t in toks[0]]

        pattern_name = f'blocks.{layer}.attn.hook_pattern'
        with torch.no_grad():
            _, cache = model.run_with_cache(toks, names_filter=[pattern_name])
            patt = cache[pattern_name][0, head].float().cpu()  # [q, k]

        # Convert to lists, also find top off-diagonal attended pairs.
        # Skip the BOS attention sink (j==0) — attention heads use the BOS token
        # as a "do nothing" target when they have nothing relevant to attend to,
        # and that uninteresting behaviour dominates raw rankings.
        seq_len = patt.shape[0]
        rows = [[round(float(patt[i, j].item()), 5) for j in range(seq_len)] for i in range(seq_len)]
        pairs = []
        for i in range(seq_len):
            for j in range(seq_len):
                if i != j and j != 0:  # skip self and BOS sink
                    pairs.append((float(patt[i, j].item()), i, j))
        pairs.sort(reverse=True)
        top_attended = [{"fromIdx": p[1], "fromToken": token_strs[p[1]],
                          "toIdx": p[2], "toToken": token_strs[p[2]],
                          "weight": round(p[0], 5)} for p in pairs[:10]]

        del cache
        if device.type == "cuda":
            torch.cuda.empty_cache()

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "layer": layer,
            "head": head,
            "nLayers": n_layers,
            "nHeads": n_heads,
            "tokens": token_strs,
            "pattern": rows,
            "maxAttention": round(float(patt.max().item()), 5),
            "topAttended": top_attended,
        })
    except Exception as e:
        print(f"[AttentionPattern] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/direct-logit-attribution/analyze', methods=['POST'])
def direct_logit_attribution_analyze():
    """
    Direct Logit Attribution (DLA) — decompose the model's prediction at the
    final token into per-attention-head and per-MLP-layer contributions by
    projecting each component's output (at the last position) onto the
    unembedding direction for a target token.

    This is the standard mechanistic-interpretability technique for asking
    "which components actually caused the model to predict token X?".

    Algorithm:
      1. Tokenize the prompt and run forward with cache (z + mlp_out).
      2. Pick a target token (top-1 prediction by default, or user-supplied
         word — first BPE token taken).
      3. unembed_dir = W_U[:, target_token]
      4. For each layer l and head h:
           per_head_out[l,h] = z[l, last, h] @ W_O[l, h]   # [d_model]
           contribution[l,h] = per_head_out[l,h] · unembed_dir
      5. For each layer l:
           mlp_contribution[l] = mlp_out[l, last] · unembed_dir
      6. Return ranked head contributions + per-layer MLP contributions +
         per-layer aggregated attention contribution.

    Note: contributions are raw dot products (not LN-scaled), so absolute
    magnitudes don't equal final logit deltas exactly, but the relative
    ranking — which is what users want — is preserved.

    Request:  { prompt, targetToken? (string, optional) }
    Response: { success, model, device, prompt, targetToken, targetTokenId,
                targetLogit, targetProb, baselineTop5,
                headContributions:[{layer, head, contribution}],   // ranked desc
                mlpContributions:[{layer, contribution}],
                attnLayerContributions:[{layer, contribution}],
                totalAttn, totalMlp }
    """
    try:
        data = request.json or {}
        prompt = (data.get('prompt') or "").strip()
        target_word = (data.get('targetToken') or "").strip()
        if not prompt:
            return jsonify({"error": "prompt required"}), 400
        if len(prompt) > 1000:
            prompt = prompt[:1000]

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device

        toks = model.to_tokens(prompt)
        n_layers = model.cfg.n_layers
        n_heads = model.cfg.n_heads

        # Names we need: per-layer z (head outputs pre-W_O), mlp_out, ln_final scale (optional).
        z_names = [f'blocks.{l}.attn.hook_z' for l in range(n_layers)]
        mlp_names = [f'blocks.{l}.hook_mlp_out' for l in range(n_layers)]
        names = z_names + mlp_names

        with torch.no_grad():
            logits, cache = model.run_with_cache(toks, names_filter=names)
            last_logits = logits[0, -1]  # [vocab]
            top5_vals, top5_idx = torch.topk(last_logits.softmax(dim=-1), 5)

            # Choose target token
            if target_word:
                target_ids = model.to_tokens(target_word, prepend_bos=False)[0]
                target_id = int(target_ids[0].item()) if target_ids.numel() > 0 else int(top5_idx[0].item())
            else:
                target_id = int(top5_idx[0].item())

            unembed_dir = model.W_U[:, target_id].float()  # [d_model]

            head_contribs = torch.zeros(n_layers, n_heads, device=device)
            mlp_contribs = torch.zeros(n_layers, device=device)
            for l in range(n_layers):
                z_l = cache[z_names[l]][0, -1].float()  # [head, d_head]
                W_O_l = model.W_O[l].float()             # [head, d_head, d_model]
                # per-head output at last pos: [head, d_model]
                per_head = torch.einsum('h d, h d m -> h m', z_l, W_O_l)
                head_contribs[l] = per_head @ unembed_dir
                mlp_contribs[l] = cache[mlp_names[l]][0, -1].float() @ unembed_dir

            head_contribs_cpu = head_contribs.cpu()
            mlp_contribs_cpu = mlp_contribs.cpu()

            head_list = []
            for l in range(n_layers):
                for h in range(n_heads):
                    head_list.append({
                        "layer": l, "head": h,
                        "contribution": round(float(head_contribs_cpu[l, h].item()), 5),
                    })
            head_list.sort(key=lambda x: abs(x['contribution']), reverse=True)

            mlp_list = [{"layer": l, "contribution": round(float(mlp_contribs_cpu[l].item()), 5)} for l in range(n_layers)]
            attn_layer_list = [{"layer": l, "contribution": round(float(head_contribs_cpu[l].sum().item()), 5)} for l in range(n_layers)]

            target_logit = float(last_logits[target_id].item())
            target_prob = float(last_logits.softmax(dim=-1)[target_id].item())
            target_str = model.to_string([target_id])

            top5 = [{"token": model.to_string([int(i.item())]),
                     "tokenId": int(i.item()),
                     "prob": round(float(p.item()), 5)} for p, i in zip(top5_vals, top5_idx)]

        del cache
        if device.type == "cuda":
            torch.cuda.empty_cache()

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "prompt": prompt,
            "targetToken": target_str,
            "targetTokenId": target_id,
            "targetLogit": round(target_logit, 5),
            "targetProb": round(target_prob, 5),
            "baselineTop5": top5,
            "headContributions": head_list,
            "mlpContributions": mlp_list,
            "attnLayerContributions": attn_layer_list,
            "totalAttn": round(float(head_contribs_cpu.sum().item()), 5),
            "totalMlp": round(float(mlp_contribs_cpu.sum().item()), 5),
        })
    except Exception as e:
        print(f"[DLA] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/induction-heads/analyze', methods=['POST'])
def induction_heads_analyze():
    """
    Detect induction heads — attention heads that implement the
    "[A][B]...[A] → predict [B]" pattern, the mechanism Anthropic identified
    as foundational to in-context learning.

    Algorithm (per Anthropic 2022, "In-context Learning and Induction Heads"):
    1. Build B random sequences of unique tokens of length seq_len, then
       concatenate each with itself to form a "repeated" sequence of length
       2*seq_len. Position p in the second half corresponds to the same token
       as position p - seq_len in the first half. The induction prediction
       at position p is "the token that came AFTER position (p - seq_len)
       in the first half" — i.e., the token at position p - seq_len + 1.
    2. Run forward with cache, collecting attention patterns from
       blocks.{l}.attn.hook_pattern for every layer.
    3. For each (layer, head): the induction score is the mean attention
       weight on the diagonal that maps query position q (in the second half)
       to key position (q - seq_len + 1) (the "prev-token-after-match"
       position in the first half), averaged over query positions in the
       second half and across batches.
    4. Top induction heads are the ones with the largest score.

    Request:  { batchSize, seqLen }
    Response: { success, model, device, nLayers, nHeads, batchSize, seqLen,
                topHeads:[{layer, head, score}],   // sorted desc
                heatmap:[L][H],                     // raw scores
                meanScore, maxScore }
    """
    try:
        data = request.json or {}
        batch = max(1, min(int(data.get('batchSize', 4)), 16))
        seq_len = max(8, min(int(data.get('seqLen', 25)), 64))

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers
        n_heads = model.cfg.n_heads
        vocab = model.cfg.d_vocab

        # Build random unique-token sequences and repeat them.
        # Avoid special tokens — sample from middle of vocab to skip BOS/EOS.
        rand = torch.randint(low=1000, high=vocab - 1000, size=(batch, seq_len), device=device)
        repeated = torch.cat([rand, rand], dim=1)  # [batch, 2*seq_len]

        pattern_names = [f'blocks.{l}.attn.hook_pattern' for l in range(n_layers)]
        with torch.no_grad():
            _, cache = model.run_with_cache(repeated, names_filter=pattern_names)

        # induction score per (layer, head):
        # mean over batch and over query positions q in [seq_len, 2*seq_len-1]
        # of attn_pattern[batch, head, q, q - seq_len + 1].
        scores = torch.zeros(n_layers, n_heads, device=device)
        for l in range(n_layers):
            patt = cache[pattern_names[l]]  # [batch, n_heads, q, k]
            # gather diag: for q in seq_len..2*seq_len-1, k = q - seq_len + 1
            qs = torch.arange(seq_len, 2 * seq_len, device=device)
            ks = qs - seq_len + 1
            # patt[:, :, qs, ks] — fancy indexing
            diag = patt[:, :, qs, ks]  # [batch, heads, seq_len]
            scores[l] = diag.mean(dim=(0, 2))
        scores_cpu = scores.cpu()

        flat = []
        for l in range(n_layers):
            for h in range(n_heads):
                flat.append({"layer": l, "head": h, "score": round(float(scores_cpu[l, h].item()), 5)})
        flat.sort(key=lambda x: x['score'], reverse=True)

        del cache
        if device.type == "cuda":
            torch.cuda.empty_cache()

        return jsonify({
            "success": True,
            "model": "gpt2",
            "device": str(device),
            "nLayers": n_layers,
            "nHeads": n_heads,
            "batchSize": batch,
            "seqLen": seq_len,
            "topHeads": flat[:12],
            "heatmap": [[round(float(scores_cpu[l, h].item()), 5) for h in range(n_heads)] for l in range(n_layers)],
            "meanScore": round(float(scores_cpu.mean().item()), 5),
            "maxScore": round(float(scores_cpu.max().item()), 5),
        })
    except Exception as e:
        print(f"[InductionHeads] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/activation-steering/analyze', methods=['POST'])
def activation_steering_analyze():
    """
    Activation Steering — classic mechanistic-interpretability technique.

    1. Build a "concept direction" by averaging the residual-stream activation
       at a chosen layer (last token position) over POSITIVE prompts that
       exemplify the target concept, then subtracting the mean over NEGATIVE
       prompts. The result is a single d_model-dim vector that points from
       "not concept" toward "concept" in activation space.
    2. Run the TEST prompt through the model twice:
         - Baseline: no intervention.
         - Steered: register a forward hook on blocks.{layer}.hook_resid_post
           that adds (coefficient * steering_vector) at every position.
    3. Compare next-token distributions: top-5 token shifts, KL divergence,
       and per-token probability changes for the most-affected tokens.

    Request:  { positivePrompts:[..], negativePrompts:[..], testPrompt,
                layer (0..n-1), coefficient (default 6.0) }
    Response: { success, device, model, layer, coefficient,
                steeringVectorNorm, baselineTop5, steeredTop5,
                klBaselineToSteered, topShifts:[{token, baselineProb, steeredProb, delta}] }
    """
    try:
        data = request.json or {}
        pos = [p for p in (data.get('positivePrompts') or []) if isinstance(p, str) and p.strip()][:30]
        neg = [p for p in (data.get('negativePrompts') or []) if isinstance(p, str) and p.strip()][:30]
        test_prompt = (data.get('testPrompt') or '').strip()
        coef = float(data.get('coefficient', 6.0))
        layer = int(data.get('layer', 6))
        if len(pos) < 2 or len(neg) < 2:
            return jsonify({"error": "need at least 2 positive AND 2 negative prompts"}), 400
        if not test_prompt:
            return jsonify({"error": "testPrompt is required"}), 400

        model = _load_hooked_transformer('gpt2')
        device = next(model.parameters()).device
        n_layers = model.cfg.n_layers
        if layer < 0 or layer >= n_layers:
            return jsonify({"error": f"layer must be in [0, {n_layers - 1}]"}), 400

        hook_name = f'blocks.{layer}.hook_resid_post'

        def mean_resid(prompts):
            acc = None
            n = 0
            for p in prompts:
                try:
                    toks = model.to_tokens(p).to(device)
                    if toks.shape[1] > 64:
                        toks = toks[:, :64]
                    with torch.no_grad():
                        _, cache = model.run_with_cache(toks, names_filter=[hook_name])
                    v = cache[hook_name][0, -1, :].float()
                    acc = v if acc is None else acc + v
                    n += 1
                    del cache
                except Exception as inner:
                    print(f"[ActivationSteering] skip: {inner}")
                    continue
            return (acc / n) if (acc is not None and n > 0) else None

        v_pos = mean_resid(pos)
        v_neg = mean_resid(neg)
        if v_pos is None or v_neg is None:
            return jsonify({"error": "all prompts failed to encode"}), 400
        steering_vec = (v_pos - v_neg).to(next(model.parameters()).dtype)
        sv_norm = float(torch.norm(steering_vec).item())

        # Baseline forward
        test_toks = model.to_tokens(test_prompt).to(device)
        if test_toks.shape[1] > 64:
            test_toks = test_toks[:, :64]
        with torch.no_grad():
            base_logits = model(test_toks)
        base_probs = torch.softmax(base_logits[0, -1, :].float(), dim=-1)

        # Steered forward — hook adds coef * steering_vec at every position.
        def steer_hook(activation, hook):
            # activation: [batch, pos, d_model] in model dtype
            return activation + (coef * steering_vec)

        with torch.no_grad():
            steered_logits = model.run_with_hooks(test_toks, fwd_hooks=[(hook_name, steer_hook)])
        steered_probs = torch.softmax(steered_logits[0, -1, :].float(), dim=-1)

        # Top-5 each side
        def top5(probs):
            vals, idxs = probs.topk(5)
            return [
                {"token": model.to_string(torch.tensor([int(i)])), "prob": round(float(v.item()), 4)}
                for v, i in zip(vals, idxs)
            ]
        baseline_top5 = top5(base_probs)
        steered_top5 = top5(steered_probs)

        # KL(steered || baseline)
        eps = 1e-12
        kl = (steered_probs * (steered_probs.clamp_min(eps).log() - base_probs.clamp_min(eps).log())).sum().item()

        # Top shifts by absolute delta
        delta = (steered_probs - base_probs).cpu()
        abs_top = torch.topk(delta.abs(), 12).indices.tolist()
        top_shifts = []
        for tid in abs_top:
            top_shifts.append({
                "token": model.to_string(torch.tensor([tid])),
                "baselineProb": round(float(base_probs[tid].item()), 5),
                "steeredProb": round(float(steered_probs[tid].item()), 5),
                "delta": round(float(delta[tid].item()), 5),
            })

        return jsonify({
            "success": True,
            "device": str(device),
            "model": "gpt2",
            "layer": layer,
            "nLayers": n_layers,
            "coefficient": coef,
            "steeringVectorNorm": round(sv_norm, 4),
            "nPositive": len(pos),
            "nNegative": len(neg),
            "baselineTop5": baseline_top5,
            "steeredTop5": steered_top5,
            "klBaselineToSteered": round(float(kl), 6),
            "topShifts": top_shifts,
        })
    except Exception as e:
        print(f"[ActivationSteering] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/quantization-impact/analyze', methods=['POST'])
def quantization_impact_analyze():
    """
    Compare a model in fp32 (full precision) vs fp16 (half precision) to measure
    the quality / latency tradeoff that production deployments care about.

    Pipeline:
    1. Load gpt2 fresh in fp32, deep-copy and cast to fp16 — both share weights but
       differ only in numerical precision.
    2. Warm up both models with one forward pass each (CUDA tensor-core JIT).
    3. For each prompt:
         - Time forward pass on fp32 and fp16 separately (CUDA-synchronized).
         - Take final-position next-token logits, softmax in fp32 for fair comparison.
         - Compute top-1 agreement, top-5 Jaccard, KL(fp16 ‖ fp32).
    4. Aggregate: mean KL, top-1 agreement rate, mean Jaccard, mean latency speedup,
       theoretical memory savings (fp32 → fp16 = 2× reduction).

    Request:  { prompts:[..] }
    Response: { success, device, model, paramCount, nPrompts,
                fp32MeanLatencyMs, fp16MeanLatencyMs, speedup,
                memBytesFp32, memBytesFp16, memReductionPct,
                top1AgreementRate, meanTop5Jaccard, meanKL,
                rows:[{prompt, top1Fp32, top1Fp16, agree, top5Jaccard, kl,
                       latencyFp32Ms, latencyFp16Ms}] }
    """
    try:
        import copy
        import time
        from transformer_lens import HookedTransformer
        data = request.json or {}
        prompts = [p for p in (data.get('prompts') or []) if isinstance(p, str) and p.strip()][:25]
        if len(prompts) < 1:
            return jsonify({"error": "at least 1 prompt required"}), 400

        device_str = "cuda" if torch.cuda.is_available() else "cpu"
        # Fresh load both copies — bypass the cached singleton so we don't corrupt
        # the global model precision. Use .to(device, dtype=...) which converts
        # parameters AND buffers (.half() alone misses some HookedTransformer buffers
        # like attn.mask / attn.IGNORE that get registered as fp32).
        model_fp32 = HookedTransformer.from_pretrained('gpt2', dtype=torch.float32).to(device_str)
        model_fp32.eval()
        model_fp16 = HookedTransformer.from_pretrained('gpt2', dtype=torch.float16).to(device_str)
        model_fp16.eval()
        device = next(model_fp32.parameters()).device

        param_count = sum(p.numel() for p in model_fp32.parameters())
        mem_fp32 = param_count * 4
        mem_fp16 = param_count * 2

        # Warmup
        with torch.no_grad():
            warm = model_fp32.to_tokens("warmup")
            _ = model_fp32(warm.to(device))
            _ = model_fp16(warm.to(device))
        if device.type == "cuda":
            torch.cuda.synchronize()

        rows = []
        kls = []
        agree_n = 0
        jaccards = []
        lat32s = []
        lat16s = []

        for p in prompts:
            try:
                tokens = model_fp32.to_tokens(p).to(device)
                if tokens.shape[1] > 64:
                    tokens = tokens[:, :64]

                with torch.no_grad():
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t0 = time.perf_counter()
                    l32 = model_fp32(tokens)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t1 = time.perf_counter()
                    l16 = model_fp16(tokens)
                    if device.type == "cuda":
                        torch.cuda.synchronize()
                    t2 = time.perf_counter()
                lat32 = (t1 - t0) * 1000.0
                lat16 = (t2 - t1) * 1000.0
                lat32s.append(lat32)
                lat16s.append(lat16)

                p32 = torch.softmax(l32[0, -1, :].float(), dim=-1)
                p16 = torch.softmax(l16[0, -1, :].float(), dim=-1)

                top1_32 = int(p32.argmax().item())
                top1_16 = int(p16.argmax().item())
                agree = (top1_32 == top1_16)
                if agree:
                    agree_n += 1

                top5_32 = set(p32.topk(5).indices.tolist())
                top5_16 = set(p16.topk(5).indices.tolist())
                inter = len(top5_32 & top5_16)
                union = len(top5_32 | top5_16)
                jacc = inter / union if union else 0.0
                jaccards.append(jacc)

                eps = 1e-12
                kl = (p16 * (p16.clamp_min(eps).log() - p32.clamp_min(eps).log())).sum().item()
                kls.append(kl)

                rows.append({
                    "prompt": p[:200],
                    "top1Fp32": model_fp32.to_string(torch.tensor([top1_32])),
                    "top1Fp16": model_fp16.to_string(torch.tensor([top1_16])),
                    "agree": agree,
                    "top5Jaccard": round(jacc, 4),
                    "kl": round(float(kl), 6),
                    "latencyFp32Ms": round(lat32, 2),
                    "latencyFp16Ms": round(lat16, 2),
                })
                del l32, l16, p32, p16
            except Exception as inner:
                print(f"[QuantizationImpact] skip prompt: {inner}")
                continue

        if not rows:
            return jsonify({"error": "no usable prompts"}), 400

        mean_lat32 = sum(lat32s) / len(lat32s)
        mean_lat16 = sum(lat16s) / len(lat16s)
        speedup = mean_lat32 / mean_lat16 if mean_lat16 > 0 else 1.0

        # Free memory
        del model_fp32, model_fp16
        if device.type == "cuda":
            torch.cuda.empty_cache()

        return jsonify({
            "success": True,
            "device": str(device),
            "model": "gpt2",
            "paramCount": param_count,
            "nPrompts": len(rows),
            "fp32MeanLatencyMs": round(mean_lat32, 2),
            "fp16MeanLatencyMs": round(mean_lat16, 2),
            "speedup": round(speedup, 3),
            "memBytesFp32": mem_fp32,
            "memBytesFp16": mem_fp16,
            "memReductionPct": round(50.0, 1),
            "top1AgreementRate": round(agree_n / len(rows), 4),
            "meanTop5Jaccard": round(sum(jaccards) / len(jaccards), 4),
            "meanKL": round(sum(kls) / len(kls), 6),
            "rows": rows,
        })
    except Exception as e:
        print(f"[QuantizationImpact] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


@app.route('/knowledge-distillation/analyze', methods=['POST'])
def knowledge_distillation_analyze():
    """
    Compare a teacher model and a student (distilled) model on the same prompts.

    Pipeline:
    1. Load teacher (gpt2 full, 12 layers) and student (distilgpt2, 6 layers) — both
       share the GPT-2 BPE tokenizer / 50257 vocab so logits align.
    2. For each prompt, run a forward pass on each model and take next-token logits at
       the final position. Convert to probability distributions in fp32.
    3. Per-prompt metrics:
         - top-1 agreement: argmax_t == argmax_s
         - top-5 jaccard: |T5 ∩ S5| / |T5 ∪ S5|
         - KL(student || teacher) = Σ s * (log s − log t) — the canonical distillation loss
    4. Aggregate: mean KL, top-1 agreement rate, mean top-5 jaccard.

    Request:  { prompts:[..] }   (always uses gpt2 vs distilgpt2)
    Response: { success, device, teacher, student, nPrompts,
                meanKL, top1AgreementRate, meanTop5Jaccard,
                rows:[{prompt, teacherTop1, studentTop1, agree, top5Jaccard, kl,
                       teacherTop3:[{tok,prob}], studentTop3:[{tok,prob}]}] }
    """
    try:
        data = request.json or {}
        prompts = [p for p in (data.get('prompts') or []) if isinstance(p, str) and p.strip()][:25]
        if len(prompts) < 1:
            return jsonify({"error": "at least 1 prompt required"}), 400

        teacher = _load_hooked_transformer('gpt2')
        student = _load_hooked_transformer('distilgpt2')
        device = next(teacher.parameters()).device

        rows = []
        kls = []
        agree_n = 0
        jaccards = []

        def top3(probs, model):
            vals, ids = probs.topk(3)
            out = []
            for v, i in zip(vals.tolist(), ids.tolist()):
                try:
                    tok = model.to_string(torch.tensor([i]))
                except Exception:
                    tok = f"[id={i}]"
                out.append({"tok": tok, "prob": round(float(v), 4)})
            return out

        for p in prompts:
            try:
                t_tokens = teacher.to_tokens(p).to(device)
                s_tokens = student.to_tokens(p).to(device)
                if t_tokens.shape[1] > 64:
                    t_tokens = t_tokens[:, :64]
                if s_tokens.shape[1] > 64:
                    s_tokens = s_tokens[:, :64]
                with torch.no_grad():
                    t_logits = teacher(t_tokens)
                    s_logits = student(s_tokens)
                t_probs = torch.softmax(t_logits[0, -1, :].float(), dim=-1)
                s_probs = torch.softmax(s_logits[0, -1, :].float(), dim=-1)

                t_top1 = int(t_probs.argmax().item())
                s_top1 = int(s_probs.argmax().item())
                agree = (t_top1 == s_top1)
                if agree:
                    agree_n += 1

                t_top5 = set(t_probs.topk(5).indices.tolist())
                s_top5 = set(s_probs.topk(5).indices.tolist())
                inter = len(t_top5 & s_top5)
                union = len(t_top5 | s_top5)
                jacc = inter / union if union else 0.0
                jaccards.append(jacc)

                eps = 1e-12
                kl = (s_probs * (s_probs.clamp_min(eps).log() - t_probs.clamp_min(eps).log())).sum().item()
                kls.append(kl)

                rows.append({
                    "prompt": p[:200],
                    "teacherTop1": teacher.to_string(torch.tensor([t_top1])),
                    "studentTop1": student.to_string(torch.tensor([s_top1])),
                    "agree": agree,
                    "top5Jaccard": round(jacc, 4),
                    "kl": round(float(kl), 6),
                    "teacherTop3": top3(t_probs, teacher),
                    "studentTop3": top3(s_probs, student),
                })
                del t_logits, s_logits, t_probs, s_probs
            except Exception as inner:
                print(f"[KnowledgeDistillation] skip prompt: {inner}")
                continue

        if not rows:
            return jsonify({"error": "no usable prompts"}), 400

        return jsonify({
            "success": True,
            "device": str(device),
            "teacher": "gpt2 (12 layers, 124M)",
            "student": "distilgpt2 (6 layers, 82M)",
            "nPrompts": len(rows),
            "meanKL": round(sum(kls) / len(kls), 6),
            "top1AgreementRate": round(agree_n / len(rows), 4),
            "meanTop5Jaccard": round(sum(jaccards) / len(jaccards), 4),
            "rows": rows,
        })
    except Exception as e:
        print(f"[KnowledgeDistillation] Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500


if __name__ == '__main__':
    port = int(os.environ.get('SAE_SERVICE_PORT', 7860))
    print(f"[SAE Service] Starting on port {port}")
    print(f"[SAE Service] Models will be loaded on first request or via /warmup")

    app.run(host='0.0.0.0', port=port, debug=False, threaded=False)