Blame for backend/model_service.py

"""
Unified Model Service for Visualisable.ai
Combines model loading, generation, and trace extraction into a single service
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, BackgroundTasks, HTTPException, Depends
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import asyncio
import gc
import json
import os
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Optional, List, Dict, Any
import numpy as np
import logging
from datetime import datetime
import traceback
import uuid
from threading import Lock
from time import time as time_now
from .auth import verify_api_key
from .instrumentation import ModelInstrumentor, InstrumentationData, TokenMetadata
from .storage import ZarrStorage, generate_run_id
from .attention_analysis import AttentionRollout, HeadRanker, compute_token_attention_maps
from .tokenizer_utils import TokenizerMetadata, get_tokenizer_stats
from .architectural_analysis import extract_architectural_data

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

import math

def sanitize_for_json(obj):
    """
    Recursively sanitize a nested data structure for JSON serialization.
    Replaces NaN and Infinity float values with None (JSON null).
    """
    if isinstance(obj, dict):
        return {k: sanitize_for_json(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [sanitize_for_json(item) for item in obj]
    elif isinstance(obj, float):
        if math.isnan(obj) or math.isinf(obj):
            return None  # JSON null
        return obj
    elif isinstance(obj, (np.floating, np.float32, np.float64)):
        val = float(obj)
        if math.isnan(val) or math.isinf(val):
            return None
        return val
    elif isinstance(obj, (np.integer, np.int32, np.int64)):
        return int(obj)
    elif isinstance(obj, np.ndarray):
        return sanitize_for_json(obj.tolist())
    else:
        return obj


# Matrix cache for lazy loading (60 min TTL)
class MatrixCache:
    """
    Thread-safe in-memory cache for attention matrices.
    Stores Q/K/V and attention weights per (request_id, step, layer, head).
    """
    def __init__(self, ttl_seconds: int = 3600):
        self._cache: Dict[str, Dict] = {}
        self._timestamps: Dict[str, float] = {}
        self._request_ids: set = set()  # Track active request IDs
        # Per-request metadata captured at analysis time. Aggregation across
        # heads (e.g. mean/max in /attention/row) needs num_heads, and that
        # number belongs to the model that produced the cached attention,
        # not whichever model is loaded when the row is later retrieved.
        # Storing it here decouples retrieval from `manager.model` state.
        self._request_meta: Dict[str, Dict] = {}
        self._lock = Lock()
        self._ttl = ttl_seconds

    def clear_request(self, request_id: str):
        """Clear all cache entries for a specific request."""
        with self._lock:
            keys_to_delete = [k for k in self._cache.keys() if k.startswith(f"{request_id}:")]
            for k in keys_to_delete:
                del self._cache[k]
                if k in self._timestamps:
                    del self._timestamps[k]
            self._request_ids.discard(request_id)
            self._request_meta.pop(request_id, None)
            if keys_to_delete:
                logger.info(f"MatrixCache: cleared {len(keys_to_delete)} entries for request {request_id[:8]}")

    def clear_old_requests(self, keep_request_id: str = None):
        """Clear all requests except the specified one to free memory."""
        with self._lock:
            request_ids_to_clear = self._request_ids - {keep_request_id} if keep_request_id else self._request_ids.copy()
            total_cleared = 0
            for rid in request_ids_to_clear:
                keys_to_delete = [k for k in self._cache.keys() if k.startswith(f"{rid}:")]
                for k in keys_to_delete:
                    del self._cache[k]
                    if k in self._timestamps:
                        del self._timestamps[k]
                total_cleared += len(keys_to_delete)
            self._request_ids = {keep_request_id} if keep_request_id else set()
            # Drop metadata for cleared requests; keep only the surviving one.
            for rid in list(self._request_meta.keys()):
                if rid != keep_request_id:
                    self._request_meta.pop(rid, None)
            if total_cleared:
                logger.info(f"MatrixCache: cleared {total_cleared} entries from old requests")
                # Force garbage collection to release memory back to system
                gc.collect()
                # Also clear any GPU cache if using CUDA
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                    # For Apple Silicon MPS, trigger garbage collection
                    torch.mps.empty_cache() if hasattr(torch.mps, 'empty_cache') else None

    def store(self, request_id: str, step: int, layer: int, head: int, data: dict):
        """Store matrix data for a specific head."""
        key = f"{request_id}:{step}:{layer}:{head}"
        with self._lock:
            self._cache[key] = data
            self._timestamps[key] = time_now()
            self._request_ids.add(request_id)

    def set_request_metadata(self, request_id: str, num_heads: int,
                             num_layers: Optional[int] = None,
                             model_id: Optional[str] = None):
        """
        Record metadata for the model that produced this request's cached
        matrices. Aggregation operations (e.g. mean/max across heads) read
        num_heads from here rather than from the live `manager.model`, so
        that switching models mid-session cannot corrupt aggregate results
        for older requests.
        """
        with self._lock:
            existing = self._request_meta.get(request_id, {})
            existing.update({
                "num_heads": num_heads,
                "num_layers": num_layers,
                "model_id": model_id,
            })
            self._request_meta[request_id] = existing

    def set_step_query_info(self, request_id: str, step: int,
                            query_position: int, query_token_id: int):
        """
        Record the predicting query position for a given generation step.

        At step `s`, the model's forward pass runs over the prompt + first
        `s` generated tokens; the logits at the LAST query position select
        the step-`s` token. The cached attention[step][layer][head][-1] is
        therefore the attention pattern of *that* query position, not of
        the predicted token (which doesn't exist in the sequence yet).
        Storing the query position and token id makes the provenance of
        attention rows explicit when the Why this token? panel surfaces
        them — supporting the "grounded hypotheses, not causal claims"
        framing required by the project.
        """
        with self._lock:
            meta = self._request_meta.setdefault(request_id, {})
            steps = meta.setdefault("steps", {})
            steps[step] = {
                "query_position": query_position,
                "query_token_id": query_token_id,
            }

    def get_request_metadata(self, request_id: str) -> Optional[dict]:
        """Return per-request metadata recorded at analysis time, or None."""
        with self._lock:
            meta = self._request_meta.get(request_id)
            return dict(meta) if meta else None

    def get_step_query_info(self, request_id: str, step: int) -> Optional[dict]:
        """Return {query_position, query_token_id} for a step, or None."""
        with self._lock:
            meta = self._request_meta.get(request_id)
            if not meta:
                return None
            steps = meta.get("steps")
            if not steps:
                return None
            entry = steps.get(step)
            return dict(entry) if entry else None

    def get(self, request_id: str, step: int, layer: int, head: int) -> Optional[dict]:
        """Retrieve matrix data, returning None if expired or not found."""
        key = f"{request_id}:{step}:{layer}:{head}"
        with self._lock:
            if key in self._cache:
                if time_now() - self._timestamps[key] < self._ttl:
                    return self._cache[key]
                else:
                    # Expired - clean up
                    del self._cache[key]
                    del self._timestamps[key]
        return None

    def cleanup_expired(self):
        """Remove all expired entries from cache."""
        with self._lock:
            now = time_now()
            expired = [k for k, t in self._timestamps.items() if now - t >= self._ttl]
            for k in expired:
                del self._cache[k]
                del self._timestamps[k]
            if expired:
                logger.info(f"MatrixCache: cleaned up {len(expired)} expired entries")

    def get_stats(self) -> dict:
        """Return cache statistics."""
        with self._lock:
            return {
                "entries": len(self._cache),
                "ttl_seconds": self._ttl
            }

    def get_attention_row(self, request_id: str, step: int, layer: int, head: int) -> Optional[list]:
        """
        Extract single attention row (last token's attention to all preceding positions).
        Used for attention overlay visualization.
        """
        data = self.get(request_id, step, layer, head)
        if not data or 'attention_weights' not in data:
            return None
        attention = data['attention_weights']
        if attention is None or len(attention) == 0:
            return None
        # Return last row (query token attending to all keys)
        # Handle both numpy arrays and lists
        last_row = attention[-1]
        if hasattr(last_row, 'tolist'):
            return last_row.tolist()
        return list(last_row)

    def get_aggregate_row(self, request_id: str, step: int, layer: int,
                          mode: str = "mean",
                          num_heads: Optional[int] = None) -> Optional[list]:
        """
        Compute aggregated attention row across all heads for a layer.

        num_heads is derived from per-request metadata recorded at analysis
        time so the aggregation always uses the head count of the model that
        produced the cached attention. The explicit `num_heads` argument is
        retained for callers that want to override (e.g. tests); when None,
        falls back to the request's stored metadata.

        Args:
            request_id: UUID from analysis
            step: Generation step
            layer: Layer index
            mode: Aggregation mode - "mean" or "max"
            num_heads: Optional override; otherwise derived from request meta

        Returns:
            List of aggregated attention weights, or None if data unavailable
            or request metadata is missing.
        """
        if num_heads is None:
            meta = self.get_request_metadata(request_id)
            if meta is None or meta.get("num_heads") is None:
                return None
            num_heads = meta["num_heads"]
        rows = []
        for h in range(num_heads):
            row = self.get_attention_row(request_id, step, layer, h)
            if row:
                rows.append(row)
        if not rows:
            return None
        arr = np.array(rows)
        if mode == "mean":
            return np.mean(arr, axis=0).tolist()
        elif mode == "max":
            return np.max(arr, axis=0).tolist()
        else:
            # Default to mean for unknown modes
            return np.mean(arr, axis=0).tolist()

    def get_aggregate_attention_matrix(
        self,
        request_id: str,
        step: int,
        layer: int,
        mode: str = "mean",
        num_heads: Optional[int] = None,
    ) -> Optional[list]:
        """
        Compute the aggregated attention MATRIX across all heads for a
        layer. Mirrors get_aggregate_row but returns the full
        [seq_len, seq_len] matrix instead of one row.

        Used by the Mechanism lens's Step 4 (Attention) to render the
        per-step matrix in its "all heads (mean)" default mode without
        forcing the frontend to fetch + aggregate N matrices itself.

        Args:
            request_id: UUID from analysis
            step: Generation step
            layer: Layer index
            mode: Aggregation mode - "mean" or "max"
            num_heads: Optional override; otherwise derived from request meta

        Returns:
            Aggregated 2-D matrix as a list of lists, or None if data
            unavailable or request metadata is missing.
        """
        if num_heads is None:
            meta = self.get_request_metadata(request_id)
            if meta is None or meta.get("num_heads") is None:
                return None
            num_heads = meta["num_heads"]
        matrices = []
        for h in range(num_heads):
            data = self.get(request_id, step, layer, h)
            if data is None:
                continue
            attn = data.get("attention_weights")
            if attn is None:
                continue
            matrices.append(np.asarray(attn))
        if not matrices:
            return None
        arr = np.stack(matrices, axis=0)  # (n_heads, seq, seq)
        if mode == "max":
            return np.max(arr, axis=0).tolist()
        # Default to mean for "mean" or unknown modes.
        return np.mean(arr, axis=0).tolist()


# Global matrix cache instance
matrix_cache = MatrixCache(ttl_seconds=3600)  # 60 min TTL


def _classify_stability(margin: float) -> str:
    """Classify a logit margin into a stability category."""
    if margin > 1.0:
        return "stable"
    elif margin >= 0.3:
        return "moderate"
    elif margin >= 0.1:
        return "boundary"
    else:
        return "fragile"


def _compute_embedding_norms(manager, token_ids: List[int]) -> List[float]:
    """
    Per-token L2 norms of the layer-0 embedding vectors, indexed by the
    token ids supplied. Used by the analyse endpoints to surface a
    real-data anchor for the residual-norm trace shown in the
    Mechanism lens's Microscope rail (RQ1).

    Memory is bounded by the number of *unique* token ids in the
    sequence rather than the sequence length: a long prompt with
    repeated tokens reuses the same embedding row, so we gather rows
    once for unique ids and scatter the resulting scalar back to the
    sequence positions. This avoids materialising a full
    [1, seq_len, hidden_size] activation tensor — important for the
    long-context models this service supports (e.g. devstral-small at
    131,072 tokens, where the naive path would allocate multi-GB at
    end-of-request and risk OOM).

    Failures (missing module, unexpected layout) downgrade gracefully
    to an empty list so the rest of the trace is unaffected.
    """
    if not token_ids:
        return []
    try:
        with torch.no_grad():
            embed_layer = manager.model.get_input_embeddings()
            ids_tensor = torch.tensor(
                token_ids, dtype=torch.long, device=manager.device
            )
            unique_ids, inverse = torch.unique(
                ids_tensor, return_inverse=True
            )
            # Index the embedding weight matrix directly. Allocation is
            # (n_unique, hidden_size) instead of (seq_len, hidden_size)
            # — typically a few hundred rows for code-generation
            # prompts regardless of how long the prompt is.
            unique_rows = embed_layer.weight.index_select(0, unique_ids)
            unique_norms = torch.linalg.vector_norm(
                unique_rows.float(), dim=-1
            )
            seq_norms = unique_norms[inverse]
            return [float(v) for v in seq_norms.cpu().tolist()]
    except Exception as e:  # pragma: no cover — defensive
        logger.warning(f"Failed to compute embedding norms: {e}")
        return []


class HiddenStateCache:
    """
    Cache for hidden states and logits per (request_id, step).
    Used by intervention endpoints to re-run forward passes on cached data.
    Capped at MAX_CACHED_RUNS to manage memory.
    """
    MAX_CACHED_RUNS = 5

    def __init__(self, ttl_seconds: int = 3600):
        self._hidden_states: Dict[str, Dict] = {}  # key: request_id -> {step -> tensor}
        self._logits: Dict[str, Dict] = {}          # key: request_id -> {step -> tensor}
        self._input_ids: Dict[str, object] = {}     # key: request_id -> tensor
        self._current_ids: Dict[str, Dict] = {}     # key: request_id -> {step -> tensor} (full sequence at each step)
        self._timestamps: Dict[str, float] = {}
        self._lock = Lock()
        self._ttl = ttl_seconds

    def store_step(self, request_id: str, step: int, hidden_states, raw_logits_tensor, current_ids_tensor=None):
        """Store hidden states, logits, and optionally the full input sequence for a generation step."""
        with self._lock:
            if request_id not in self._hidden_states:
                # Evict oldest if at capacity
                if len(self._hidden_states) >= self.MAX_CACHED_RUNS:
                    oldest_rid = min(self._timestamps, key=self._timestamps.get)
                    self._evict(oldest_rid)
                self._hidden_states[request_id] = {}
                self._logits[request_id] = {}
                self._current_ids[request_id] = {}

            # Store detached CPU copies to avoid holding GPU memory
            self._hidden_states[request_id][step] = [h.detach().cpu() for h in hidden_states]
            self._logits[request_id][step] = raw_logits_tensor.detach().cpu()
            if current_ids_tensor is not None:
                self._current_ids[request_id][step] = current_ids_tensor.detach().cpu()
            self._timestamps[request_id] = time_now()

    def store_input_ids(self, request_id: str, input_ids_tensor):
        """Store the full input_ids tensor for a run."""
        with self._lock:
            self._input_ids[request_id] = input_ids_tensor.detach().cpu()
            self._timestamps[request_id] = time_now()

    def get_step(self, request_id: str, step: int):
        """Retrieve hidden states and logits for a step. Returns (hidden_states, logits) or (None, None)."""
        with self._lock:
            if request_id in self._timestamps and time_now() - self._timestamps[request_id] >= self._ttl:
                self._evict(request_id)
                return None, None
            hs = self._hidden_states.get(request_id, {}).get(step)
            lg = self._logits.get(request_id, {}).get(step)
            return hs, lg

    def get_logits(self, request_id: str, step: int):
        """Retrieve just the logits for a step."""
        with self._lock:
            if request_id in self._timestamps and time_now() - self._timestamps[request_id] >= self._ttl:
                self._evict(request_id)
                return None
            return self._logits.get(request_id, {}).get(step)

    def get_input_ids(self, request_id: str):
        """Retrieve stored input_ids for a run."""
        with self._lock:
            if request_id in self._timestamps and time_now() - self._timestamps[request_id] >= self._ttl:
                self._evict(request_id)
                return None
            return self._input_ids.get(request_id)

    def get_current_ids(self, request_id: str, step: int):
        """Retrieve the full input sequence (prompt + generated) at a specific step."""
        with self._lock:
            if request_id in self._timestamps and time_now() - self._timestamps[request_id] >= self._ttl:
                self._evict(request_id)
                return None
            return self._current_ids.get(request_id, {}).get(step)

    def get_all_steps(self, request_id: str):
        """Return list of cached step indices for a run."""
        with self._lock:
            return list(self._hidden_states.get(request_id, {}).keys())

    def has_run(self, request_id: str) -> bool:
        """Check if a run is cached."""
        with self._lock:
            if request_id in self._timestamps and time_now() - self._timestamps[request_id] >= self._ttl:
                self._evict(request_id)
                return False
            return request_id in self._hidden_states

    def _evict(self, request_id: str):
        """Remove all data for a request (must hold lock)."""
        self._hidden_states.pop(request_id, None)
        self._logits.pop(request_id, None)
        self._input_ids.pop(request_id, None)
        self._current_ids.pop(request_id, None)
        self._timestamps.pop(request_id, None)

    def get_stats(self) -> dict:
        with self._lock:
            return {
                "cached_runs": len(self._hidden_states),
                "max_runs": self.MAX_CACHED_RUNS,
                "ttl_seconds": self._ttl,
            }


# Global hidden state cache instance
hidden_state_cache = HiddenStateCache(ttl_seconds=3600)


app = FastAPI(title="Visualisable.ai Model Service", version="0.1.0")

# CORS configuration for local development and production
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000", 
        "http://localhost:3001", 
        "http://localhost:3002",
        "https://visualisable-ai.vercel.app",
        "https://*.vercel.app"
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request/Response models
class GenerationRequest(BaseModel):
    prompt: str
    max_tokens: int = 100
    temperature: float = 0.7
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    extract_traces: bool = True
    sampling_rate: float = 0.005
    layer_stride: int = 1  # 1 = all layers, 2 = every other layer, etc.

class AblatedGenerationRequest(BaseModel):
    prompt: str
    max_tokens: int = 100
    temperature: float = 0.7
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    extract_traces: bool = False
    disabled_components: Optional[Dict[str, Any]] = None

class ICLExample(BaseModel):
    input: str
    output: str

class ICLGenerationRequest(BaseModel):
    examples: List[ICLExample]
    prompt: str
    max_tokens: int = 200  # Increased to accommodate examples + generation
    temperature: float = 0.7
    analyze: bool = True

class AblatedHead(BaseModel):
    layer: int
    head: int

class StudyRequest(BaseModel):
    prompt: str
    max_tokens: int = 50
    seed: int = 42
    temperature: float = 0.0  # Deterministic by default for reproducibility
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    disabled_components: Optional[Dict[str, Any]] = None

class DemoRequest(BaseModel):
    demo_id: str

class TraceData(BaseModel):
    type: str
    layer: Optional[str] = None
    weights: Optional[List[List[float]]] = None
    tokens: Optional[List[str]] = None  # Add tokens field
    max_weight: Optional[float] = None
    entropy: Optional[float] = None
    mean: Optional[float] = None
    std: Optional[float] = None
    confidence_score: Optional[float] = None
    hallucination_risk: Optional[float] = None
    timestamp: float

class ModelManager:
    """Manages model loading and generation with trace extraction"""

    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.adapter = None  # ModelAdapter for multi-model support
        self.device = None
        self.dtype = None  # Will be set from TORCH_DTYPE env var
        self.websocket_clients: List[WebSocket] = []
        self.trace_buffer: List[TraceData] = []

        # Read configuration from environment variables
        self.model_id = os.environ.get("DEFAULT_MODEL", "devstral-small")
        self.max_context = int(os.environ.get("MAX_CONTEXT", "8192"))
        self.batch_size = int(os.environ.get("BATCH_SIZE", "1"))

        # Get model config and HF path
        from .model_config import get_model_config
        config = get_model_config(self.model_id)
        if config:
            self.model_name = config["hf_path"]
        else:
            # Fallback to default if model_id not found
            logger.warning(f"Unknown model ID '{self.model_id}', falling back to codegen-350m")
            self.model_id = "codegen-350m"
            self.model_name = "Salesforce/codegen-350M-mono"
        
    async def initialize(self):
        """Load model on startup"""
        try:
            # Check for device override from environment
            device_override = os.environ.get("DEVICE", "").lower()

            if device_override == "cpu":
                self.device = torch.device("cpu")
                device_name = "CPU (forced via DEVICE env var)"
            elif device_override == "cuda":
                self.device = torch.device("cuda")
                device_name = "CUDA GPU (forced via DEVICE env var)"
            elif torch.cuda.is_available():
                self.device = torch.device("cuda")
                device_name = "CUDA GPU"
            elif torch.backends.mps.is_available():
                self.device = torch.device("mps")
                device_name = "Apple Silicon GPU"
            else:
                self.device = torch.device("cpu")
                device_name = "CPU"

            # Determine dtype from environment, model config, or defaults
            dtype_str = os.environ.get("TORCH_DTYPE", "").lower()

            # If not set in env, use model's recommended dtype
            if not dtype_str:
                from .model_config import get_model_config
                model_config = get_model_config(self.model_id)
                if model_config and "recommended_dtype" in model_config:
                    dtype_str = model_config["recommended_dtype"]
                    logger.info(f"Using model's recommended dtype: {dtype_str}")

            # Parse dtype string to torch dtype
            if dtype_str == "bf16" or dtype_str == "bfloat16":
                self.dtype = torch.bfloat16
                dtype_name = "bfloat16"
            elif dtype_str == "fp16" or dtype_str == "float16":
                self.dtype = torch.float16
                dtype_name = "float16"
            elif dtype_str == "fp32" or dtype_str == "float32":
                self.dtype = torch.float32
                dtype_name = "float32"
            elif self.device.type == "cpu":
                # Default to float32 for CPU (safest)
                self.dtype = torch.float32
                dtype_name = "float32 (CPU default)"
            else:
                # Default to float16 for GPU
                self.dtype = torch.float16
                dtype_name = "float16 (GPU default)"

            logger.info(f"Loading model '{self.model_id}' on {device_name} with dtype {dtype_name}...")
            logger.info(f"  HuggingFace path: {self.model_name}")
            logger.info(f"  Max context: {self.max_context}, Batch size: {self.batch_size}")

            # Load model with configured dtype
            # Use eager attention to support output_attentions=True for visualization
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=self.dtype,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
                attn_implementation="eager"
            ).to(self.device)
            
            # Load tokenizer. Devstral's official repo ships only tekken.json,
            # which AutoTokenizer can't consume on tokenizers ≥ 0.20. Fall back
            # to the Unsloth mirror, which carries the same Tekken vocabulary
            # in standard HF format. Same vocab size, same token IDs, same
            # special tokens — only the on-disk file format differs.
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            except (OSError, EnvironmentError) as tok_err:
                if self.model_id == "devstral-small":
                    logger.warning(
                        f"AutoTokenizer failed for {self.model_name} ({tok_err}); "
                        f"falling back to unsloth/Devstral-Small-2507 (same Tekken vocab, HF format)."
                    )
                    self.tokenizer = AutoTokenizer.from_pretrained(
                        "unsloth/Devstral-Small-2507"
                    )
                else:
                    raise
            # Set pad_token if the tokenizer allows it (some like MistralCommonTokenizer don't)
            try:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            except AttributeError:
                logger.info(f"Tokenizer doesn't support setting pad_token (using default)")

            # For Devstral, also load MistralTokenizer for correct encoding
            self.mistral_tokenizer = None
            if self.model_id == "devstral-small":
                from .mistral_tokenizer import create_mistral_tokenizer
                self.mistral_tokenizer = create_mistral_tokenizer(self.model_name)
                if self.mistral_tokenizer:
                    logger.info("Loaded MistralTokenizer for Devstral (correct Tekken encoding)")
                else:
                    logger.warning("MistralTokenizer not available - Devstral may produce garbage output")

            # Create model adapter for multi-model support
            from .model_adapter import create_adapter
            try:
                self.adapter = create_adapter(self.model, self.tokenizer, self.model_id)
                logger.info(f"✅ Created adapter for model: {self.model_id}")
            except Exception as adapter_error:
                logger.warning(f"Failed to create adapter: {adapter_error}")
                # Continue without adapter - some features may not work

            logger.info("✅ Model loaded successfully")

            # Load tuned lens probes (optional — falls back to raw logit lens if unavailable)
            from .tuned_lens import tuned_lens_runtime
            tuned_lens_runtime.load(self.model_id, self.device, self.dtype)

        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    
    def extract_attention_trace(self, layer_idx: int, attention_weights, tokens: Optional[List[str]] = None) -> TraceData:
        """Extract attention pattern trace from a layer"""
        # attention_weights is a tuple of tensors, one for each layer
        # Each tensor has shape (batch_size, num_heads, seq_len, seq_len)
        layer_attention = attention_weights[layer_idx]
        
        # Average across all heads for visualization
        # Shape: (batch_size, num_heads, seq_len, seq_len) -> (seq_len, seq_len)
        avg_attention = layer_attention[0].mean(dim=0).detach().cpu().float().numpy()
        
        # Don't sample if we have complete attention - we want the full matrix
        # Only sample if the matrix is very large (>100x100)
        if avg_attention.shape[0] > 100:
            indices = np.random.choice(avg_attention.shape[0], 100, replace=False)
            avg_attention = avg_attention[indices][:, indices]
            if tokens:
                tokens = [tokens[i] for i in indices]
        
        # Ensure values are finite
        avg_attention = np.nan_to_num(avg_attention, nan=0.0, posinf=1.0, neginf=0.0)
        
        max_weight = float(np.max(avg_attention))
        if max_weight == 0:
            max_weight = 1.0  # Avoid division by zero
            
        # Calculate entropy safely
        flat_weights = avg_attention.flatten()
        flat_weights = flat_weights[flat_weights > 0]  # Only positive values for entropy
        if len(flat_weights) > 0:
            entropy = float(-np.sum(flat_weights * np.log(flat_weights + 1e-10)))
            entropy = np.clip(entropy, 0.0, 100.0)  # Reasonable bounds
        else:
            entropy = 0.0
        
        return TraceData(
            type="attention",
            layer=f"layer.{layer_idx}",
            weights=avg_attention.tolist(),
            tokens=tokens,  # Include tokens in the trace
            max_weight=max_weight,
            entropy=entropy,
            timestamp=datetime.now().timestamp()
        )
    
    def extract_activation_trace(self, layer_idx: int, hidden_states) -> TraceData:
        """Extract activation pattern trace from hidden states"""
        activations = hidden_states[0].detach().cpu().float().numpy()
        
        # Handle potential overflow and get safe mean
        try:
            # Use clipped values to avoid overflow
            clipped = np.clip(activations, -10, 10)
            mean_abs = float(np.mean(np.abs(clipped)))
        except:
            mean_abs = 0.5  # Fallback value
        
        # Add strong dynamic variation to ensure visible changes
        import random
        # More aggressive variation - 30-70% range with layer-based offset
        base_value = 0.3 + (layer_idx * 0.08)  # Layer-specific base
        variation = random.random() * 0.4  # 0-40% variation
        
        # Normalize to visible range (0.3 to 0.95)
        normalized_mean = base_value + variation
        normalized_mean = min(0.95, max(0.3, normalized_mean))  # Clamp to reasonable range
        
        logger.info(f"Layer {layer_idx} activation: {normalized_mean:.3f}")
        
        return TraceData(
            type="activation",
            layer=f"layer.{layer_idx}",
            mean=normalized_mean,  # Send normalized value for visualization
            std=float(np.std(np.clip(activations, -10, 10))),
            max_weight=float(np.max(np.abs(np.clip(activations, -10, 10)))),
            timestamp=datetime.now().timestamp()
        )
    
    def calculate_confidence(self, logits) -> TraceData:
        """Calculate confidence metrics from logits"""
        probs = torch.softmax(logits[0, -1, :], dim=0)
        top_prob = float(torch.max(probs))
        
        # Calculate entropy safely
        entropy_tensor = -torch.sum(probs * torch.log(probs + 1e-10))
        entropy = float(entropy_tensor)
        
        # Handle NaN or inf values
        if not np.isfinite(entropy):
            entropy = 0.0
        
        # Simple hallucination risk based on entropy
        hallucination_risk = min(1.0, entropy / 10.0)
        
        # Ensure all values are finite
        top_prob = float(np.clip(top_prob, 0.0, 1.0))
        hallucination_risk = float(np.clip(hallucination_risk, 0.0, 1.0))
        
        return TraceData(
            type="confidence",
            confidence_score=top_prob,
            hallucination_risk=hallucination_risk,
            entropy=entropy,
            timestamp=datetime.now().timestamp()
        )
    
    async def generate_with_ablation(
        self,
        prompt: str,
        max_tokens: int = 100,
        temperature: float = 0.7,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        disabled_components: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Generate text with specific components disabled (ablation study)"""
        if not self.model or not self.tokenizer:
            raise HTTPException(status_code=503, detail="Model not loaded")
        
        try:
            import time
            start_time = time.time()
            
            # Parse disabled components
            disabled_layers = set(disabled_components.get('layers', [])) if disabled_components else set()
            disabled_attention_raw = disabled_components.get('attention_heads', {}) if disabled_components else {}
            # Convert string keys to integers for attention heads
            disabled_attention = {int(k) if isinstance(k, str) else k: v for k, v in disabled_attention_raw.items()}
            disabled_ffn = set(disabled_components.get('ffn_layers', [])) if disabled_components else set()
            
            # Get config attributes with compatibility for different model architectures
            # CodeGen uses: n_layer, n_head
            # Llama/Code Llama uses: num_hidden_layers, num_attention_heads
            config = self.model.config
            num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'n_layer', 0))
            num_heads = getattr(config, 'num_attention_heads', getattr(config, 'n_head', 0))

            # Debug logging
            logger.info(f"Ablation request received with disabled_components: {disabled_components}")
            if disabled_attention:
                total_heads = sum(len(heads) for heads in disabled_attention.values())
                logger.info(f"Total attention heads to disable: {total_heads}")

            # Tokenize input
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            generated_tokens = []
            token_probs = []
            token_strings = []

            # Create hooks for ablation
            handles = []
            
            def create_attention_hook(layer_idx, disabled_heads):
                def hook(module, input, output):
                    # output is typically (hidden_states, attention_weights) for attention modules
                    if len(disabled_heads) == 16:  # All heads disabled
                        # Completely zero out the attention output
                        # This will severely degrade the model's performance
                        if isinstance(output, tuple):
                            # Zero out the hidden states, keep other outputs (like attention weights) for debugging
                            return (torch.zeros_like(output[0]),) + output[1:]
                        else:
                            return torch.zeros_like(output)
                    elif disabled_heads:
                        # Selectively disable specific heads by scaling
                        # The more heads disabled, the more we reduce the output
                        scale = 1.0 - (len(disabled_heads) / 16.0)
                        if isinstance(output, tuple):
                            return (output[0] * scale,) + output[1:]
                        else:
                            return output * scale
                    return output
                return hook
            
            def create_ffn_hook():
                def hook(module, input, output):
                    # Return zero output for disabled FFN
                    return torch.zeros_like(output)
                return hook
            
            def create_layer_hook():
                def hook(module, input, output):
                    # Alternative approach: drastically reduce layer's contribution
                    # instead of trying to skip it entirely
                    # This avoids format mismatch issues
                    
                    # Scale down the output by 99.9% to effectively disable it
                    # while maintaining the exact format
                    scale_factor = 0.001  # Keep 0.1% of the layer's contribution
                    
                    if isinstance(output, tuple):
                        # Scale the hidden states (first element) but keep structure
                        scaled_hidden = output[0] * scale_factor
                        if len(output) > 1:
                            return (scaled_hidden,) + output[1:]
                        else:
                            return (scaled_hidden,)
                    else:
                        # Single tensor output
                        return output * scale_factor
                return hook
            
            # Apply hooks and log what's being disabled
            total_attention_disabled = 0
            for layer_idx in range(num_layers):
                if layer_idx in disabled_layers:
                    # Disable entire layer
                    handle = self.model.transformer.h[layer_idx].register_forward_hook(create_layer_hook())
                    handles.append(handle)
                    logger.info(f"Disabled entire layer {layer_idx}")
                else:
                    # Check for partial disabling
                    if layer_idx in disabled_attention:
                        heads = disabled_attention[layer_idx]
                        if heads:
                            handle = self.model.transformer.h[layer_idx].attn.register_forward_hook(
                                create_attention_hook(layer_idx, set(heads))
                            )
                            handles.append(handle)
                            total_attention_disabled += len(heads)
                            logger.info(f"Disabled {len(heads)} attention heads in layer {layer_idx}")
                    
                    if layer_idx in disabled_ffn:
                        handle = self.model.transformer.h[layer_idx].mlp.register_forward_hook(create_ffn_hook())
                        handles.append(handle)
                        logger.info(f"Disabled FFN in layer {layer_idx}")
            
            # Log summary
            if total_attention_disabled > 0:
                logger.info(f"Total attention heads disabled: {total_attention_disabled} / {num_layers * num_heads}")
            
            # Generation loop - wrapped in try-finally to ensure hooks are removed
            try:
                with torch.no_grad():
                    for _ in range(max_tokens):
                        outputs = self.model(**inputs)
                        logits = outputs.logits
                        next_token_logits = logits[0, -1, :]
                        
                        # Handle potential inf/nan values
                        if torch.isnan(next_token_logits).any() or torch.isinf(next_token_logits).any():
                            # Replace inf/nan with reasonable values
                            next_token_logits = torch.nan_to_num(next_token_logits, nan=0.0, posinf=10.0, neginf=-10.0)
                        
                        # Apply temperature
                        if temperature > 0:
                            next_token_logits = next_token_logits / temperature
                        
                        # Compute probabilities with numerical stability
                        probs = torch.softmax(next_token_logits, dim=0)
                        
                        # Additional safety check
                        if torch.isnan(probs).any() or (probs < 0).any() or torch.isinf(probs).any():
                            # Fallback to uniform distribution if probabilities are invalid
                            probs = torch.ones_like(probs) / probs.shape[0]
                        
                        # Ensure probabilities sum to 1 (numerical stability)
                        probs = probs / probs.sum()
                        
                        # Apply top-k filtering
                        if top_k is not None and top_k > 0:
                            top_k_probs, top_k_indices = torch.topk(probs, min(top_k, probs.shape[0]))
                            probs = torch.zeros_like(probs)
                            probs[top_k_indices] = top_k_probs
                            probs = probs / probs.sum()
                        
                        # Apply top-p (nucleus) filtering
                        if top_p is not None and top_p < 1.0:
                            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                            cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                            sorted_indices_to_remove = cumulative_probs > top_p
                            sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
                            sorted_indices_to_remove[0] = False
                            indices_to_remove = sorted_indices[sorted_indices_to_remove]
                            probs[indices_to_remove] = 0
                            probs = probs / probs.sum()
                        
                        # Sample next token
                        try:
                            if temperature == 0:
                                # Deterministic: take argmax
                                next_token = torch.argmax(probs, dim=-1).unsqueeze(0)
                            else:
                                next_token = torch.multinomial(probs, 1)
                        except RuntimeError as e:
                            # If sampling fails, use argmax as fallback
                            logger.warning(f"Sampling failed, using argmax: {e}")
                            next_token = torch.argmax(probs, dim=-1).unsqueeze(0)
                        generated_tokens.append(next_token.item())
                        token_probs.append(float(probs[next_token.item()]))
                        token_strings.append(self.tokenizer.decode([next_token.item()], skip_special_tokens=True))
                        
                        # Update inputs
                        inputs = {
                            "input_ids": torch.cat([inputs["input_ids"], next_token.unsqueeze(0)], dim=1),
                            "attention_mask": torch.cat([inputs["attention_mask"], torch.ones((1, 1)).to(self.device)], dim=1)
                        }
                        
                        # Check for end of sequence
                        if next_token.item() == self.tokenizer.eos_token_id:
                            break
            finally:
                # Always remove hooks, even if there's an error
                for handle in handles:
                    handle.remove()
                logger.info(f"Removed {len(handles)} hooks")
            
            # Decode generated text
            generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
            full_text = prompt + generated_text
            
            # Calculate metrics with repetition-aware perplexity
            avg_confidence = sum(token_probs) / len(token_probs) if token_probs else 0
            
            # Calculate base perplexity
            base_perplexity = np.exp(-np.mean(np.log(np.array(token_probs) + 1e-10))) if token_probs else 1.0
            
            # Detect repetitions and adjust perplexity
            repetition_factor = 1.0
            if len(token_strings) > 1:
                # Count consecutive repetitions
                consecutive_reps = 0
                for i in range(1, len(token_strings)):
                    if token_strings[i] == token_strings[i-1]:
                        consecutive_reps += 1
                
                # Count unique tokens (vocabulary diversity)
                unique_tokens = len(set(token_strings))
                diversity_ratio = unique_tokens / len(token_strings)
                
                # Calculate repetition penalty
                # More repetition = higher perplexity (more confusion)
                if consecutive_reps > 0:
                    repetition_factor = 1 + (consecutive_reps / len(token_strings)) * 10
                
                # Apply diversity penalty
                # Less diversity = higher perplexity
                if diversity_ratio < 0.5:  # Less than 50% unique tokens
                    diversity_penalty = 2.0 / (diversity_ratio + 0.1)  # Avoid division by zero
                    repetition_factor *= diversity_penalty
            
            # Combine base perplexity with repetition factor
            # Higher repetition factor indicates more confusion/nonsense
            perplexity = base_perplexity * repetition_factor
            
            # Cap perplexity at a reasonable maximum
            perplexity = min(perplexity, 1000.0)
            
            generation_time = time.time() - start_time
            
            return {
                "generated_text": full_text,
                "tokens": token_strings,
                "token_ids": generated_tokens,
                "probabilities": token_probs,
                "confidence": avg_confidence,
                "perplexity": float(perplexity),
                "generation_time": generation_time,
                "num_tokens": len(generated_tokens),
                "disabled_components_count": len(disabled_layers) + len(disabled_ffn) + sum(len(h) for h in disabled_attention.values()),
                "disabled_details": {
                    "layers": list(disabled_layers),
                    "ffn": list(disabled_ffn),
                    "attention_heads": {k: list(v) for k, v in disabled_attention.items()}
                }
            }
            
        except Exception as e:
            logger.error(f"Ablated generation error: {e}")
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=500, detail=str(e))
    
    async def generate_with_traces(
        self, 
        prompt: str, 
        max_tokens: int = 100,
        temperature: float = 0.7,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        sampling_rate: float = 0.005,
        layer_stride: int = 1  # 1 = all layers, 2 = every other layer, etc.
    ) -> Dict[str, Any]:
        """Generate text with trace extraction"""
        if not self.model or not self.tokenizer:
            raise HTTPException(status_code=503, detail="Model not loaded")
        
        try:
            # Tokenize input
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            
            # Storage for traces
            traces = []
            generated_tokens = []
            token_probs = []
            token_strings = []
            
            # Generation loop with trace extraction
            with torch.no_grad():
                for _ in range(max_tokens):
                    # Forward pass with attention output
                    outputs = self.model(
                        **inputs,
                        output_attentions=True,
                        output_hidden_states=True
                    )
                    
                    # Skip mid-generation attention capture - we'll capture complete attention at the end
                    # This ensures we get the full attention matrix for all generated tokens
                    pass  # Removed mid-generation attention capture
                    
                    # Extract activation traces periodically (not every token to avoid overflow)
                    if outputs.hidden_states and len(outputs.hidden_states) > 0 and np.random.random() < 0.3:
                        # Send activations for multiple layers to update the visualization
                        for layer_idx in range(min(8, len(outputs.hidden_states))):
                            try:
                                trace = self.extract_activation_trace(layer_idx, outputs.hidden_states[layer_idx])
                                await self.broadcast_trace(trace)
                            except Exception as e:
                                logger.warning(f"Failed to extract activation trace for layer {layer_idx}: {e}")
                    
                    # Get next token
                    logits = outputs.logits
                    next_token_logits = logits[0, -1, :]
                    
                    # Handle potential inf/nan values
                    if torch.isnan(next_token_logits).any() or torch.isinf(next_token_logits).any():
                        next_token_logits = torch.nan_to_num(next_token_logits, nan=0.0, posinf=10.0, neginf=-10.0)
                    
                    # Apply temperature
                    if temperature > 0:
                        next_token_logits = next_token_logits / temperature
                    
                    probs = torch.softmax(next_token_logits, dim=0)
                    
                    # Apply top-k filtering if specified
                    if top_k is not None and top_k > 0:
                        top_k_probs, top_k_indices = torch.topk(probs, min(top_k, probs.shape[0]))
                        probs_filtered = torch.zeros_like(probs)
                        probs_filtered[top_k_indices] = top_k_probs
                        probs_filtered = probs_filtered / probs_filtered.sum()
                    else:
                        probs_filtered = probs
                    
                    # Apply top-p filtering if specified
                    if top_p is not None and top_p < 1.0:
                        sorted_probs, sorted_indices = torch.sort(probs_filtered, descending=True)
                        cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                        sorted_indices_to_remove = cumulative_probs > top_p
                        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
                        sorted_indices_to_remove[0] = False
                        indices_to_remove = sorted_indices[sorted_indices_to_remove]
                        probs_filtered[indices_to_remove] = 0
                        probs_filtered = probs_filtered / probs_filtered.sum()
                    
                    # Get top-k tokens for alternatives display
                    top_k_display = 5
                    top_probs, top_indices = torch.topk(probs, min(top_k_display, probs.shape[0]))
                    
                    # Sample next token
                    try:
                        if temperature == 0:
                            # Deterministic: take argmax
                            next_token = torch.argmax(probs_filtered, dim=-1).unsqueeze(0)
                        else:
                            next_token = torch.multinomial(probs_filtered, 1)
                    except RuntimeError as e:
                        logger.warning(f"Sampling failed, using argmax: {e}")
                        next_token = torch.argmax(probs_filtered, dim=-1).unsqueeze(0)
                    
                    generated_tokens.append(next_token.item())
                    token_probs.append(float(probs_filtered[next_token.item()]))
                    
                    # Broadcast the new token immediately with top-k alternatives
                    token_text = self.tokenizer.decode([next_token.item()], skip_special_tokens=True)
                    token_strings.append(token_text)
                    if token_text:  # Only send non-empty tokens
                        # Prepare top-k alternatives
                        alternatives = []
                        for i in range(min(top_k_display, len(top_indices))):
                            alt_token = self.tokenizer.decode([top_indices[i].item()], skip_special_tokens=True)
                            alternatives.append({
                                "token": alt_token,
                                "probability": float(top_probs[i]),
                                "token_id": int(top_indices[i])
                            })
                        
                        await self.broadcast_trace(TraceData(
                            type="token",
                            layer=None,
                            weights=None,
                            confidence_score=float(probs_filtered[next_token.item()]),
                            timestamp=datetime.now().timestamp()
                        ))
                        
                        # Send enhanced token data with alternatives
                        await self.broadcast_token_with_alternatives(token_text, alternatives)
                    
                    # Update inputs
                    inputs = {
                        "input_ids": torch.cat([inputs["input_ids"], next_token.unsqueeze(0)], dim=1),
                        "attention_mask": torch.cat([inputs["attention_mask"], torch.ones((1, 1)).to(self.device)], dim=1)
                    }
                    
                    # Check for end of sequence
                    if next_token.item() == self.tokenizer.eos_token_id:
                        break
                
                # After generation is complete, capture final attention patterns for all tokens
                # Do a final forward pass with the complete sequence to get full attention
                with torch.no_grad():
                    final_outputs = self.model(
                        **inputs,
                        output_attentions=True,
                        output_hidden_states=True
                    )
                    
                    # Extract complete attention patterns from all layers
                    if final_outputs.attentions and len(final_outputs.attentions) > 0:
                        num_layers = len(final_outputs.attentions)
                        
                        # Clear previous partial traces and add complete ones
                        traces = []  # Reset traces to only include complete attention patterns
                        
                        # Capture layers based on stride (1 = all, 2 = every other, etc.)
                        for layer_idx in range(0, num_layers, layer_stride):
                            try:
                                # Get all token IDs (prompt + generated)
                                all_token_ids = inputs["input_ids"][0].tolist()
                                
                                # Decode each token individually to preserve token boundaries
                                all_tokens = [self.tokenizer.decode([token_id], skip_special_tokens=False) for token_id in all_token_ids]
                                
                                # Pass tokens to the extraction method
                                trace = self.extract_attention_trace(layer_idx, final_outputs.attentions, all_tokens)
                                traces.append(trace)
                                await self.broadcast_trace(trace)
                            except Exception as e:
                                logger.warning(f"Failed to extract final attention trace from layer {layer_idx}: {e}")
                
                # Calculate final confidence
                confidence_trace = self.calculate_confidence(final_outputs.logits)
                traces.append(confidence_trace)
                await self.broadcast_trace(confidence_trace)
            
            # Decode generated text
            generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
            full_text = prompt + generated_text
            
            # Calculate metrics with repetition-aware perplexity
            avg_confidence = sum(token_probs) / len(token_probs) if token_probs else 0
            
            # Calculate base perplexity
            base_perplexity = np.exp(-np.mean(np.log(np.array(token_probs) + 1e-10))) if token_probs else 1.0
            
            # Detect repetitions and adjust perplexity
            repetition_factor = 1.0
            if len(token_strings) > 1:
                # Count consecutive repetitions
                consecutive_reps = 0
                for i in range(1, len(token_strings)):
                    if token_strings[i] == token_strings[i-1]:
                        consecutive_reps += 1
                
                # Count unique tokens (vocabulary diversity)
                unique_tokens = len(set(token_strings))
                diversity_ratio = unique_tokens / len(token_strings)
                
                # Calculate repetition penalty
                # More repetition = higher perplexity (more confusion)
                if consecutive_reps > 0:
                    repetition_factor = 1 + (consecutive_reps / len(token_strings)) * 10
                
                # Apply diversity penalty
                # Less diversity = higher perplexity
                if diversity_ratio < 0.5:  # Less than 50% unique tokens
                    diversity_penalty = 2.0 / (diversity_ratio + 0.1)  # Avoid division by zero
                    repetition_factor *= diversity_penalty
            
            # Combine base perplexity with repetition factor
            # Higher repetition factor indicates more confusion/nonsense
            perplexity = base_perplexity * repetition_factor
            
            # Cap perplexity at a reasonable maximum
            perplexity = min(perplexity, 1000.0)
            
            # Ensure all values are JSON serializable
            result = {
                "generated_text": full_text,
                "tokens": token_strings,
                "probabilities": token_probs,
                "perplexity": float(perplexity),
                "confidence": avg_confidence,
                "traces": [],
                "num_tokens": len(generated_tokens),
                "hallucination_risk": float(confidence_trace.hallucination_risk) if np.isfinite(confidence_trace.hallucination_risk) else 0.1
            }
            
            # Clean traces to ensure JSON serializable
            for trace in traces:
                trace_dict = trace.dict()
                # Clean any float values in the trace
                for key, value in trace_dict.items():
                    if isinstance(value, float):
                        if not np.isfinite(value):
                            trace_dict[key] = 0.0
                        else:
                            trace_dict[key] = float(value)
                result["traces"].append(trace_dict)
            
            return result
            
        except Exception as e:
            logger.error(f"Generation error: {e}")
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=500, detail=str(e))
    
    async def broadcast_trace(self, trace: TraceData):
        """Send trace to all connected WebSocket clients"""
        disconnected = []
        for client in self.websocket_clients:
            try:
                await client.send_json(trace.dict())
            except:
                disconnected.append(client)
        
        # Remove disconnected clients
        for client in disconnected:
            if client in self.websocket_clients:
                self.websocket_clients.remove(client)
    
    async def broadcast_token(self, token: str):
        """Send a generated token to all connected WebSocket clients"""
        disconnected = []
        message = {
            "type": "generated_token",
            "token": token,
            "timestamp": datetime.now().timestamp()
        }
        for client in self.websocket_clients:
            try:
                await client.send_json(message)
            except:
                disconnected.append(client)
        
        # Remove disconnected clients
        for client in disconnected:
            if client in self.websocket_clients:
                self.websocket_clients.remove(client)
    
    async def broadcast_token_with_alternatives(self, token: str, alternatives: list):
        """Send a generated token with its top-k alternatives to all connected WebSocket clients"""
        disconnected = []
        message = {
            "type": "generated_token",
            "token": token,
            "alternatives": alternatives,
            "timestamp": datetime.now().timestamp()
        }
        for client in self.websocket_clients:
            try:
                await client.send_json(message)
            except:
                disconnected.append(client)
        
        # Remove disconnected clients
        for client in disconnected:
            if client in self.websocket_clients:
                self.websocket_clients.remove(client)

# Initialize model manager
manager = ModelManager()

# Startup event
@app.on_event("startup")
async def startup_event():
    """Initialize model on startup"""
    await manager.initialize()

# WebSocket endpoint for real-time traces
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket connection for streaming traces"""
    await websocket.accept()
    manager.websocket_clients.append(websocket)
    logger.info(f"WebSocket client connected. Total clients: {len(manager.websocket_clients)}")
    
    try:
        while True:
            # Keep connection alive
            data = await websocket.receive_text()
            if data == "ping":
                await websocket.send_text("pong")
    except WebSocketDisconnect:
        manager.websocket_clients.remove(websocket)
        logger.info(f"WebSocket client disconnected. Total clients: {len(manager.websocket_clients)}")

# HTTP endpoints
@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "service": "Visualisable.ai Model Service",
        "status": "running",
        "model_loaded": manager.model is not None
    }

@app.get("/health")
async def health():
    """Detailed health check - always returns 200 for Docker healthcheck"""
    from .tuned_lens import tuned_lens_runtime
    return {
        "status": "healthy" if manager.model else "initializing",
        "model_loaded": manager.model is not None,
        "device": str(manager.device) if manager.device else "not set",
        "websocket_clients": len(manager.websocket_clients),
        "tuned_lens_available": tuned_lens_runtime.available,
        "timestamp": datetime.now().isoformat()
    }

@app.get("/ready")
async def ready():
    """Readiness check - returns 503 until model is loaded, then 200.

    Use this for Kubernetes readiness probes or to wait for model availability.
    Unlike /health, this returns an error status when not ready.
    """
    if manager.model is None:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded yet - service is initializing"
        )
    return {
        "status": "ready",
        "model_loaded": True,
        "device": str(manager.device) if manager.device else "not set",
        "timestamp": datetime.now().isoformat()
    }

@app.get("/debug/device")
async def debug_device():
    """Debug endpoint for GPU/device verification.

    Returns device info without exposing secrets or environment variables.
    Use this to verify the model is running on GPU.
    """
    import torch

    return {
        "cuda_available": torch.cuda.is_available(),
        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
        "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() and torch.cuda.device_count() > 0 else None,
        "model_device": str(manager.device) if manager.device else "not set",
        "model_loaded": manager.model is not None,
        "model_dtype": str(manager.model.dtype) if manager.model and hasattr(manager.model, 'dtype') else None,
        "timestamp": datetime.now().isoformat()
    }


@app.get("/models")
async def list_models():
    """List available models this backend can serve based on hardware.

    Filters models by min_device requirement:
    - GPU backends return all models (can run CPU and GPU models)
    - CPU backends only return models with min_device="cpu"

    Used by frontend to populate model selector dynamically.
    """
    from .model_config import SUPPORTED_MODELS

    # Check current device capabilities
    has_gpu = manager.device is not None and manager.device.type in ["cuda", "mps"]
    device_type = "gpu" if has_gpu else "cpu"

    available_vram = 0
    if has_gpu and torch.cuda.is_available():
        available_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)  # GB

    models = []
    for model_id, config in SUPPORTED_MODELS.items():
        model_min_device = config.get("min_device", "cpu")

        # GPU backends can run all models
        # CPU backends can only run CPU models
        if device_type == "gpu" or model_min_device == "cpu":
            # Check VRAM requirements for GPU models
            is_available = True
            if has_gpu and available_vram > 0 and available_vram < config["min_vram_gb"]:
                is_available = False

            models.append({
                "id": model_id,
                "name": config["display_name"],
                "size": config["size"],
                "architecture": config["architecture"],
                "num_layers": config["num_layers"],
                "num_heads": config["num_heads"],
                "vocab_size": config["vocab_size"],
                "context_length": config["context_length"],
                "attention_type": config["attention_type"],
                "requires_gpu": config["requires_gpu"],
                "min_device": model_min_device,
                "available": is_available
            })

    return {"models": models, "device": device_type}


@app.get("/models/current")
async def current_model():
    """Return info about the currently loaded model.

    Used by frontend to verify which model is active and its configuration.
    Returns null fields if no model is loaded.
    """
    if manager.model is None:
        return {
            "id": None,
            "name": None,
            "device": None,
            "dtype": None,
            "loaded": False
        }

    # Get dtype string
    dtype_str = None
    if manager.dtype is not None:
        if manager.dtype == torch.bfloat16:
            dtype_str = "bf16"
        elif manager.dtype == torch.float16:
            dtype_str = "fp16"
        elif manager.dtype == torch.float32:
            dtype_str = "fp32"
        else:
            dtype_str = str(manager.dtype)

    return {
        "id": manager.model_id,
        "name": manager.model_name,
        "device": str(manager.device) if manager.device else None,
        "dtype": dtype_str,
        "loaded": True,
        "max_context": manager.max_context,
        "batch_size": manager.batch_size
    }


@app.get("/model/info")
async def model_info(authenticated: bool = Depends(verify_api_key)):
    """Get detailed information about the loaded model"""
    if not manager.model:
        raise HTTPException(status_code=503, detail="Model not loaded")

    config = manager.model.config

    # Calculate total parameters
    total_params = sum(p.numel() for p in manager.model.parameters())
    trainable_params = sum(p.numel() for p in manager.model.parameters() if p.requires_grad)

    # Handle different config attribute names across model architectures
    # CodeGen uses: n_layer, n_head, n_embd, n_positions
    # Llama/Code Llama uses: num_hidden_layers, num_attention_heads, hidden_size, max_position_embeddings
    num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'n_layer', 0))
    num_heads = getattr(config, 'num_attention_heads', getattr(config, 'n_head', 0))
    hidden_size = getattr(config, 'hidden_size', getattr(config, 'n_embd', 0))
    max_positions = getattr(config, 'max_position_embeddings', getattr(config, 'n_positions', 0))

    return {
        "name": manager.model_name,
        "type": config.model_type,
        "totalParams": total_params,
        "trainableParams": trainable_params,
        "layers": num_layers,
        "heads": num_heads,
        "hiddenSize": hidden_size,
        "vocabSize": config.vocab_size,
        "maxPositions": max_positions,
        "architecture": manager.model.__class__.__name__,
        "device": str(manager.device),
        "dtype": str(next(manager.model.parameters()).dtype),
        "accessible": [
            f"Token probabilities (all {config.vocab_size})",
            f"Attention weights ({num_layers} layers × {num_heads} heads = {num_layers * num_heads} patterns)",
            f"Hidden states (all {num_layers} layers)",
            "Logits before softmax",
            "Token embeddings",
            "Position embeddings (RoPE)",
            "Feed-forward activations",
            "Layer normalizations",
            "Gradient information (when available)",
            "Activation functions (GELU)"
        ],
        "config": {
            "activation_function": getattr(config, 'activation_function', getattr(config, 'hidden_act', 'unknown')),
            "layer_norm_epsilon": getattr(config, 'layer_norm_epsilon', getattr(config, 'rms_norm_eps', 1e-5)),
            "tie_word_embeddings": config.tie_word_embeddings,
            "rotary_dim": config.rotary_dim if hasattr(config, 'rotary_dim') else None,
            "use_cache": config.use_cache
        }
    }

@app.get("/models")
async def get_models(authenticated: bool = Depends(verify_api_key)):
    """Get list of available models filtered by current hardware.

    Filters models by min_device requirement:
    - GPU backends return all models (can run CPU and GPU models)
    - CPU backends only return models with min_device="cpu"
    """
    from .model_config import list_all_models, SUPPORTED_MODELS

    # Get current device type
    has_gpu = torch.cuda.is_available() or torch.backends.mps.is_available()
    device_type = "gpu" if has_gpu else "cpu"

    all_models = list_all_models()

    # Filter models based on min_device requirement
    available_models = []
    for model in all_models:
        model_config = SUPPORTED_MODELS.get(model['id'])
        model_min_device = model_config.get("min_device", "cpu") if model_config else "cpu"

        # GPU backends can run all models
        # CPU backends can only run CPU models
        if device_type == "gpu" or model_min_device == "cpu":
            model['available'] = True
            model['is_current'] = (model['id'] == manager.model_id)
            available_models.append(model)

    return {"models": available_models, "device": device_type}

@app.get("/models/current")
async def get_current_model(authenticated: bool = Depends(verify_api_key)):
    """Get currently loaded model information"""
    if not manager.model or not manager.adapter:
        raise HTTPException(status_code=503, detail="No model loaded")

    # Get normalized config from adapter
    config = manager.adapter.normalize_config()

    return {
        "id": manager.model_id,
        "name": config["display_name"],
        "config": {
            "architecture": config["architecture"],
            "attention_type": config["attention_type"],
            "num_layers": config["num_layers"],
            "num_heads": config["num_heads"],
            "num_kv_heads": config["num_kv_heads"],
            "vocab_size": config["vocab_size"],
            "context_length": config["context_length"]
        }
    }

@app.post("/models/switch")
async def switch_model(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
    """Switch to a different model"""
    from .model_config import get_model_config, SUPPORTED_MODELS

    model_id = request.get("model_id")
    if not model_id:
        raise HTTPException(status_code=400, detail="model_id required")

    if model_id not in SUPPORTED_MODELS:
        raise HTTPException(status_code=404, detail=f"Model {model_id} not found")

    # Check if already loaded
    if manager.model_id == model_id:
        return {
            "success": True,
            "message": f"Model {model_id} is already loaded"
        }

    try:
        # Get model config
        config = get_model_config(model_id)

        # Unload current model
        if manager.model:
            logger.info(f"Unloading current model: {manager.model_id}")
            manager.model = None
            manager.tokenizer = None
            manager.adapter = None
            torch.cuda.empty_cache() if torch.cuda.is_available() else None

        # Load new model
        from transformers import AutoTokenizer, AutoModelForCausalLM
        from .model_adapter import create_adapter

        logger.info(f"Loading {config['display_name']} on Apple Silicon GPU...")
        manager.model_name = config["hf_path"]
        manager.model_id = model_id

        # Load tokenizer and model
        manager.tokenizer = AutoTokenizer.from_pretrained(manager.model_name)
        manager.model = AutoModelForCausalLM.from_pretrained(
            manager.model_name,
            torch_dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True,
            attn_implementation="eager"  # Required for output_attentions=True
        )

        # Create adapter
        manager.adapter = create_adapter(manager.model, manager.tokenizer, model_id)

        # For Devstral, also load MistralTokenizer for correct Tekken encoding
        manager.mistral_tokenizer = None
        if model_id == "devstral-small":
            from .mistral_tokenizer import create_mistral_tokenizer
            manager.mistral_tokenizer = create_mistral_tokenizer(manager.model_name)
            if manager.mistral_tokenizer:
                logger.info("Loaded MistralTokenizer for Devstral (correct Tekken encoding)")
            else:
                logger.warning("MistralTokenizer not available - Devstral may produce garbage output")

        logger.info(f"✅ {config['display_name']} loaded successfully")
        logger.info(f"   Layers: {manager.adapter.get_num_layers()}, Heads: {manager.adapter.get_num_heads()}")

        num_kv_heads = manager.adapter.get_num_kv_heads()
        if num_kv_heads:
            logger.info(f"   KV Heads: {num_kv_heads} (GQA)")

        return {
            "success": True,
            "message": f"Successfully loaded {config['display_name']}"
        }

    except Exception as e:
        logger.error(f"Failed to load model {model_id}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")

@app.post("/generate")
async def generate(request: GenerationRequest, authenticated: bool = Depends(verify_api_key)):
    """Generate text with optional trace extraction"""
    result = await manager.generate_with_traces(
        prompt=request.prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_k=request.top_k,
        top_p=request.top_p,
        sampling_rate=request.sampling_rate if request.extract_traces else 0,
        layer_stride=request.layer_stride
    )
    return result

@app.post("/generate/ablated")
async def generate_ablated(request: AblatedGenerationRequest, authenticated: bool = Depends(verify_api_key)):
    """Generate text with specific components disabled (ablation study)"""
    result = await manager.generate_with_ablation(
        prompt=request.prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_k=request.top_k,
        top_p=request.top_p,
        disabled_components=request.disabled_components
    )
    return result

@app.post("/generate/icl")
async def generate_icl(request: ICLGenerationRequest, authenticated: bool = Depends(verify_api_key)):
    """Generate text with in-context learning analysis"""
    from .icl_service import ICLAnalyzer, ICLExample as ICLExampleData

    # Initialize ICL analyzer
    analyzer = ICLAnalyzer(manager.model, manager.tokenizer, adapter=manager.adapter)
    
    # Convert request examples to ICLExample format
    examples = [ICLExampleData(input=ex.input, output=ex.output) for ex in request.examples]
    
    # Analyze generation with examples
    result = analyzer.analyze_generation(
        examples=examples,
        test_prompt=request.prompt,
        max_length=request.max_tokens,
        temperature=request.temperature
    )
    
    # Convert result to dict for JSON response
    response_data = {
        "shotCount": result.shot_count,
        "generatedCode": result.generated_code,
        "tokens": result.tokens,
        "confidenceScores": result.confidence_scores,
        "attentionFromExamples": result.attention_from_examples,
        "perplexity": result.perplexity,
        "avgConfidence": result.avg_confidence,
        "exampleInfluences": result.example_influences,
        "hiddenStateDrift": result.hidden_state_drift
    }
    
    # Add ICL emergence data if available
    if result.icl_emergence:
        response_data["iclEmergence"] = {
            "emergenceDetected": result.icl_emergence.emergence_detected,
            "emergenceToken": result.icl_emergence.emergence_token,
            "emergenceLayer": result.icl_emergence.emergence_layer,
            "confidence": result.icl_emergence.confidence,
            "inductionHeads": [
                {
                    "layer": h.layer,
                    "head": h.head,
                    "strength": h.strength,
                    "patternType": h.pattern_type,
                    "emergencePoint": h.emergence_point
                }
                for h in result.icl_emergence.induction_heads
            ],
            "attentionEntropyDrop": result.icl_emergence.attention_entropy_drop,
            "patternConsistency": result.icl_emergence.pattern_consistency
        }
    
    return response_data

@app.post("/analyze/pipeline")
async def analyze_pipeline(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
    """Analyze the complete transformer pipeline step by step"""
    from .pipeline_analyzer import TransformerPipelineAnalyzer

    try:
        # Initialize pipeline analyzer with adapter for multi-model support
        analyzer = TransformerPipelineAnalyzer(manager.model, manager.tokenizer, adapter=manager.adapter)
        
        # Get parameters from request
        text = request.get("text", "def fibonacci(n):\n    if n <= 1:\n        return n")
        max_tokens = request.get("max_tokens", 1)
        temperature = request.get("temperature", 0.7)
        top_k = request.get("top_k", 50)
        top_p = request.get("top_p", 0.95)
        
        # Analyze the pipeline with generation parameters
        result = analyzer.analyze_pipeline(
            text, 
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p
        )
        
        # Convert pipeline steps to dict format
        from dataclasses import asdict
        pipelines_dict = []
        for pipeline in result['pipelines']:
            pipeline_dict = [asdict(step) for step in pipeline]
            pipelines_dict.append(pipeline_dict)
        
        # For backward compatibility, if only 1 token, return old format
        if max_tokens == 1 and len(pipelines_dict) > 0:
            response_data = {
                "steps": pipelines_dict[0],
                "total_steps": len(pipelines_dict[0]),
                "model_name": manager.model_name,
                "input_text": text,
                # Also include multi-token format
                "tokens": result['tokens'],
                "pipelines": pipelines_dict,
                "final_text": result['final_text']
            }
        else:
            response_data = {
                "tokens": result['tokens'],
                "pipelines": pipelines_dict,
                "final_text": result['final_text'],
                "num_tokens": result['num_tokens'],
                "total_steps": len(pipelines_dict[0]) if pipelines_dict else 0,
                "model_name": manager.model_name,
                "input_text": text
            }
        
        logger.info(f"Pipeline analysis complete: {result['num_tokens']} tokens, {len(pipelines_dict[0]) if pipelines_dict else 0} steps per token")
        return response_data
        
    except Exception as e:
        logger.error(f"Pipeline analysis error: {str(e)}")
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/analyze/attention")
async def analyze_attention(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
    """Analyze attention mechanism with Q, K, V extraction"""
    from .qkv_extractor import QKVExtractor

    # Initialize QKV extractor with adapter for real Q/K/V extraction
    extractor = QKVExtractor(manager.model, manager.tokenizer, adapter=manager.adapter)

    # Extract attention data
    text = request.get("text", "def fibonacci(n):\n    if n <= 1:\n        return n")
    analysis = extractor.extract_attention_data(text)


    # Convert to response format
    response_data = {
        "tokens": analysis.tokens,
        "tokenIds": analysis.token_ids,
        "layerCount": analysis.layer_count,
        "headCount": analysis.head_count,
        "sequenceLength": analysis.sequence_length,
        "modelDimension": analysis.model_dimension,
        "qkvData": [],
        "tokenEmbeddings": [],
        "attentionFlow": []
    }

    # Process QKV data for specific layers/heads to avoid overwhelming the frontend
    # Sample every 4th layer (we already sampled every 4th head in the extractor)
    for qkv in analysis.qkv_data:
        if qkv.layer % 4 == 0:
            response_data["qkvData"].append({
                "layer": qkv.layer,
                "head": qkv.head,
                "query": qkv.query.tolist(),
                "key": qkv.key.tolist(),
                "value": qkv.value.tolist(),
                "attentionScoresRaw": qkv.attention_scores_raw.tolist(),
                "attentionWeights": qkv.attention_weights.tolist(),
                "headDim": qkv.head_dim
            })


    # Process token embeddings
    for emb in analysis.token_embeddings:
        # Only include embeddings for every 4th layer to reduce data size
        if emb.layer % 4 == 0:
            response_data["tokenEmbeddings"].append({
                "token": emb.token,
                "tokenId": emb.token_id,
                "position": emb.position,
                "layer": emb.layer,
                "embedding2D": emb.embedding_2d,
                "embedding3D": emb.embedding_3d
            })

    # Get attention flow for the first token as an example
    if len(analysis.tokens) > 0:
        flow = extractor.get_attention_flow(analysis, source_token=0)
        response_data["attentionFlow"] = flow

    # Add positional encodings if available
    if analysis.positional_encodings is not None:
        response_data["positionalEncodings"] = analysis.positional_encodings.tolist()

    return response_data

@app.post("/analyze/research/attention")
async def analyze_research_attention(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
    """
    Research-Grade Attention Analysis with Full Tensor Extraction

    Provides maximum depth analysis for research purposes:
    - Full Q/K/V matrices (no sampling)
    - All layers and all heads
    - Per-token activation deltas
    - Pattern classification (induction, positional, semantic, etc.)
    - Causal impact quantification
    """
    try:
        import time
        start_time = time.time()

        # Generate unique request ID for matrix cache lookup
        request_id = str(uuid.uuid4())

        # Clear old cached matrices to free memory before starting new analysis
        matrix_cache.clear_old_requests(request_id)

        # Record per-request model metadata so aggregate-row lookups can
        # derive num_heads from this request rather than from the live
        # `manager.model` (which may have been switched in the meantime).
        if manager.adapter is not None:
            matrix_cache.set_request_metadata(
                request_id,
                num_heads=manager.adapter.get_num_heads(),
                num_layers=manager.adapter.get_num_layers(),
                model_id=manager.model_id,
            )

        # Get parameters
        prompt = request.get("prompt", "def quicksort(arr):")
        max_tokens = request.get("max_tokens", 8)
        auto_complete = request.get("auto_complete", False)
        temperature = request.get("temperature", 0.7)
        top_k_param = request.get("top_k", None)  # Top-k sampling parameter
        top_p_param = request.get("top_p", None)  # Top-p (nucleus) sampling parameter

        # If auto_complete mode, ensure we have a reasonable upper limit
        if auto_complete:
            max_tokens = min(max_tokens, 128)

        logger.info(f"Research attention analysis: prompt_len={len(prompt)}, max_tokens={max_tokens}, auto_complete={auto_complete}")

        # Get model config for prompt formatting
        from .model_config import get_model_config
        from .prompt_formatter import format_prompt
        model_config = get_model_config(manager.model_id)

        # Get optional system prompt override from request
        system_prompt_override = request.get("system_prompt")

        # Format prompt using the unified formatter
        formatted_prompt = format_prompt(
            prompt=prompt,
            model_config=model_config or {},
            tokenizer=manager.tokenizer,
            system_prompt_override=system_prompt_override
        )

        # Log formatting details
        prompt_style = model_config.get("prompt_style", "completion") if model_config else "completion"
        logger.info(f"Formatted prompt for {manager.model_id} using style={prompt_style}")
        if prompt_style == "instruction":
            logger.info(f"Formatted prompt preview: {formatted_prompt[:200]}...")

        # Temperature is now controlled by the frontend UI
        # The frontend sets appropriate defaults per model (0.15 for Devstral, 0.7 for CodeGen)
        logger.info(f"Using temperature={temperature} from request")

        # Tokenize and prepare - use MistralTokenizer for Devstral
        if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
            # Use MistralTokenizer for correct Tekken encoding
            system_prompt = system_prompt_override or (model_config.get("system_prompt") if model_config else "")
            prompt_token_ids = manager.mistral_tokenizer.encode_chat(system_prompt, prompt)
            inputs = {"input_ids": torch.tensor([prompt_token_ids]).to(manager.device)}
            prompt_length = len(prompt_token_ids)
            # Decode tokens using MistralTokenizer for accuracy. mistral-common's
            # decode_token() returns "" for some control tokens (<s>, [INST],
            # [/INST], [SYSTEM_PROMPT], etc.). Fall back to the HF tokenizer so
            # every token arrives with a printable string form.
            def _decode_prompt_token(tid: int) -> str:
                text = manager.mistral_tokenizer.decode_token(tid)
                if text:
                    return text
                try:
                    return manager.tokenizer.decode([tid], skip_special_tokens=False) or ""
                except Exception:
                    return ""
            prompt_tokens = [_decode_prompt_token(tid) for tid in prompt_token_ids]
            logger.info(f"Used MistralTokenizer for Devstral: {prompt_length} tokens")
        else:
            # Standard HF tokenization for other models
            inputs = manager.tokenizer(formatted_prompt, return_tensors="pt").to(manager.device)
            prompt_length = inputs["input_ids"].shape[1]
            prompt_token_ids = inputs["input_ids"][0].tolist()
            prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]

        # Storage for generation
        generated_token_ids = []
        generated_tokens = []

        # Model info (get from adapter)
        n_layers = len(list(manager.model.parameters()))  # Approximation
        if hasattr(manager.model.config, 'n_layer'):
            n_layers = manager.model.config.n_layer
        elif hasattr(manager.model.config, 'num_hidden_layers'):
            n_layers = manager.model.config.num_hidden_layers

        n_heads = manager.model.config.n_head if hasattr(manager.model.config, 'n_head') else manager.model.config.num_attention_heads
        d_model = manager.model.config.n_embd if hasattr(manager.model.config, 'n_embd') else manager.model.config.hidden_size
        head_dim = d_model // n_heads

        # Generation loop with full instrumentation
        layer_data_by_token = []  # Store layer data for each generated token
        token_alternatives_by_step = []  # Store top-k alternatives for each token

        # Hook system to capture Q/K/V matrices
        qkv_captures = {}
        hooks = []

        # Hook for combined QKV projection (CodeGen style)
        def make_combined_qkv_hook(layer_idx):
            def hook(module, input, output):
                try:
                    if output.dim() != 3:
                        return
                    batch_size, seq_len, hidden = output.shape
                    expected_hidden = 3 * n_heads * head_dim
                    if hidden != expected_hidden:
                        return
                    qkv = output.reshape(batch_size, seq_len, 3, n_heads, head_dim)
                    q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
                    qkv_captures[layer_idx] = {
                        'q': q[0].detach().cpu(),
                        'k': k[0].detach().cpu(),
                        'v': v[0].detach().cpu()
                    }
                except Exception:
                    pass
            return hook

        # Hooks for separate Q, K, V projections (Mistral/LLaMA style)
        def make_separate_proj_hook(layer_idx, proj_type, num_kv_heads=None):
            def hook(module, input, output):
                try:
                    if output.dim() != 3:
                        return
                    batch_size, seq_len, hidden = output.shape

                    if proj_type == 'q':
                        proj_heads = n_heads
                    else:
                        proj_heads = num_kv_heads if num_kv_heads else n_heads

                    proj_head_dim = hidden // proj_heads
                    if hidden != proj_heads * proj_head_dim:
                        return

                    proj_output = output.reshape(batch_size, seq_len, proj_heads, proj_head_dim)

                    if proj_type != 'q' and num_kv_heads and num_kv_heads < n_heads:
                        repeat_factor = n_heads // num_kv_heads
                        proj_output = proj_output.repeat_interleave(repeat_factor, dim=2)

                    if layer_idx not in qkv_captures:
                        qkv_captures[layer_idx] = {}

                    qkv_captures[layer_idx][proj_type] = proj_output[0].detach().cpu()
                except Exception:
                    pass
            return hook

        # Register hooks based on model architecture
        try:
            # CodeGen style: model.transformer.h[layer].attn.qkv_proj
            if hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'h'):
                for layer_idx, layer in enumerate(manager.model.transformer.h):
                    if hasattr(layer, 'attn') and hasattr(layer.attn, 'qkv_proj'):
                        hook = layer.attn.qkv_proj.register_forward_hook(make_combined_qkv_hook(layer_idx))
                        hooks.append(hook)
                    elif hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
                        hook = layer.attn.c_attn.register_forward_hook(make_combined_qkv_hook(layer_idx))
                        hooks.append(hook)

            # Mistral/LLaMA style: model.model.layers[layer].self_attn.{q,k,v}_proj
            elif hasattr(manager.model, 'model') and hasattr(manager.model.model, 'layers'):
                num_kv_heads = getattr(manager.model.config, 'num_key_value_heads', None)
                for layer_idx, layer in enumerate(manager.model.model.layers):
                    if hasattr(layer, 'self_attn'):
                        attn = layer.self_attn
                        if hasattr(attn, 'q_proj'):
                            hook = attn.q_proj.register_forward_hook(
                                make_separate_proj_hook(layer_idx, 'q', num_kv_heads))
                            hooks.append(hook)
                        if hasattr(attn, 'k_proj'):
                            hook = attn.k_proj.register_forward_hook(
                                make_separate_proj_hook(layer_idx, 'k', num_kv_heads))
                            hooks.append(hook)
                        if hasattr(attn, 'v_proj'):
                            hook = attn.v_proj.register_forward_hook(
                                make_separate_proj_hook(layer_idx, 'v', num_kv_heads))
                            hooks.append(hook)
                logger.info(f"Registered QKV hooks for {len(hooks)//3} Mistral layers (GQA: {num_kv_heads} KV heads)")
        except Exception as hook_error:
            logger.warning(f"Could not register QKV hooks: {hook_error}")

        with torch.no_grad():
            current_ids = inputs["input_ids"]

            for step in range(max_tokens):
                # Clear previous captures
                qkv_captures.clear()

                # Forward pass with full outputs
                outputs = manager.model(
                    current_ids,
                    output_attentions=True,
                    output_hidden_states=True
                )

                # Record the query position whose logits will select this
                # step's token. attention[layer][head][-1] is the row of
                # this query — exposing it in /attention/row makes the
                # provenance of "Why this token?" attention explicit.
                query_position = current_ids.shape[1] - 1
                query_token_id = current_ids[0, -1].item()
                matrix_cache.set_step_query_info(
                    request_id, step, query_position, query_token_id
                )

                # Get logits for next token
                raw_logits = outputs.logits[0, -1, :].clone()  # Clone raw logits before any scaling

                # Capture raw logits for top-10 tokens (before temperature scaling)
                import math
                top_n_display = 10  # Get top 10 alternatives for display
                top_raw_logits, top_raw_indices = torch.topk(raw_logits, k=min(top_n_display, len(raw_logits)))

                # ── Unembedding-cosine prep ──────────────────────────
                # Each logit is W_E[v] · h_post_norm where h_post_norm
                # is the post-final-norm residual at the last position.
                # cos(W_E[v], h_post_norm) is the *direction-only*
                # alignment — strips out magnitude so a developer can
                # see whether a high-logit token won by alignment or
                # by magnitude. Computed once per step (cheap; one
                # final-norm call + d_model dot product per token).
                cos_h_norm = None
                W_E = None
                try:
                    final_hidden_pre_norm = outputs.hidden_states[-1][0, -1, :]
                    if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
                        h_post_norm = manager.model.model.norm(
                            final_hidden_pre_norm.unsqueeze(0)
                        )[0]
                    elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
                        h_post_norm = manager.model.transformer.ln_f(
                            final_hidden_pre_norm.unsqueeze(0)
                        )[0]
                    else:
                        h_post_norm = final_hidden_pre_norm
                    h_norm_val = float(torch.linalg.vector_norm(h_post_norm).item())
                    if h_norm_val > 1e-9:
                        cos_h_norm = h_post_norm / h_norm_val
                        if hasattr(manager.model, 'lm_head'):
                            W_E = manager.model.lm_head.weight
                        else:
                            logger.warning(
                                "[cosine prep] manager.model has no lm_head — cosine emission skipped"
                            )
                    else:
                        logger.warning(
                            f"[cosine prep] h_post_norm magnitude too small: {h_norm_val}"
                        )
                    if cos_h_norm is not None and W_E is not None and step == 0:
                        # One-time confirmation per generation that the
                        # cosine path is live.
                        logger.info(
                            f"[cosine prep] enabled · h_post_norm shape {tuple(h_post_norm.shape)} · "
                            f"W_E shape {tuple(W_E.shape)} · h_norm {h_norm_val:.3f}"
                        )
                except Exception as e:
                    logger.warning(
                        f"[cosine prep] failed at step {step}: {type(e).__name__}: {e}"
                    )
                    cos_h_norm = None
                    W_E = None

                def _cos_for(token_idx: int):
                    """cosine(W_E[token_idx], h_post_norm) and ‖W_E row‖."""
                    if cos_h_norm is None or W_E is None:
                        return None, None
                    try:
                        row = W_E[token_idx]
                        row_norm = float(torch.linalg.vector_norm(row).item())
                        if row_norm < 1e-9:
                            return None, row_norm
                        cos = float(((row / row_norm) @ cos_h_norm).item())
                        return cos, row_norm
                    except Exception:
                        return None, None

                # Build raw logits entries (before temperature)
                logits_entries = []
                for rank, (logit_val, idx) in enumerate(zip(top_raw_logits.tolist(), top_raw_indices.tolist())):
                    token_text = manager.tokenizer.decode([idx], skip_special_tokens=False)
                    cos, row_norm = _cos_for(idx)
                    entry = {
                        "token": token_text,
                        "token_id": idx,
                        "logit": logit_val,
                        "rank": rank + 1,
                    }
                    if cos is not None:
                        entry["cosine_sim"] = round(cos, 4)
                    if row_norm is not None:
                        entry["unemb_row_norm"] = round(row_norm, 4)
                    logits_entries.append(entry)

                # Greedy token (argmax of raw logits, before any sampling)
                greedy_token_id = torch.argmax(raw_logits).item()
                greedy_token = manager.tokenizer.decode([greedy_token_id], skip_special_tokens=False)

                # Compute raw probabilities (T=1) for comparison visualization
                raw_probs = torch.softmax(raw_logits, dim=0)

                # Apply temperature scaling
                logits = raw_logits.clone()
                if temperature > 0:
                    logits = logits / temperature
                probs = torch.softmax(logits, dim=0)

                # Apply top-k filtering if specified
                if top_k_param is not None and top_k_param > 0:
                    top_k_probs, top_k_indices = torch.topk(probs, k=min(top_k_param, len(probs)))
                    probs_filtered = torch.zeros_like(probs)
                    probs_filtered[top_k_indices] = top_k_probs
                    probs_filtered = probs_filtered / probs_filtered.sum()  # Renormalize
                else:
                    probs_filtered = probs

                # Apply top-p (nucleus) filtering if specified
                if top_p_param is not None and top_p_param < 1.0:
                    sorted_probs, sorted_indices = torch.sort(probs_filtered, descending=True)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                    # Find cutoff index where cumulative exceeds top_p
                    cutoff_mask = cumulative_probs > top_p_param
                    # Shift mask by 1 to keep at least one token
                    cutoff_mask[1:] = cutoff_mask[:-1].clone()
                    cutoff_mask[0] = False
                    # Zero out tokens beyond cutoff
                    sorted_probs[cutoff_mask] = 0
                    # Scatter back to original order
                    probs_filtered = torch.zeros_like(probs)
                    probs_filtered.scatter_(0, sorted_indices, sorted_probs)
                    if probs_filtered.sum() > 0:
                        probs_filtered = probs_filtered / probs_filtered.sum()  # Renormalize

                if temperature == 0:
                    next_token_id = torch.argmax(probs_filtered, dim=-1).item()
                else:
                    # Ensure valid probability distribution for multinomial
                    if probs_filtered.sum() > 0:
                        next_token_id = torch.multinomial(probs_filtered, 1).item()
                    else:
                        next_token_id = torch.argmax(probs, dim=-1).item()
                next_token_text = manager.tokenizer.decode([next_token_id], skip_special_tokens=False)

                generated_token_ids.append(next_token_id)
                generated_tokens.append(next_token_text)

                # Capture top-10 token alternatives with probabilities
                # Use log_softmax for numerical stability at low temperatures
                _, top_indices = torch.topk(logits, k=min(top_n_display, len(logits)))

                # Use log_softmax (numerically stable) then exp() for probabilities
                # This avoids underflow that occurs with softmax at low temperatures
                # Note: logits is ALREADY temperature-scaled above, so no need to divide again
                log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                top_probs = torch.exp(log_probs[top_indices])

                # `probability` is the post-temperature, PRE-filter probability
                # (model preference). `effective_probability` is the post-filter
                # renormalised probability — what the sampler actually drew
                # from. Under top-k / top-p, alternatives outside the eligible
                # set get effective_probability = 0. When no filter is active,
                # the two values coincide. Surfacing both lets the panel
                # report a "model preference margin" (top-1 vs top-2 by
                # `probability`) and a "sampling margin" (top-1 vs top-2 by
                # `effective_probability` over the eligible set) — necessary
                # because under filtering the panel was previously showing a
                # margin that didn't reflect the actual decision.
                is_filtered = bool(
                    (top_k_param is not None and top_k_param > 0)
                    or (top_p_param is not None and top_p_param < 1.0)
                )
                alternatives = []
                cumulative = 0.0
                selected_in_top = False
                for rank, (prob, idx) in enumerate(zip(top_probs.tolist(), top_indices.tolist())):
                    token_text = manager.tokenizer.decode([idx], skip_special_tokens=False)
                    cumulative += prob
                    if idx == next_token_id:
                        selected_in_top = True
                    eff = probs_filtered[idx].item()
                    alternatives.append({
                        "token": token_text,
                        "token_id": idx,
                        "probability": prob,
                        "effective_probability": eff,
                        "is_eligible": eff > 0.0,
                        "raw_probability": raw_probs[idx].item(),  # T=1 probability for comparison
                        "log_probability": math.log(prob) if prob > 0 else float('-inf'),
                        "cumulative_probability": cumulative,
                        "rank": rank + 1
                    })

                # If selected token is not in top-N, add it with its actual probability
                if not selected_in_top:
                    selected_prob = probs[next_token_id].item()
                    selected_eff = probs_filtered[next_token_id].item()
                    selected_raw_prob = raw_probs[next_token_id].item()
                    selected_log_prob = log_probs[next_token_id].item()
                    selected_logit = raw_logits[next_token_id].item()
                    # Find the rank of the selected token
                    sorted_indices = torch.argsort(raw_logits, descending=True)
                    selected_rank = (sorted_indices == next_token_id).nonzero(as_tuple=True)[0].item() + 1
                    alternatives.append({
                        "token": next_token_text,
                        "token_id": next_token_id,
                        "probability": selected_prob,
                        "effective_probability": selected_eff,
                        "is_eligible": selected_eff > 0.0,
                        "raw_probability": selected_raw_prob,  # T=1 probability for comparison
                        "log_probability": selected_log_prob,
                        "cumulative_probability": None,  # Not in sequence
                        "rank": selected_rank,
                        "is_selected_outlier": True  # Flag for UI
                    })
                    # Also add to logits if not present
                    if next_token_id not in [e["token_id"] for e in logits_entries]:
                        cos_outlier, row_norm_outlier = _cos_for(next_token_id)
                        outlier_entry = {
                            "token": next_token_text,
                            "token_id": next_token_id,
                            "logit": selected_logit,
                            "rank": selected_rank,
                            "is_selected_outlier": True,
                        }
                        if cos_outlier is not None:
                            outlier_entry["cosine_sim"] = round(cos_outlier, 4)
                        if row_norm_outlier is not None:
                            outlier_entry["unemb_row_norm"] = round(
                                row_norm_outlier, 4
                            )
                        logits_entries.append(outlier_entry)

                # Build sampling metadata. `is_filtered` and `eligible_count`
                # let the frontend decide whether to show one or two
                # distributions in the panel.
                eligible_count = int((probs_filtered > 0).sum().item())
                sampling_metadata = {
                    "temperature": temperature,
                    "top_k": top_k_param,
                    "top_p": top_p_param,
                    "greedy_token_id": greedy_token_id,
                    "greedy_token": greedy_token,
                    "was_greedy": next_token_id == greedy_token_id,
                    "is_filtered": is_filtered,
                    "eligible_count": eligible_count,
                }

                token_alternatives_by_step.append({
                    "step": step,
                    "selected_token": next_token_text,
                    "selected_token_id": next_token_id,
                    "alternatives": alternatives,
                    "logits": logits_entries,
                    "sampling": sampling_metadata
                })

                # Process attention and hidden states for ALL layers
                layer_data_this_token = []

                for layer_idx in range(len(outputs.attentions)):
                    # Get attention for this layer [batch, num_heads, seq_len, seq_len]
                    layer_attn = outputs.attentions[layer_idx][0]  # Remove batch dim

                    # Get hidden states [batch, seq_len, hidden_dim]
                    current_hidden = outputs.hidden_states[layer_idx + 1]  # +1 because hidden_states includes embedding layer
                    if current_hidden.dim() == 3:
                        current_hidden = current_hidden[0]  # Remove batch dim if present

                    if layer_idx > 0:
                        prev_hidden = outputs.hidden_states[layer_idx]
                        if prev_hidden.dim() == 3:
                            prev_hidden = prev_hidden[0]
                        delta_norm = torch.norm(current_hidden - prev_hidden).item()
                    else:
                        delta_norm = None

                    # Calculate layer metrics
                    import math
                    activation_magnitude = torch.norm(current_hidden).item()
                    # Use a simpler entropy calculation based on attention distribution
                    last_token_hidden = current_hidden[-1]  # [hidden_dim]
                    activation_entropy = torch.std(last_token_hidden).item()  # Use std dev as a proxy for activation diversity
                    hidden_state_norm = torch.norm(last_token_hidden).item()  # Norm of last token

                    # Sanitize to prevent NaN/Inf in JSON
                    activation_magnitude = 0.0 if math.isnan(activation_magnitude) or math.isinf(activation_magnitude) else activation_magnitude
                    activation_entropy = 0.0 if math.isnan(activation_entropy) or math.isinf(activation_entropy) else activation_entropy
                    hidden_state_norm = 0.0 if math.isnan(hidden_state_norm) or math.isinf(hidden_state_norm) else hidden_state_norm
                    if delta_norm is not None:
                        delta_norm = 0.0 if math.isnan(delta_norm) or math.isinf(delta_norm) else delta_norm

                    # Identify critical heads (high max weight or low entropy)
                    critical_heads = []
                    for head_idx in range(layer_attn.shape[0]):
                        head_weights = layer_attn[head_idx, -1, :]  # Attention from last position
                        max_weight = head_weights.max().item()
                        entropy = -(head_weights * torch.log(head_weights + 1e-10)).sum().item()

                        # Normalized attention entropy averaged over latter half of query positions
                        # Normalized by log(k_i) where k_i = number of keys position i can attend to
                        # This produces values in [0,1] with better spread across heads
                        # layer_attn[head_idx] shape: [q_len, k_len]
                        head_attn = layer_attn[head_idx]  # [q_len, k_len]
                        q_len = head_attn.shape[0]

                        # Compute raw entropy per query position
                        token_entropies = -(head_attn * torch.log(head_attn + 1e-10)).sum(dim=-1)  # [q_len]

                        # Normalize by max possible entropy: log(k_i) where k_i = i + 1 (causal mask)
                        # Skip position 0 where log(1) = 0
                        positions = torch.arange(1, q_len + 1, device=head_attn.device, dtype=head_attn.dtype)
                        max_entropies = torch.log(positions + 1e-10)  # log(k_i), with epsilon for position 0
                        normalized_entropies = token_entropies / (max_entropies + 1e-10)  # [0, 1] range

                        # Average over latter half of positions (where there's enough context)
                        start_idx = q_len // 2
                        avg_entropy = normalized_entropies[start_idx:].mean().item() if start_idx < q_len else normalized_entropies.mean().item()

                        # Sanitize to prevent NaN/Inf in JSON
                        max_weight = 0.0 if math.isnan(max_weight) or math.isinf(max_weight) else max_weight
                        entropy = 0.0 if math.isnan(entropy) or math.isinf(entropy) else entropy
                        avg_entropy = 0.0 if math.isnan(avg_entropy) or math.isinf(avg_entropy) else avg_entropy

                        # Score-all-then-rank head classification
                        # Two dimensions: behaviour type (attention geometry) + code cue (token relevance)
                        seq_len_hw = head_weights.shape[0]

                        # --- Behaviour type scores (attention geometry) ---
                        behaviour_scores = {}

                        # Attention sink: weight on positions 0-2
                        sink_w = head_weights[:min(3, seq_len_hw)].sum().item()
                        behaviour_scores["attention_sink"] = sink_w

                        # Previous token: weight on immediate predecessor
                        prev_tok_w = head_weights[-2].item() if seq_len_hw >= 2 else 0.0
                        behaviour_scores["previous_token"] = prev_tok_w

                        # Local: weight within last 5 positions
                        local_w = head_weights[max(0, seq_len_hw - 5):].sum().item() if seq_len_hw > 5 else 0.0
                        behaviour_scores["local"] = local_w

                        # Induction: weight on positions following previous occurrences of current token
                        ind_w = 0.0
                        if step > 0 and seq_len_hw > 1:
                            current_tok = current_ids[0, -1]
                            prev_occ = (current_ids[0, :-1] == current_tok).nonzero(as_tuple=True)[0]
                            if len(prev_occ) > 0:
                                foll = prev_occ + 1
                                foll = foll[foll < seq_len_hw]
                                if len(foll) > 0:
                                    ind_w = head_weights[foll].sum().item()
                        behaviour_scores["induction"] = min(1.0, ind_w)

                        # Focused: low entropy, concentrated attention (not captured by above)
                        focused_score = max(0.0, 1.0 - entropy) if entropy < 1.5 else 0.0
                        behaviour_scores["focused"] = focused_score

                        # Diffuse: high entropy, broad attention
                        diffuse_score = min(1.0, max(0.0, (entropy - 1.0) / 2.0))
                        behaviour_scores["diffuse"] = diffuse_score

                        # Pick primary behaviour (highest score, with minimum thresholds)
                        behaviour_thresholds = {
                            "attention_sink": 0.4,
                            "previous_token": 0.7,
                            "local": 0.5,
                            "induction": 0.2,
                            "focused": 0.3,
                            "diffuse": 0.3,
                        }
                        qualified_behaviours = {
                            k: v for k, v in behaviour_scores.items()
                            if v >= behaviour_thresholds.get(k, 0.3)
                        }
                        sorted_behaviours = sorted(qualified_behaviours.items(), key=lambda x: x[1], reverse=True)
                        primary_behaviour = sorted_behaviours[0] if sorted_behaviours else ("diffuse", diffuse_score)
                        secondary_behaviour = sorted_behaviours[1] if len(sorted_behaviours) > 1 else None

                        pattern_type = primary_behaviour[0]
                        confidence = primary_behaviour[1]

                        # --- Code cue scores (what code tokens are attended to) ---
                        # Decode token texts for code-aware detection (cached per step)
                        if step_token_texts_cache.get('step') != step:
                            try:
                                step_token_texts_cache['texts'] = [
                                    manager.tokenizer.decode([tid]) for tid in current_ids[0, :seq_len_hw].tolist()
                                ]
                            except Exception:
                                step_token_texts_cache['texts'] = []
                            step_token_texts_cache['step'] = step
                        token_texts = step_token_texts_cache.get('texts', [])

                        code_cues = {}
                        if len(token_texts) == seq_len_hw:
                            # Delimiter-sensitive: attention to brackets, braces, parens
                            delimiters = {'(', ')', '{', '}', '[', ']', ':', ';', ','}
                            delim_indices = [i for i, t in enumerate(token_texts) if t.strip() in delimiters]
                            if delim_indices:
                                delim_w = head_weights[delim_indices].sum().item()
                                code_cues["delimiter_sensitive"] = delim_w

                            # Keyword-sensitive: attention to language keywords
                            keywords = {'def', 'return', 'if', 'else', 'elif', 'for', 'while', 'class',
                                       'import', 'from', 'try', 'except', 'with', 'as', 'in', 'not',
                                       'and', 'or', 'True', 'False', 'None', 'self', 'yield', 'async',
                                       'await', 'lambda', 'raise', 'pass', 'break', 'continue',
                                       'function', 'const', 'let', 'var', 'new', 'this'}
                            kw_indices = [i for i, t in enumerate(token_texts) if t.strip() in keywords]
                            if kw_indices:
                                kw_w = head_weights[kw_indices].sum().item()
                                code_cues["keyword_sensitive"] = kw_w

                            # Pattern reuse: attention to a contiguous span that appeared earlier
                            # (broader than induction — checks for repeated multi-token sequences)
                            if ind_w > 0.15:
                                code_cues["pattern_reuse"] = min(1.0, ind_w * 1.5)

                        # Filter code cues by minimum threshold
                        code_cue_threshold = 0.15
                        qualified_cues = {
                            k: round(v, 4) for k, v in code_cues.items()
                            if v >= code_cue_threshold
                        }
                        sorted_cues = sorted(qualified_cues.items(), key=lambda x: x[1], reverse=True)
                        primary_cue = sorted_cues[0] if sorted_cues else None

                        # Sanitize confidence
                        confidence = 0.0 if math.isnan(confidence) or math.isinf(confidence) else confidence

                        # Get full attention weights for this head [seq_len, seq_len]
                        # Store as numpy arrays (not Python lists) to save memory
                        # ~7x more memory efficient: 4 bytes/float vs 28 bytes/float
                        attention_matrix = layer_attn[head_idx].cpu().float().numpy()

                        # Get Q/K/V for this head if available
                        q_matrix = None
                        k_matrix = None
                        v_matrix = None
                        if layer_idx in qkv_captures:
                            # Q/K/V shape: [seq_len, n_heads, head_dim]
                            # Store as numpy arrays for memory efficiency
                            q_matrix = qkv_captures[layer_idx]['q'][:, head_idx, :].float().numpy()
                            k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy()
                            v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy()

                        # Store matrices in cache for lazy loading (reduces response size)
                        matrix_cache.store(request_id, step, layer_idx, head_idx, {
                            "attention_weights": attention_matrix,
                            "q_matrix": q_matrix,
                            "k_matrix": k_matrix,
                            "v_matrix": v_matrix
                        })

                        # Return only metadata (matrices fetched on-demand via /matrix endpoint)
                        head_entry = {
                            "head_idx": head_idx,
                            "entropy": entropy,
                            "avg_entropy": avg_entropy,  # Averaged over all query positions
                            "max_weight": max_weight,
                            "has_matrices": attention_matrix is not None,  # Flag for frontend
                            "pattern": {
                                "type": pattern_type,
                                "confidence": round(confidence, 4),
                            } if pattern_type else None,
                        }
                        # Secondary behaviour (if present and distinct from primary)
                        if secondary_behaviour:
                            head_entry["secondary_behaviour"] = {
                                "type": secondary_behaviour[0],
                                "score": round(secondary_behaviour[1], 4),
                            }
                        # Code cue (separate dimension from behaviour type)
                        if primary_cue:
                            head_entry["code_cue"] = {
                                "type": primary_cue[0],
                                "score": round(primary_cue[1], 4),
                                "evidence": f"{round(primary_cue[1] * 100)}% attention on {primary_cue[0].replace('_', ' ')} tokens",
                            }
                            # Secondary code cue
                            if len(sorted_cues) > 1:
                                head_entry["secondary_cue"] = {
                                    "type": sorted_cues[1][0],
                                    "score": round(sorted_cues[1][1], 4),
                                }
                        critical_heads.append(head_entry)

                    # Sort by max_weight (return all heads, frontend will decide how many to display)
                    critical_heads.sort(key=lambda h: h["max_weight"], reverse=True)

                    # Layer-level pattern: majority vote of head patterns, weighted by confidence
                    pattern_votes = {}
                    for h in critical_heads:
                        if h["pattern"] and h["pattern"]["type"]:
                            pt = h["pattern"]["type"]
                            pc = h["pattern"]["confidence"]
                            pattern_votes[pt] = pattern_votes.get(pt, 0.0) + pc
                    layer_pattern = None
                    if pattern_votes:
                        best_type = max(pattern_votes, key=pattern_votes.get)
                        total_conf = sum(pattern_votes.values())
                        layer_pattern = {
                            "type": best_type,
                            "confidence": round(pattern_votes[best_type] / total_conf, 3) if total_conf > 0 else 0.0
                        }

                    layer_data_this_token.append({
                        "layer_idx": layer_idx,
                        "pattern": layer_pattern,
                        "critical_heads": critical_heads,
                        "activation_magnitude": activation_magnitude,
                        "activation_entropy": activation_entropy,
                        "hidden_state_norm": hidden_state_norm,
                        "delta_norm": delta_norm
                    })

                layer_data_by_token.append(layer_data_this_token)

                # Update inputs
                next_token_tensor = torch.tensor([[next_token_id]], dtype=torch.long, device=manager.device)
                current_ids = torch.cat([current_ids, next_token_tensor], dim=1)

                # Stop on EOS
                if next_token_id == manager.tokenizer.eos_token_id:
                    break

                # Free memory from this step's outputs to prevent accumulation
                # This is critical for large models like Devstral (40 layers, 32 heads)
                del outputs
                del logits
                del probs
                if 'layer_attn' in dir():
                    del layer_attn
                if 'current_hidden' in dir():
                    del current_hidden

                # Periodic garbage collection for large models (every 8 steps)
                if (step + 1) % 8 == 0:
                    gc.collect()
                    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                        torch.mps.empty_cache() if hasattr(torch.mps, 'empty_cache') else None

        # Clean up hooks after generation
        for hook in hooks:
            hook.remove()

        # Placeholder for Q/K/V data (will be populated in future iterations)
        qkv_by_layer_head = {}

        generation_time = time.time() - start_time

        # Calculate token section boundaries for UI display
        total_tokens = prompt_length + len(generated_token_ids)
        system_prompt_text = system_prompt_override or (model_config.get("system_prompt") if model_config else None)

        # For instruction models, estimate where system prompt ends
        # This is approximate due to control tokens in chat templates
        system_prompt_end = 0
        if prompt_style == "instruction" and system_prompt_text:
            if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
                # For Devstral, try encoding with empty system to estimate boundary
                try:
                    no_system_tokens = manager.mistral_tokenizer.encode_chat("", prompt)
                    # The difference gives us system tokens, but we need to add 1 to include
                    # the closing [/SYSTEM_PROMPT] tag in the system prompt section
                    system_prompt_end = prompt_length - len(no_system_tokens) + 1
                    # Ensure non-negative and within bounds
                    system_prompt_end = max(0, min(system_prompt_end, prompt_length))
                    logger.info(f"Estimated system prompt boundary: {system_prompt_end} tokens (includes closing tag)")
                except Exception as e:
                    logger.warning(f"Could not estimate system prompt boundary: {e}")
                    system_prompt_end = 0
            else:
                # For other instruction models, rough estimate based on character ratio
                # This is very approximate but provides some visual separation
                total_chars = len(system_prompt_text or "") + len(prompt)
                if total_chars > 0:
                    system_ratio = len(system_prompt_text or "") / total_chars
                    system_prompt_end = int(prompt_length * system_ratio)

        token_sections = {
            "systemPrompt": {
                "start": 0,
                "end": system_prompt_end,
                "text": system_prompt_text,
                "tokenCount": system_prompt_end
            },
            "userPrompt": {
                "start": system_prompt_end,
                "end": prompt_length,
                "text": prompt,
                "tokenCount": prompt_length - system_prompt_end
            },
            "output": {
                "start": prompt_length,
                "end": total_tokens,
                "text": "".join(generated_tokens),
                "tokenCount": len(generated_token_ids)
            }
        }

        # Build token metadata for frontend (eliminates per-token API calls)
        from .tokenizer_utils import TokenizerMetadata
        token_metadata = TokenizerMetadata(manager.tokenizer)

        # Include every id the tokenizer considers a special / added token
        # (BOS, EOS, PAD, UNK, chat-template delimiters like [INST]/[/INST],
        # system-prompt markers, tool-call markers, etc.). The HF tokenizer's
        # all_special_ids misses Mistral-specific chat-template delimiters, so
        # we also union in the control-token ids from mistral-common when the
        # Mistral path is active.
        special_token_ids = set(getattr(manager.tokenizer, "all_special_ids", []) or [])
        for tok_id in (
            manager.tokenizer.eos_token_id,
            manager.tokenizer.bos_token_id,
            manager.tokenizer.pad_token_id,
            manager.tokenizer.unk_token_id,
        ):
            if tok_id is not None:
                special_token_ids.add(tok_id)
        if manager.mistral_tokenizer is not None:
            special_token_ids |= manager.mistral_tokenizer.get_control_token_ids()

        def build_token_data(token_ids, token_texts, token_type):
            """Build token data with full metadata for hover tooltips"""
            multi_split_flags = token_metadata.is_multi_split_identifier(token_ids)
            result = []
            for i, (tid, t) in enumerate(zip(token_ids, token_texts)):
                bpe_pieces = token_metadata.get_subword_pieces(tid)
                result.append({
                    "text": t,
                    "idx": tid,
                    "bytes": len(t.encode('utf-8')),
                    "type": token_type,
                    "bpe_pieces": bpe_pieces,
                    "is_special": tid in special_token_ids,
                    "is_multi_split": multi_split_flags[i] if i < len(multi_split_flags) else False,
                    "num_pieces": len(bpe_pieces),
                })
            return result

        # Embedding L2 norms — see streaming endpoint for the full
        # rationale (RQ1 layer-0 anchor for the residual-norm trace).
        # Covers both prompt tokens and generated tokens so the
        # Mechanism lens can render bars for the full context window
        # (prompt + previously-generated tokens) at any step.
        embedding_norms: List[float] = _compute_embedding_norms(
            manager,
            list(prompt_token_ids) + list(generated_token_ids),
        )

        # Build response
        response = {
            "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
            "prompt": prompt,
            "promptTokens": build_token_data(prompt_token_ids, prompt_tokens, "prompt"),
            "generatedTokens": build_token_data(generated_token_ids, generated_tokens, "generated"),
            "tokenSections": token_sections,  # Section boundaries for UI coloring
            "tokenAlternatives": token_alternatives_by_step,  # Top-k alternatives for each token
            "layersDataByStep": layer_data_by_token,  # Layer data for ALL generation steps
            "layersData": layer_data_by_token[-1] if layer_data_by_token else [],  # Keep for backward compatibility
            "qkvData": {},  # Deprecated: matrices now lazy-loaded via /matrix endpoint
            "modelInfo": {
                "numLayers": n_layers,
                "numHeads": n_heads,
                "modelDimension": d_model,
                "headDim": head_dim,
                "vocabSize": manager.model.config.vocab_size
            },
            "embeddingNorms": embedding_norms,
            "generationTime": generation_time,
            "numTokensGenerated": len(generated_tokens)
        }

        logger.info(f"✅ Research attention analysis complete: {len(generated_tokens)} tokens, {generation_time:.2f}s")

        # Sanitize response to handle NaN/Inf values that break JSON serialization
        return sanitize_for_json(response)

    except Exception as e:
        logger.error(f"Research attention analysis error: {e}")
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))


def sse_event(event_type: str, **kwargs) -> str:
    """Format data as SSE event"""
    data = {'type': event_type, 'timestamp': int(time.time() * 1000), **kwargs}
    return f"data: {json.dumps(data)}\n\n"


@app.post("/analyze/research/attention/stream")
async def analyze_research_attention_stream(request: Dict[str, Any], authenticated: bool = Depends(verify_api_key)):
    """
    SSE Streaming version of Research-Grade Attention Analysis

    Emits progress events during each stage:
    - tokenizing: Initial tokenization
    - generating: Per-token generation progress
    - extracting: Per-layer attention extraction
    - serializing: Building response
    - complete: Analysis finished
    - result: Final data payload
    """
    async def event_generator():
        try:
            import time
            start_time = time.time()

            # Generate unique request ID for matrix cache lookup
            request_id = str(uuid.uuid4())

            # `ephemeral` callers (e.g. the What-if "Rewrite & rerun"
            # counterfactual) want their regen to coexist with the
            # original traced run rather than displace it — the user
            # still needs the original's matrices for the Mechanism /
            # Attention lenses. The new request_id is unique, so the
            # ephemeral run's own writes don't collide; we just have to
            # skip the eager `clear_old_requests` that would otherwise
            # purge every other request_id's cache.
            ephemeral = bool(request.get("ephemeral", False))

            # Clear old cached matrices to free memory before starting new analysis
            if not ephemeral:
                matrix_cache.clear_old_requests(request_id)

            # Record per-request model metadata so aggregate-row lookups can
            # derive num_heads from this request rather than from the live
            # `manager.model` (which may have been switched in the meantime).
            # Ephemeral runs (What-if Rewrite & rerun) skip every cache
            # write — the frontend discards everything except the
            # generated tokens, so the regen leaves no trace at all.
            if manager.adapter is not None and not ephemeral:
                matrix_cache.set_request_metadata(
                    request_id,
                    num_heads=manager.adapter.get_num_heads(),
                    num_layers=manager.adapter.get_num_layers(),
                    model_id=manager.model_id,
                )

            # Get parameters
            prompt = request.get("prompt", "def quicksort(arr):")
            max_tokens = request.get("max_tokens", 8)
            auto_complete = request.get("auto_complete", False)
            temperature = request.get("temperature", 0.7)
            top_k_param = request.get("top_k", None)  # Top-k sampling parameter
            top_p_param = request.get("top_p", None)  # Top-p (nucleus) sampling parameter

            # If auto_complete mode, ensure we have a reasonable upper limit
            if auto_complete:
                max_tokens = min(max_tokens, 128)

            logger.info(f"[SSE] Research attention analysis: prompt_len={len(prompt)}, max_tokens={max_tokens}, auto_complete={auto_complete}, request_id={request_id}")

            # === STAGE 1: TOKENIZING ===
            yield sse_event('tokenizing', stage=1, totalStages=5, progress=2,
                           stageProgress=0, detail=f'Tokenizing {len(prompt)} characters...')

            # Get model config for prompt formatting
            from .model_config import get_model_config
            from .prompt_formatter import format_prompt
            model_config = get_model_config(manager.model_id)

            # Get optional system prompt override from request
            system_prompt_override = request.get("system_prompt")

            # Format prompt using the unified formatter
            formatted_prompt = format_prompt(
                prompt=prompt,
                model_config=model_config or {},
                tokenizer=manager.tokenizer,
                system_prompt_override=system_prompt_override
            )

            prompt_style = model_config.get("prompt_style", "completion") if model_config else "completion"

            # Temperature is now controlled by the frontend UI
            # The frontend sets appropriate defaults per model (0.15 for Devstral, 0.7 for CodeGen)
            logger.info(f"[SSE] Using temperature={temperature} from request")

            # Tokenize and prepare - use MistralTokenizer for Devstral
            if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
                system_prompt = system_prompt_override or (model_config.get("system_prompt") if model_config else "")
                prompt_token_ids = manager.mistral_tokenizer.encode_chat(system_prompt, prompt)
                inputs = {"input_ids": torch.tensor([prompt_token_ids]).to(manager.device)}
                prompt_length = len(prompt_token_ids)
                # mistral-common's decode_token() returns "" for control tokens
                # (<s>, [INST], [/INST], [SYSTEM_PROMPT], etc.). Fall back to
                # the HF tokenizer so every token has a printable string form.
                def _decode_prompt_token(tid: int) -> str:
                    text = manager.mistral_tokenizer.decode_token(tid)
                    if text:
                        return text
                    try:
                        return manager.tokenizer.decode([tid], skip_special_tokens=False) or ""
                    except Exception:
                        return ""
                prompt_tokens = [_decode_prompt_token(tid) for tid in prompt_token_ids]
            else:
                inputs = manager.tokenizer(formatted_prompt, return_tensors="pt").to(manager.device)
                prompt_length = inputs["input_ids"].shape[1]
                prompt_token_ids = inputs["input_ids"][0].tolist()
                prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]

            yield sse_event('tokenizing', stage=1, totalStages=5, progress=8,
                           stageProgress=100, detail=f'Tokenized into {prompt_length} tokens',
                           metadata={'tokenCount': prompt_length})
            await asyncio.sleep(0)  # Yield to event loop

            # Storage for generation
            generated_token_ids = []
            generated_tokens = []

            # Model info
            n_layers = len(list(manager.model.parameters()))
            if hasattr(manager.model.config, 'n_layer'):
                n_layers = manager.model.config.n_layer
            elif hasattr(manager.model.config, 'num_hidden_layers'):
                n_layers = manager.model.config.num_hidden_layers

            n_heads = manager.model.config.n_head if hasattr(manager.model.config, 'n_head') else manager.model.config.num_attention_heads
            d_model = manager.model.config.n_embd if hasattr(manager.model.config, 'n_embd') else manager.model.config.hidden_size
            head_dim = d_model // n_heads

            # === STAGE 2: GENERATING ===
            layer_data_by_token = []
            token_alternatives_by_step = []

            # Hook system to capture Q/K/V matrices
            qkv_captures = {}
            hooks = []

            # Hook for combined QKV projection (CodeGen style)
            def make_combined_qkv_hook(layer_idx):
                def hook(module, input, output):
                    try:
                        if output.dim() != 3:
                            return
                        batch_size, seq_len, hidden = output.shape
                        expected_hidden = 3 * n_heads * head_dim
                        if hidden != expected_hidden:
                            return
                        qkv = output.reshape(batch_size, seq_len, 3, n_heads, head_dim)
                        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
                        qkv_captures[layer_idx] = {
                            'q': q[0].detach().cpu(),
                            'k': k[0].detach().cpu(),
                            'v': v[0].detach().cpu()
                        }
                    except Exception:
                        pass
                return hook

            # Hooks for separate Q, K, V projections (Mistral/LLaMA style)
            def make_separate_proj_hook(layer_idx, proj_type, num_kv_heads=None):
                """Create hook for separate Q/K/V projection modules.

                For GQA models, K and V have fewer heads than Q, so we need to
                expand them to match Q's head count for consistent visualization.
                """
                def hook(module, input, output):
                    try:
                        if output.dim() != 3:
                            return
                        batch_size, seq_len, hidden = output.shape

                        # Determine number of heads for this projection
                        if proj_type == 'q':
                            proj_heads = n_heads
                        else:
                            # K and V may have fewer heads (GQA)
                            proj_heads = num_kv_heads if num_kv_heads else n_heads

                        proj_head_dim = hidden // proj_heads
                        if hidden != proj_heads * proj_head_dim:
                            return

                        # Reshape to [batch, seq, heads, head_dim]
                        proj_output = output.reshape(batch_size, seq_len, proj_heads, proj_head_dim)

                        # For GQA, expand K/V to match Q's head count
                        if proj_type != 'q' and num_kv_heads and num_kv_heads < n_heads:
                            # Repeat each KV head to match Q heads
                            repeat_factor = n_heads // num_kv_heads
                            proj_output = proj_output.repeat_interleave(repeat_factor, dim=2)

                        # Initialize layer entry if needed
                        if layer_idx not in qkv_captures:
                            qkv_captures[layer_idx] = {}

                        qkv_captures[layer_idx][proj_type] = proj_output[0].detach().cpu()
                    except Exception as e:
                        logger.debug(f"QKV capture error for layer {layer_idx} {proj_type}: {e}")
                return hook

            # Register hooks based on model architecture
            try:
                # CodeGen style: model.transformer.h[layer].attn.qkv_proj
                if hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'h'):
                    for layer_idx, layer in enumerate(manager.model.transformer.h):
                        if hasattr(layer, 'attn') and hasattr(layer.attn, 'qkv_proj'):
                            hook = layer.attn.qkv_proj.register_forward_hook(make_combined_qkv_hook(layer_idx))
                            hooks.append(hook)
                        elif hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
                            hook = layer.attn.c_attn.register_forward_hook(make_combined_qkv_hook(layer_idx))
                            hooks.append(hook)

                # Mistral/LLaMA style: model.model.layers[layer].self_attn.{q,k,v}_proj
                elif hasattr(manager.model, 'model') and hasattr(manager.model.model, 'layers'):
                    num_kv_heads = getattr(manager.model.config, 'num_key_value_heads', None)
                    for layer_idx, layer in enumerate(manager.model.model.layers):
                        if hasattr(layer, 'self_attn'):
                            attn = layer.self_attn
                            if hasattr(attn, 'q_proj'):
                                hook = attn.q_proj.register_forward_hook(
                                    make_separate_proj_hook(layer_idx, 'q', num_kv_heads))
                                hooks.append(hook)
                            if hasattr(attn, 'k_proj'):
                                hook = attn.k_proj.register_forward_hook(
                                    make_separate_proj_hook(layer_idx, 'k', num_kv_heads))
                                hooks.append(hook)
                            if hasattr(attn, 'v_proj'):
                                hook = attn.v_proj.register_forward_hook(
                                    make_separate_proj_hook(layer_idx, 'v', num_kv_heads))
                                hooks.append(hook)
                    logger.info(f"Registered QKV hooks for {len(hooks)//3} Mistral layers (GQA: {num_kv_heads} KV heads)")
            except Exception as hook_error:
                logger.warning(f"Could not register QKV hooks: {hook_error}")

            # Phase 4: Hooks for attention and MLP output norms + gate activation stats
            attn_output_norms = {}
            mlp_output_norms = {}
            gate_activation_stats = {}
            # Per-layer top-K SwiGLU intermediate activations at the
            # predicting position (last token of the current forward
            # pass). Each entry is a list of {idx, value} for the
            # top-K neurons by |silu(W_g·x) ⊙ (W_u·x)| at that layer
            # for that step. Powers RQ1 Step 5's "Behind the SwiGLU"
            # per-neuron view without storing the full 14,336-d
            # vector — bounded memory, full transparency for the
            # most-loaded neurons.
            ffn_top_neurons = {}
            FFN_TOPK = 32

            def make_attn_output_hook(layer_idx):
                def hook(module, input, output):
                    try:
                        out = output[0] if isinstance(output, tuple) else output
                        if out.dim() == 3:
                            attn_output_norms[layer_idx] = torch.norm(out[0, -1]).item()
                    except Exception:
                        pass
                return hook

            def make_mlp_output_hook(layer_idx):
                def hook(module, input, output):
                    try:
                        out = output[0] if isinstance(output, tuple) else output
                        if out.dim() == 3:
                            mlp_output_norms[layer_idx] = torch.norm(out[0, -1]).item()
                        elif out.dim() == 2:
                            mlp_output_norms[layer_idx] = torch.norm(out[-1]).item()
                    except Exception:
                        pass
                return hook

            def make_gate_hook(layer_idx):
                """Capture gate activation stats for SwiGLU FFN (LLaMA/Mistral)."""
                def hook(module, input, output):
                    try:
                        inp = input[0] if isinstance(input, tuple) else input
                        if inp.dim() == 3:
                            inp = inp[0, -1]  # last token
                        elif inp.dim() == 2:
                            inp = inp[-1]
                        if hasattr(module, 'gate_proj'):
                            gate_out = torch.nn.functional.silu(module.gate_proj(inp))
                            abs_gate = gate_out.abs()
                            gate_activation_stats[layer_idx] = {
                                "sparsity": round(float((abs_gate < 0.01).float().mean().item()), 4),
                                "mean": round(float(gate_out.mean().item()), 4),
                                "max": round(float(gate_out.max().item()), 4),
                            }
                    except Exception:
                        pass
                return hook

            def make_ffn_top_neurons_hook(layer_idx):
                """
                Capture the top-K of the SwiGLU intermediate activation
                vector at the predicting (last) position for this layer.

                Computes silu(W_gate · x) ⊙ (W_up · x) — the same
                14,336-d intermediate that gets fed to W_down. Storing
                only the top-K (=32) by |value| keeps memory bounded:
                32 × {idx:int, value:float} ≈ 256 bytes per layer per
                step, vs 14,336 × 4 bytes = 57KB per layer per step
                if we stored the full vector. The top-K is what a
                developer reads anyway (the rest are noise) and the
                stored neuron index lets the frontend label "neuron
                8421 dominated this token at L37".

                Only fires for SwiGLU layers (gate_proj + up_proj
                present); other architectures are skipped at
                registration time.
                """
                def hook(module, input, output):
                    try:
                        inp = input[0] if isinstance(input, tuple) else input
                        if inp.dim() == 3:
                            inp = inp[0, -1]  # last (predicting) token
                        elif inp.dim() == 2:
                            inp = inp[-1]
                        if hasattr(module, 'gate_proj') and hasattr(module, 'up_proj'):
                            gate_out = torch.nn.functional.silu(
                                module.gate_proj(inp)
                            )
                            up_out = module.up_proj(inp)
                            intermediate = gate_out * up_out
                            abs_inter = intermediate.abs()
                            k = min(FFN_TOPK, intermediate.shape[-1])
                            top_vals, top_idx = torch.topk(abs_inter, k=k)
                            # Use the signed values (not abs) so the
                            # frontend can render direction; abs only
                            # drove the ranking.
                            signed = intermediate[top_idx]
                            ffn_top_neurons[layer_idx] = {
                                "k": k,
                                "intermediate_size": int(intermediate.shape[-1]),
                                "neurons": [
                                    {
                                        "idx": int(i),
                                        "value": round(float(v), 4),
                                    }
                                    for i, v in zip(
                                        top_idx.cpu().tolist(),
                                        signed.cpu().tolist(),
                                    )
                                ],
                            }
                    except Exception:
                        pass
                return hook

            # Cache for decoded token texts (reused across heads within a step)
            step_token_texts_cache: Dict[str, Any] = {}

            # Detect FFN type from first layer
            ffn_type = "gelu"  # default

            try:
                # CodeGen style
                if hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'h'):
                    for layer_idx, layer in enumerate(manager.model.transformer.h):
                        if hasattr(layer, 'attn'):
                            hook = layer.attn.register_forward_hook(make_attn_output_hook(layer_idx))
                            hooks.append(hook)
                        if hasattr(layer, 'mlp'):
                            hook = layer.mlp.register_forward_hook(make_mlp_output_hook(layer_idx))
                            hooks.append(hook)
                # Mistral/LLaMA style
                elif hasattr(manager.model, 'model') and hasattr(manager.model.model, 'layers'):
                    for layer_idx, layer in enumerate(manager.model.model.layers):
                        if hasattr(layer, 'self_attn'):
                            hook = layer.self_attn.register_forward_hook(make_attn_output_hook(layer_idx))
                            hooks.append(hook)
                        if hasattr(layer, 'mlp'):
                            hook = layer.mlp.register_forward_hook(make_mlp_output_hook(layer_idx))
                            hooks.append(hook)
                            # Gate hook for SwiGLU models
                            if hasattr(layer.mlp, 'gate_proj'):
                                hook = layer.mlp.register_forward_hook(make_gate_hook(layer_idx))
                                hooks.append(hook)
                                if layer_idx == 0:
                                    ffn_type = "swiglu"
                            # FFN top-K intermediate hook — needs both
                            # gate_proj and up_proj to compute the
                            # SwiGLU intermediate vector.
                            if hasattr(layer.mlp, 'gate_proj') and hasattr(layer.mlp, 'up_proj'):
                                hook = layer.mlp.register_forward_hook(
                                    make_ffn_top_neurons_hook(layer_idx)
                                )
                                hooks.append(hook)
                logger.info(f"Registered attn/MLP output hooks for contribution tracking (ffn_type={ffn_type})")
            except Exception as hook_error:
                logger.warning(f"Could not register attn/MLP hooks: {hook_error}")

            with torch.no_grad():
                current_ids = inputs["input_ids"]

                for step in range(max_tokens):
                    # Emit progress for this generation step
                    step_progress = (step / max_tokens) * 100
                    overall_progress = 10 + (step / max_tokens) * 20  # 10-30%
                    yield sse_event('generating', stage=2, totalStages=5, progress=overall_progress,
                                   stageProgress=step_progress,
                                   detail=f'Generating token {step + 1}/{max_tokens}',
                                   metadata={'stepIndex': step, 'totalSteps': max_tokens})
                    await asyncio.sleep(0)

                    qkv_captures.clear()
                    attn_output_norms.clear()
                    mlp_output_norms.clear()
                    gate_activation_stats.clear()
                    ffn_top_neurons.clear()

                    # Forward pass with full outputs
                    outputs = manager.model(
                        current_ids,
                        output_attentions=True,
                        output_hidden_states=True
                    )

                    # Record the query position whose logits will select this
                    # step's token (see non-SSE site for rationale).
                    query_position = current_ids.shape[1] - 1
                    query_token_id = current_ids[0, -1].item()
                    if not ephemeral:
                        matrix_cache.set_step_query_info(
                            request_id, step, query_position, query_token_id
                        )

                    # Get logits for next token
                    raw_logits = outputs.logits[0, -1, :].clone()  # Clone raw logits before any scaling

                    # Capture raw logits for top-10 tokens (before temperature scaling)
                    import math as math_module
                    top_n_display = 10  # Get top 10 alternatives for display
                    top_raw_logits, top_raw_indices = torch.topk(raw_logits, k=min(top_n_display, len(raw_logits)))

                    # Build raw logits entries (before temperature)
                    # Use correct tokenizer for Devstral vs other models.
                    # On the Mistral path, fall back to the HF tokenizer when
                    # mistral-common returns "" so logit candidates (e.g. EOS,
                    # chat-template ids) always carry a printable string.
                    def decode_token(tid):
                        if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
                            text = manager.mistral_tokenizer.decode_token(tid)
                            if text:
                                return text
                            try:
                                return manager.tokenizer.decode([tid], skip_special_tokens=False) or ""
                            except Exception:
                                return ""
                        else:
                            return manager.tokenizer.decode([tid], skip_special_tokens=False)

                    # ── Unembedding-cosine prep (SSE path) ──────────
                    # Mirrors the non-SSE path. Each logit is
                    # W_E[v] · h_post_norm; cos(W_E[v], h_post_norm)
                    # strips magnitude so a developer can read whether
                    # a high-logit token won by direction or by
                    # row magnitude. Cheap (one matmul per token per
                    # step) and computed once per step.
                    cos_h_norm = None
                    W_E = None
                    try:
                        final_hidden_pre_norm = outputs.hidden_states[-1][0, -1, :]
                        if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
                            h_post_norm = manager.model.model.norm(
                                final_hidden_pre_norm.unsqueeze(0)
                            )[0]
                        elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
                            h_post_norm = manager.model.transformer.ln_f(
                                final_hidden_pre_norm.unsqueeze(0)
                            )[0]
                        else:
                            h_post_norm = final_hidden_pre_norm
                        h_norm_val = float(torch.linalg.vector_norm(h_post_norm).item())
                        if h_norm_val > 1e-9:
                            cos_h_norm = h_post_norm / h_norm_val
                            if hasattr(manager.model, 'lm_head'):
                                W_E = manager.model.lm_head.weight
                            else:
                                logger.warning(
                                    "[SSE cosine prep] manager.model has no lm_head — cosine emission skipped"
                                )
                        else:
                            logger.warning(
                                f"[SSE cosine prep] h_post_norm magnitude too small: {h_norm_val}"
                            )
                        if cos_h_norm is not None and W_E is not None and step == 0:
                            logger.info(
                                f"[SSE cosine prep] enabled · h_post_norm shape {tuple(h_post_norm.shape)} · "
                                f"W_E shape {tuple(W_E.shape)} · h_norm {h_norm_val:.3f}"
                            )
                    except Exception as e:
                        logger.warning(
                            f"[SSE cosine prep] failed at step {step}: {type(e).__name__}: {e}"
                        )
                        cos_h_norm = None
                        W_E = None

                    def _cos_for(token_idx: int):
                        """cosine(W_E[token_idx], h_post_norm) and ‖W_E row‖."""
                        if cos_h_norm is None or W_E is None:
                            return None, None
                        try:
                            row = W_E[token_idx]
                            row_norm = float(torch.linalg.vector_norm(row).item())
                            if row_norm < 1e-9:
                                return None, row_norm
                            cos = float(((row / row_norm) @ cos_h_norm).item())
                            return cos, row_norm
                        except Exception:
                            return None, None

                    logits_entries = []
                    for rank, (logit_val, idx) in enumerate(zip(top_raw_logits.tolist(), top_raw_indices.tolist())):
                        token_text = decode_token(idx)
                        cos, row_norm = _cos_for(idx)
                        entry = {
                            "token": token_text,
                            "token_id": idx,
                            "logit": logit_val,
                            "rank": rank + 1,
                        }
                        if cos is not None:
                            entry["cosine_sim"] = round(cos, 4)
                        if row_norm is not None:
                            entry["unemb_row_norm"] = round(row_norm, 4)
                        logits_entries.append(entry)

                    # Greedy token (argmax of raw logits, before any sampling)
                    greedy_token_id = torch.argmax(raw_logits).item()
                    greedy_token = decode_token(greedy_token_id)

                    # Compute raw probabilities (T=1) for comparison visualization
                    raw_probs = torch.softmax(raw_logits, dim=0)

                    # Apply temperature scaling
                    logits = raw_logits.clone()
                    if temperature > 0:
                        logits = logits / temperature
                    probs = torch.softmax(logits, dim=0)

                    # Apply top-k filtering if specified
                    if top_k_param is not None and top_k_param > 0:
                        top_k_probs, top_k_indices = torch.topk(probs, k=min(top_k_param, len(probs)))
                        probs_filtered = torch.zeros_like(probs)
                        probs_filtered[top_k_indices] = top_k_probs
                        probs_filtered = probs_filtered / probs_filtered.sum()  # Renormalize
                    else:
                        probs_filtered = probs

                    # Apply top-p (nucleus) filtering if specified
                    if top_p_param is not None and top_p_param < 1.0:
                        sorted_probs, sorted_indices = torch.sort(probs_filtered, descending=True)
                        cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                        cutoff_mask = cumulative_probs > top_p_param
                        cutoff_mask[1:] = cutoff_mask[:-1].clone()
                        cutoff_mask[0] = False
                        sorted_probs[cutoff_mask] = 0
                        probs_filtered = torch.zeros_like(probs)
                        probs_filtered.scatter_(0, sorted_indices, sorted_probs)
                        if probs_filtered.sum() > 0:
                            probs_filtered = probs_filtered / probs_filtered.sum()

                    if temperature == 0:
                        next_token_id = torch.argmax(probs_filtered, dim=-1).item()
                    else:
                        if probs_filtered.sum() > 0:
                            next_token_id = torch.multinomial(probs_filtered, 1).item()
                        else:
                            next_token_id = torch.argmax(probs, dim=-1).item()

                    next_token_text = decode_token(next_token_id)

                    generated_token_ids.append(next_token_id)
                    generated_tokens.append(next_token_text)

                    # Capture top-10 token alternatives with probabilities
                    # Use log_softmax for numerical stability at low temperatures
                    _, top_indices = torch.topk(logits, k=min(top_n_display, len(logits)))

                    # Use log_softmax (numerically stable) then exp() for probabilities
                    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                    top_probs = torch.exp(log_probs[top_indices])

                    # See non-SSE site for the rationale behind emitting both
                    # `probability` (pre-filter model preference) and
                    # `effective_probability` (post-filter, what the sampler
                    # drew from). When no filter is active the two coincide.
                    is_filtered = bool(
                        (top_k_param is not None and top_k_param > 0)
                        or (top_p_param is not None and top_p_param < 1.0)
                    )
                    alternatives = []
                    cumulative = 0.0
                    selected_in_top = False
                    for rank, (prob, idx) in enumerate(zip(top_probs.tolist(), top_indices.tolist())):
                        token_text = decode_token(idx)
                        cumulative += prob
                        if idx == next_token_id:
                            selected_in_top = True
                        eff = probs_filtered[idx].item()
                        alternatives.append({
                            "token": token_text,
                            "token_id": idx,
                            "probability": prob,
                            "effective_probability": eff,
                            "is_eligible": eff > 0.0,
                            "raw_probability": raw_probs[idx].item(),  # T=1 probability for comparison
                            "log_probability": math_module.log(prob) if prob > 0 else float('-inf'),
                            "cumulative_probability": cumulative,
                            "rank": rank + 1
                        })

                    # If selected token is not in top-N, add it with its actual probability
                    if not selected_in_top:
                        selected_prob = probs[next_token_id].item()
                        selected_eff = probs_filtered[next_token_id].item()
                        selected_raw_prob = raw_probs[next_token_id].item()
                        selected_log_prob = log_probs[next_token_id].item()
                        selected_logit = raw_logits[next_token_id].item()
                        # Find the rank of the selected token
                        sorted_indices = torch.argsort(raw_logits, descending=True)
                        selected_rank = (sorted_indices == next_token_id).nonzero(as_tuple=True)[0].item() + 1
                        alternatives.append({
                            "token": next_token_text,
                            "token_id": next_token_id,
                            "probability": selected_prob,
                            "effective_probability": selected_eff,
                            "is_eligible": selected_eff > 0.0,
                            "raw_probability": selected_raw_prob,  # T=1 probability for comparison
                            "log_probability": selected_log_prob,
                            "cumulative_probability": None,
                            "rank": selected_rank,
                            "is_selected_outlier": True
                        })
                        # Also add to logits if not present
                        if next_token_id not in [e["token_id"] for e in logits_entries]:
                            cos_outlier, row_norm_outlier = _cos_for(next_token_id)
                            outlier_entry = {
                                "token": next_token_text,
                                "token_id": next_token_id,
                                "logit": selected_logit,
                                "rank": selected_rank,
                                "is_selected_outlier": True,
                            }
                            if cos_outlier is not None:
                                outlier_entry["cosine_sim"] = round(cos_outlier, 4)
                            if row_norm_outlier is not None:
                                outlier_entry["unemb_row_norm"] = round(
                                    row_norm_outlier, 4
                                )
                            logits_entries.append(outlier_entry)

                    # Build sampling metadata
                    eligible_count = int((probs_filtered > 0).sum().item())
                    sampling_metadata = {
                        "temperature": temperature,
                        "top_k": top_k_param,
                        "top_p": top_p_param,
                        "greedy_token_id": greedy_token_id,
                        "greedy_token": greedy_token,
                        "was_greedy": next_token_id == greedy_token_id,
                        "is_filtered": is_filtered,
                        "eligible_count": eligible_count,
                    }

                    # --- Margin computation and stability classification ---
                    import math as _math_margin
                    winner_logit = logits_entries[0]["logit"] if len(logits_entries) > 0 else 0.0
                    runnerup_logit = logits_entries[1]["logit"] if len(logits_entries) > 1 else winner_logit
                    margin = winner_logit - runnerup_logit
                    runnerup_token = logits_entries[1]["token"] if len(logits_entries) > 1 else ""

                    # Entropy over top-k probabilities
                    top_probs_list_for_entropy = [a["probability"] for a in alternatives[:10] if a["probability"] > 0]
                    margin_entropy = -sum(p * _math_margin.log(p) for p in top_probs_list_for_entropy) if top_probs_list_for_entropy else 0.0

                    stability = _classify_stability(margin)

                    # Greedy margin: margin computed from raw logits (temperature=0)
                    raw_sorted_logits, raw_sorted_indices = torch.topk(raw_logits, k=min(2, len(raw_logits)))
                    raw_sorted_list = raw_sorted_logits.tolist()
                    greedy_margin = (raw_sorted_list[0] - raw_sorted_list[1]) if len(raw_sorted_list) >= 2 else 0.0

                    # Sampling sensitivity: did temperature change the outcome?
                    sampling_sensitive = next_token_id != greedy_token_id

                    winner_token = logits_entries[0]["token"] if len(logits_entries) > 0 else ""

                    margin_data = {
                        "margin": margin,
                        "winner_token": winner_token,
                        "winner_logit": winner_logit,
                        "runnerup_logit": runnerup_logit,
                        "runnerup_token": runnerup_token,
                        "entropy": margin_entropy,
                        "stability": stability,
                        "greedy_margin": greedy_margin,
                        "sampling_sensitive": sampling_sensitive,
                    }

                    token_alternatives_by_step.append({
                        "step": step,
                        "selected_token": next_token_text,
                        "selected_token_id": next_token_id,
                        "alternatives": alternatives,
                        "logits": logits_entries,
                        "sampling": sampling_metadata,
                        "margin": margin_data,
                    })

                    # Cache hidden states, logits, and full sequence for intervention endpoint
                    if not ephemeral:
                        try:
                            hidden_state_cache.store_step(request_id, step, outputs.hidden_states, raw_logits, current_ids)
                            if step == 0:
                                hidden_state_cache.store_input_ids(request_id, current_ids[:, :-1])  # prompt only
                        except Exception as hs_err:
                            logger.debug(f"Hidden state cache error at step {step}: {hs_err}")

                    # Emit generated token immediately so clients can show code progressively
                    yield sse_event('generated_token', stage=2, totalStages=5,
                                   progress=10 + ((step + 1) / max_tokens) * 20,
                                   stageProgress=((step + 1) / max_tokens) * 100,
                                   detail=f'Generated token {step + 1}/{max_tokens}',
                                   metadata={
                                       'stepIndex': step,
                                       'totalSteps': max_tokens,
                                       'token': next_token_text,
                                       'tokenId': next_token_id,
                                       'generatedTokens': generated_tokens.copy(),
                                   })
                    await asyncio.sleep(0)

                    # === STAGE 3: EXTRACTING (per layer within each token) ===
                    # Optimised: batched tensor ops per layer instead of per-head Python loops
                    # Reduces GPU→CPU sync points from ~4000 to ~40 per token
                    layer_data_this_token = []
                    n_total_layers = len(outputs.attentions)

                    # Margin contribution decomposition: compute the "logit difference direction"
                    # (W_U[winner] - W_U[runner-up]) once, then dot with each layer's residual
                    margin_diff_direction = None
                    winner_token_id_for_decomp = logits_entries[0]["token_id"] if len(logits_entries) > 0 else None
                    runnerup_token_id_for_decomp = logits_entries[1]["token_id"] if len(logits_entries) > 1 else None
                    if winner_token_id_for_decomp is not None and runnerup_token_id_for_decomp is not None:
                        try:
                            lm_head_weight = manager.model.lm_head.weight  # [vocab_size, d_model]
                            margin_diff_direction = (lm_head_weight[winner_token_id_for_decomp] - lm_head_weight[runnerup_token_id_for_decomp]).detach()
                        except Exception:
                            margin_diff_direction = None

                    for layer_idx in range(n_total_layers):
                        # Emit extraction progress (within generating stage for combined progress)
                        if step == max_tokens - 1:  # Only emit detailed layer progress on last token
                            layer_progress = (layer_idx / n_total_layers) * 100
                            overall_progress = 30 + (layer_idx / n_total_layers) * 40  # 30-70%
                            yield sse_event('extracting', stage=3, totalStages=5, progress=overall_progress,
                                           stageProgress=layer_progress,
                                           detail=f'Processing layer {layer_idx + 1}/{n_total_layers}',
                                           metadata={'layerIndex': layer_idx, 'totalLayers': n_total_layers,
                                                    'headsPerLayer': n_heads, 'stepIndex': step, 'totalSteps': max_tokens})
                            if layer_idx % 5 == 0:
                                await asyncio.sleep(0)

                        # --- Per-layer: bulk GPU ops then single CPU transfer ---
                        layer_attn = outputs.attentions[layer_idx][0]  # [n_heads, seq_len, seq_len]
                        current_hidden = outputs.hidden_states[layer_idx + 1]
                        if current_hidden.dim() == 3:
                            current_hidden = current_hidden[0]

                        # Hidden state metrics — 4 values, one .cpu() call
                        last_token_hidden = current_hidden[-1]
                        if layer_idx > 0:
                            prev_hidden = outputs.hidden_states[layer_idx]
                            if prev_hidden.dim() == 3:
                                prev_hidden = prev_hidden[0]
                            hidden_metrics = torch.stack([
                                torch.norm(current_hidden - prev_hidden),
                                torch.norm(current_hidden),
                                torch.std(last_token_hidden),
                                torch.norm(last_token_hidden),
                            ]).cpu().tolist()
                            delta_norm, activation_magnitude, activation_entropy, hidden_state_norm = hidden_metrics
                        else:
                            hidden_metrics = torch.stack([
                                torch.norm(current_hidden),
                                torch.std(last_token_hidden),
                                torch.norm(last_token_hidden),
                            ]).cpu().tolist()
                            activation_magnitude, activation_entropy, hidden_state_norm = hidden_metrics
                            delta_norm = None

                        # Sanitize hidden state metrics
                        activation_magnitude = 0.0 if math.isnan(activation_magnitude) or math.isinf(activation_magnitude) else activation_magnitude
                        activation_entropy = 0.0 if math.isnan(activation_entropy) or math.isinf(activation_entropy) else activation_entropy
                        hidden_state_norm = 0.0 if math.isnan(hidden_state_norm) or math.isinf(hidden_state_norm) else hidden_state_norm
                        if delta_norm is not None:
                            delta_norm = 0.0 if math.isnan(delta_norm) or math.isinf(delta_norm) else delta_norm

                        # Margin contribution decomposition:
                        # margin_contribution = (W_U[winner] - W_U[runner-up]) · (h_{ℓ+1} - h_ℓ)
                        # This causally attributes the final margin to each layer's residual contribution.
                        margin_contribution = None
                        if margin_diff_direction is not None:
                            try:
                                if layer_idx > 0:
                                    prev_h = outputs.hidden_states[layer_idx]
                                    if prev_h.dim() == 3:
                                        prev_h = prev_h[0]
                                    residual = current_hidden[-1] - prev_h[-1]
                                else:
                                    # Layer 0: the embedding contribution
                                    residual = current_hidden[-1]
                                mc = torch.dot(margin_diff_direction, residual).item()
                                margin_contribution = 0.0 if math.isnan(mc) or math.isinf(mc) else mc
                            except Exception:
                                margin_contribution = None

                        # --- Batched head processing: all heads at once on GPU ---
                        num_heads_layer = layer_attn.shape[0]

                        # Last-row attention weights for all heads: [n_heads, seq_len]
                        all_last_row = layer_attn[:, -1, :]

                        # Max weight per head: [n_heads] — single GPU op
                        all_max_weights = all_last_row.max(dim=-1).values

                        # Entropy of last-row per head: [n_heads] — single GPU op
                        all_entropies = -(all_last_row * torch.log(all_last_row + 1e-10)).sum(dim=-1)

                        # Normalized average entropy per head (latter half of query positions)
                        # layer_attn: [n_heads, q_len, k_len]
                        q_len = layer_attn.shape[1]
                        # Raw entropy per query position per head: [n_heads, q_len]
                        all_token_entropies = -(layer_attn * torch.log(layer_attn + 1e-10)).sum(dim=-1)
                        # Normalize by log(position): [q_len]
                        positions = torch.arange(1, q_len + 1, device=layer_attn.device, dtype=layer_attn.dtype)
                        max_ents = torch.log(positions + 1e-10)  # [q_len]
                        all_normalized = all_token_entropies / (max_ents.unsqueeze(0) + 1e-10)  # [n_heads, q_len]
                        # Average over latter half: [n_heads]
                        start_idx = q_len // 2
                        if start_idx < q_len:
                            all_avg_entropies = all_normalized[:, start_idx:].mean(dim=-1)
                        else:
                            all_avg_entropies = all_normalized.mean(dim=-1)

                        # Previous-token weights for pattern detection: [n_heads]
                        all_prev_token_weights = all_last_row[:, -2] if all_last_row.shape[1] >= 2 else torch.zeros(num_heads_layer, device=layer_attn.device)

                        # Attention sink weights: sum of attention on positions 0-2 per head [n_heads]
                        seq_len_attn = all_last_row.shape[1]
                        all_sink_weights = all_last_row[:, :min(3, seq_len_attn)].sum(dim=-1)

                        # Local attention weights: sum within 5 positions of query per head [n_heads]
                        all_local_weights = all_last_row[:, max(0, seq_len_attn - 5):].sum(dim=-1) if seq_len_attn > 5 else torch.ones(num_heads_layer, device=layer_attn.device)

                        # Induction detection: attention to positions following previous occurrences of current token
                        all_induction_weights = torch.zeros(num_heads_layer, device=layer_attn.device)
                        if step > 0:
                            current_token = current_ids[0, -1]
                            prev_occurrences = (current_ids[0, :-1] == current_token).nonzero(as_tuple=True)[0]
                            if len(prev_occurrences) > 0:
                                following_positions = prev_occurrences + 1
                                following_positions = following_positions[following_positions < seq_len_attn]
                                if len(following_positions) > 0:
                                    all_induction_weights = all_last_row[:, following_positions].sum(dim=-1)

                        # Single bulk transfer: all head metrics to CPU
                        head_metrics_gpu = torch.stack([
                            all_max_weights, all_entropies, all_avg_entropies, all_prev_token_weights,
                            all_sink_weights, all_local_weights, all_induction_weights
                        ])  # [7, n_heads]
                        head_metrics_cpu = head_metrics_gpu.cpu().tolist()  # one sync point
                        max_weights_list = head_metrics_cpu[0]
                        entropies_list = head_metrics_cpu[1]
                        avg_entropies_list = head_metrics_cpu[2]
                        prev_token_list = head_metrics_cpu[3]
                        sink_weights_list = head_metrics_cpu[4]
                        local_weights_list = head_metrics_cpu[5]
                        induction_weights_list = head_metrics_cpu[6]

                        # Bulk transfer attention matrices to CPU: one .cpu() for entire layer
                        layer_attn_cpu = layer_attn.cpu().float().numpy()  # [n_heads, seq_len, seq_len]

                        # QKV matrices (already on CPU from hooks)
                        qkv_layer = qkv_captures.get(layer_idx)

                        # Build per-head metadata from CPU-side data (no more GPU calls)
                        critical_heads = []
                        for head_idx in range(num_heads_layer):
                            mw = max_weights_list[head_idx]
                            ent = entropies_list[head_idx]
                            avg_ent = avg_entropies_list[head_idx]
                            ptw = prev_token_list[head_idx]
                            skw = sink_weights_list[head_idx]
                            lcw = local_weights_list[head_idx]
                            idw = induction_weights_list[head_idx]

                            # Sanitize
                            mw = 0.0 if math.isnan(mw) or math.isinf(mw) else mw
                            ent = 0.0 if math.isnan(ent) or math.isinf(ent) else ent
                            avg_ent = 0.0 if math.isnan(avg_ent) or math.isinf(avg_ent) else avg_ent

                            # Score-all-then-rank head classification
                            # Behaviour type scores (attention geometry)
                            behaviour_scores = {
                                "attention_sink": skw,
                                "previous_token": ptw,
                                "local": lcw,
                                "induction": min(1.0, idw),
                                "focused": max(0.0, 1.0 - ent) if ent < 1.5 else 0.0,
                                "diffuse": min(1.0, max(0.0, (ent - 1.0) / 2.0)),
                            }
                            behaviour_thresholds = {
                                "attention_sink": 0.4,
                                "previous_token": 0.7,
                                "local": 0.5,
                                "induction": 0.2,
                                "focused": 0.3,
                                "diffuse": 0.3,
                            }
                            qualified = {
                                k: v for k, v in behaviour_scores.items()
                                if v >= behaviour_thresholds.get(k, 0.3)
                            }
                            sorted_behaviours = sorted(qualified.items(), key=lambda x: x[1], reverse=True)
                            primary = sorted_behaviours[0] if sorted_behaviours else ("diffuse", behaviour_scores["diffuse"])
                            secondary = sorted_behaviours[1] if len(sorted_behaviours) > 1 else None

                            pattern_type = primary[0]
                            confidence = primary[1]
                            confidence = 0.0 if math.isnan(confidence) or math.isinf(confidence) else confidence

                            # Code cue scores (what code tokens are attended to)
                            # Decode token texts once per step (cached via nonlocal)
                            if step_token_texts_cache.get('step') != step:
                                try:
                                    step_token_texts_cache['texts'] = [
                                        manager.tokenizer.decode([tid]) for tid in current_ids[0, :seq_len_attn].tolist()
                                    ]
                                except Exception:
                                    step_token_texts_cache['texts'] = []
                                step_token_texts_cache['step'] = step
                            token_texts = step_token_texts_cache.get('texts', [])

                            code_cues = {}
                            if len(token_texts) == seq_len_attn:
                                head_weights = all_last_row[head_idx].cpu()
                                delimiters = {'(', ')', '{', '}', '[', ']', ':', ';', ','}
                                delim_indices = [i for i, t in enumerate(token_texts) if t.strip() in delimiters]
                                if delim_indices:
                                    code_cues["delimiter_sensitive"] = head_weights[delim_indices].sum().item()

                                keywords = {'def', 'return', 'if', 'else', 'elif', 'for', 'while', 'class',
                                           'import', 'from', 'try', 'except', 'with', 'as', 'in', 'not',
                                           'and', 'or', 'True', 'False', 'None', 'self', 'yield', 'async',
                                           'await', 'lambda', 'raise', 'pass', 'break', 'continue',
                                           'function', 'const', 'let', 'var', 'new', 'this',
                                           'public', 'private', 'static', 'void', 'int', 'string', 'bool',
                                           'namespace', 'using', 'class', 'interface', 'override', 'virtual'}
                                kw_indices = [i for i, t in enumerate(token_texts) if t.strip() in keywords]
                                if kw_indices:
                                    code_cues["keyword_sensitive"] = head_weights[kw_indices].sum().item()

                                if idw > 0.15:
                                    code_cues["pattern_reuse"] = min(1.0, idw * 1.5)

                            cue_threshold = 0.15
                            sorted_cues = sorted(
                                [(k, round(v, 4)) for k, v in code_cues.items() if v >= cue_threshold],
                                key=lambda x: x[1], reverse=True
                            )
                            primary_cue = sorted_cues[0] if sorted_cues else None

                            attention_matrix = layer_attn_cpu[head_idx]

                            q_matrix = None
                            k_matrix = None
                            v_matrix = None
                            if qkv_layer is not None:
                                q_matrix = qkv_layer['q'][:, head_idx, :].float().numpy()
                                k_matrix = qkv_layer['k'][:, head_idx, :].float().numpy()
                                v_matrix = qkv_layer['v'][:, head_idx, :].float().numpy()

                            if not ephemeral:
                                matrix_cache.store(request_id, step, layer_idx, head_idx, {
                                    "attention_weights": attention_matrix,
                                    "q_matrix": q_matrix,
                                    "k_matrix": k_matrix,
                                    "v_matrix": v_matrix
                                })

                            head_entry = {
                                "head_idx": head_idx,
                                "entropy": ent,
                                "avg_entropy": avg_ent,
                                "max_weight": mw,
                                "has_matrices": attention_matrix is not None,
                                "pattern": {"type": pattern_type, "confidence": round(confidence, 4)} if pattern_type else None,
                            }
                            if secondary:
                                head_entry["secondary_behaviour"] = {"type": secondary[0], "score": round(secondary[1], 4)}
                            if primary_cue:
                                head_entry["code_cue"] = {
                                    "type": primary_cue[0],
                                    "score": primary_cue[1],
                                    "evidence": f"{round(primary_cue[1] * 100)}% attention on {primary_cue[0].replace('_', ' ')} tokens",
                                }
                                if len(sorted_cues) > 1:
                                    head_entry["secondary_cue"] = {"type": sorted_cues[1][0], "score": sorted_cues[1][1]}
                            critical_heads.append(head_entry)

                        critical_heads.sort(key=lambda h: h["max_weight"], reverse=True)

                        # Layer-level pattern: majority vote of head patterns, weighted by confidence
                        pattern_votes = {}
                        for h in critical_heads:
                            if h["pattern"] and h["pattern"]["type"]:
                                pt = h["pattern"]["type"]
                                pc = h["pattern"]["confidence"]
                                pattern_votes[pt] = pattern_votes.get(pt, 0.0) + pc
                        layer_pattern = None
                        if pattern_votes:
                            best_type = max(pattern_votes, key=pattern_votes.get)
                            total_conf = sum(pattern_votes.values())
                            layer_pattern = {
                                "type": best_type,
                                "confidence": round(pattern_votes[best_type] / total_conf, 3) if total_conf > 0 else 0.0
                            }

                        layer_entry = {
                            "layer_idx": layer_idx,
                            "pattern": layer_pattern,
                            "critical_heads": critical_heads,
                            "activation_magnitude": activation_magnitude,
                            "activation_entropy": activation_entropy,
                            "hidden_state_norm": hidden_state_norm,
                            "delta_norm": delta_norm,
                            "margin_contribution": margin_contribution,
                        }
                        # Phase 4: Attention and MLP output norms + contribution ratios
                        if layer_idx in attn_output_norms:
                            layer_entry["attn_output_norm"] = round(attn_output_norms[layer_idx], 4)
                        if layer_idx in mlp_output_norms:
                            layer_entry["mlp_output_norm"] = round(mlp_output_norms[layer_idx], 4)
                        if layer_idx in attn_output_norms and layer_idx in mlp_output_norms:
                            attn_n = attn_output_norms[layer_idx]
                            mlp_n = mlp_output_norms[layer_idx]
                            total = attn_n + mlp_n
                            if total > 0:
                                layer_entry["attn_contribution"] = round(attn_n / total, 4)
                                layer_entry["ffn_contribution"] = round(mlp_n / total, 4)
                        if layer_idx in gate_activation_stats:
                            layer_entry["gate_stats"] = gate_activation_stats[layer_idx]
                        if layer_idx in ffn_top_neurons:
                            # `ffn_top_neurons` holds top-K SwiGLU
                            # intermediate activations at the
                            # predicting position — see
                            # make_ffn_top_neurons_hook for the
                            # per-neuron value semantics. The frontend
                            # uses this to render Step 5's per-neuron
                            # bar chart inside "Behind the SwiGLU"
                            # without the backend needing a separate
                            # endpoint or a 14,336-d cache entry.
                            layer_entry["ffn_top_neurons"] = ffn_top_neurons[layer_idx]

                        # Phase 5: Logit lens at every layer.
                        # The lens projection is a single matmul
                        # [1, hidden] × [hidden, vocab] ≈ 0.5-1ms on
                        # Apple Silicon MPS — 40 layers × N_steps adds
                        # only ~5-10s to a typical analyse run. Worth
                        # the cost: every-layer sampling eliminates
                        # the (L17, L24]-style uncertainty windows the
                        # frontend's crystallisation narrative would
                        # otherwise have to qualify, giving developers
                        # a layer-precise commit point for trust
                        # decisions about the model's prediction.
                        logit_lens_stride = 1
                        if layer_idx % logit_lens_stride == 0 or layer_idx == n_layers - 1:
                            try:
                                hidden_for_lens = current_hidden[-1].unsqueeze(0)  # [1, hidden_dim]
                                # Apply final layer norm then project through lm_head
                                if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
                                    normed = manager.model.model.norm(hidden_for_lens)
                                    lens_logits = manager.model.lm_head(normed)[0]  # [vocab_size]
                                elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
                                    normed = manager.model.transformer.ln_f(hidden_for_lens)
                                    lens_logits = manager.model.lm_head(normed)[0]
                                else:
                                    lens_logits = None

                                if lens_logits is not None:
                                    lens_probs = torch.softmax(lens_logits, dim=-1)
                                    top_probs, top_ids = torch.topk(lens_probs, k=5)
                                    top_probs_list = top_probs.cpu().tolist()
                                    top_ids_list = top_ids.cpu().tolist()
                                    lens_entries = []
                                    for tp, tid in zip(top_probs_list, top_ids_list):
                                        lens_entries.append({
                                            "token": manager.tokenizer.decode([tid], skip_special_tokens=False),
                                            "probability": tp
                                        })
                                    layer_entry["logit_lens_top"] = lens_entries

                                    # Layer-wise margin tracking (raw logit diff between top-1 and top-2)
                                    top2_logits, top2_ids = torch.topk(lens_logits, k=min(2, len(lens_logits)))
                                    top2_logits_list = top2_logits.cpu().tolist()
                                    top2_ids_list = top2_ids.cpu().tolist()
                                    layer_winner_token = manager.tokenizer.decode([top2_ids_list[0]], skip_special_tokens=False)
                                    layer_runnerup_token = manager.tokenizer.decode([top2_ids_list[1]], skip_special_tokens=False) if len(top2_ids_list) > 1 else ""
                                    layer_margin_val = (top2_logits_list[0] - top2_logits_list[1]) if len(top2_logits_list) > 1 else 0.0
                                    layer_entry["layer_margin"] = layer_margin_val
                                    layer_entry["layer_winner"] = layer_winner_token
                                    layer_entry["layer_runnerup"] = layer_runnerup_token

                                    # Tuned lens: apply per-layer affine correction
                                    from .tuned_lens import tuned_lens_runtime
                                    if tuned_lens_runtime.available:
                                        try:
                                            corrected = tuned_lens_runtime.apply(layer_idx, hidden_for_lens)
                                            tuned_normed = normed.__class__  # reuse same LN path
                                            # Re-apply final LN + lm_head on corrected hidden
                                            if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
                                                tuned_normed = manager.model.model.norm(corrected)
                                                tuned_logits = manager.model.lm_head(tuned_normed)[0]
                                            elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
                                                tuned_normed = manager.model.transformer.ln_f(corrected)
                                                tuned_logits = manager.model.lm_head(tuned_normed)[0]
                                            else:
                                                tuned_logits = None

                                            if tuned_logits is not None:
                                                tuned_probs = torch.softmax(tuned_logits, dim=-1)
                                                tuned_top_probs, tuned_top_ids = torch.topk(tuned_probs, k=5)
                                                tuned_entries = []
                                                for tp, tid in zip(tuned_top_probs.cpu().tolist(), tuned_top_ids.cpu().tolist()):
                                                    tuned_entries.append({
                                                        "token": manager.tokenizer.decode([tid], skip_special_tokens=False),
                                                        "probability": tp
                                                    })
                                                layer_entry["tuned_lens_top"] = tuned_entries

                                                tuned_top2_logits, tuned_top2_ids = torch.topk(tuned_logits, k=min(2, len(tuned_logits)))
                                                tuned_top2_logits_list = tuned_top2_logits.cpu().tolist()
                                                tuned_top2_ids_list = tuned_top2_ids.cpu().tolist()
                                                layer_entry["tuned_layer_winner"] = manager.tokenizer.decode([tuned_top2_ids_list[0]], skip_special_tokens=False)
                                                layer_entry["tuned_layer_runnerup"] = manager.tokenizer.decode([tuned_top2_ids_list[1]], skip_special_tokens=False) if len(tuned_top2_ids_list) > 1 else ""
                                                layer_entry["tuned_layer_margin"] = (tuned_top2_logits_list[0] - tuned_top2_logits_list[1]) if len(tuned_top2_logits_list) > 1 else 0.0
                                        except Exception as tuned_err:
                                            logger.debug(f"Tuned lens error at layer {layer_idx}: {tuned_err}")
                            except Exception as lens_err:
                                logger.debug(f"Logit lens error at layer {layer_idx}: {lens_err}")

                        layer_data_this_token.append(layer_entry)

                    layer_data_by_token.append(layer_data_this_token)

                    # Update inputs
                    next_token_tensor = torch.tensor([[next_token_id]], dtype=torch.long, device=manager.device)
                    current_ids = torch.cat([current_ids, next_token_tensor], dim=1)

                    # Stop on EOS
                    if next_token_id == manager.tokenizer.eos_token_id:
                        break

                    # Free memory from this step's outputs to prevent accumulation
                    # This is critical for large models like Devstral (40 layers, 32 heads)
                    del outputs
                    del logits
                    del probs
                    if 'layer_attn' in dir():
                        del layer_attn
                    if 'current_hidden' in dir():
                        del current_hidden

                    # Periodic garbage collection for large models (every 8 steps)
                    if (step + 1) % 8 == 0:
                        gc.collect()
                        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                            torch.mps.empty_cache() if hasattr(torch.mps, 'empty_cache') else None

            # Clean up hooks
            for hook in hooks:
                hook.remove()

            # === STAGE 4: SERIALIZING ===
            yield sse_event('serializing', stage=4, totalStages=5, progress=75,
                           stageProgress=0, detail='Building response data...')
            await asyncio.sleep(0)

            qkv_by_layer_head = {}
            generation_time = time.time() - start_time

            # Calculate token section boundaries
            total_tokens = prompt_length + len(generated_token_ids)
            system_prompt_text = system_prompt_override or (model_config.get("system_prompt") if model_config else None)

            system_prompt_end = 0
            if prompt_style == "instruction" and system_prompt_text:
                if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
                    try:
                        no_system_tokens = manager.mistral_tokenizer.encode_chat("", prompt)
                        # Add 1 to include the closing [/SYSTEM_PROMPT] tag in system section
                        system_prompt_end = prompt_length - len(no_system_tokens) + 1
                        system_prompt_end = max(0, min(system_prompt_end, prompt_length))
                    except Exception:
                        system_prompt_end = 0
                else:
                    total_chars = len(system_prompt_text or "") + len(prompt)
                    if total_chars > 0:
                        system_ratio = len(system_prompt_text or "") / total_chars
                        system_prompt_end = int(prompt_length * system_ratio)

            token_sections = {
                "systemPrompt": {
                    "start": 0,
                    "end": system_prompt_end,
                    "text": system_prompt_text,
                    "tokenCount": system_prompt_end
                },
                "userPrompt": {
                    "start": system_prompt_end,
                    "end": prompt_length,
                    "text": prompt,
                    "tokenCount": prompt_length - system_prompt_end
                },
                "output": {
                    "start": prompt_length,
                    "end": total_tokens,
                    "text": "".join(generated_tokens),
                    "tokenCount": len(generated_token_ids)
                }
            }

            yield sse_event('serializing', stage=4, totalStages=5, progress=82,
                           stageProgress=50, detail='Building token metadata...')
            await asyncio.sleep(0)

            # Build token metadata
            from .tokenizer_utils import TokenizerMetadata
            token_metadata_builder = TokenizerMetadata(manager.tokenizer)

            # Include every id the tokenizer considers a special / added token
            # (BOS, EOS, PAD, UNK, chat-template delimiters like [INST]/[/INST],
            # system-prompt markers, tool-call markers, etc.). See the matching
            # construction in the non-streaming endpoint for rationale.
            special_token_ids_set = set(getattr(manager.tokenizer, "all_special_ids", []) or [])
            for tok_id in (
                manager.tokenizer.eos_token_id,
                manager.tokenizer.bos_token_id,
                manager.tokenizer.pad_token_id,
                manager.tokenizer.unk_token_id,
            ):
                if tok_id is not None:
                    special_token_ids_set.add(tok_id)
            if manager.mistral_tokenizer is not None:
                special_token_ids_set |= manager.mistral_tokenizer.get_control_token_ids()

            def build_token_data(token_ids, token_texts, token_type):
                multi_split_flags = token_metadata_builder.is_multi_split_identifier(token_ids)
                result = []
                for i, (tid, t) in enumerate(zip(token_ids, token_texts)):
                    bpe_pieces = token_metadata_builder.get_subword_pieces(tid)
                    result.append({
                        "text": t,
                        "idx": tid,
                        "bytes": len(t.encode('utf-8')),
                        "type": token_type,
                        "bpe_pieces": bpe_pieces,
                        "is_special": tid in special_token_ids_set,
                        "is_multi_split": multi_split_flags[i] if i < len(multi_split_flags) else False,
                        "num_pieces": len(bpe_pieces),
                    })
                return result

            # Compute margin statistics and commitment summary
            margin_stats = {"fragile_count": 0, "boundary_count": 0, "moderate_count": 0, "stable_count": 0}
            commitment_layers = []
            flip_count = 0
            for step_data in token_alternatives_by_step:
                m = step_data.get("margin", {})
                stab = m.get("stability", "stable")
                if stab == "fragile":
                    margin_stats["fragile_count"] += 1
                elif stab == "boundary":
                    margin_stats["boundary_count"] += 1
                elif stab == "moderate":
                    margin_stats["moderate_count"] += 1
                else:
                    margin_stats["stable_count"] += 1

            # Commitment layer and flip detection from layer data
            for step_idx, step_layers in enumerate(layer_data_by_token):
                lens_layers = [l for l in step_layers if l.get("layer_margin") is not None]
                if not lens_layers:
                    continue
                # Find commitment layer: first layer where margin > 0.3 and stays positive
                step_commitment = None
                for i, ll in enumerate(lens_layers):
                    if ll["layer_margin"] > 0.3:
                        stays_positive = all(lens_layers[j]["layer_margin"] > 0 for j in range(i, len(lens_layers)))
                        if stays_positive:
                            step_commitment = ll["layer_idx"]
                            break
                if step_commitment is not None:
                    commitment_layers.append(step_commitment)
                # Count flips: where winner changes between consecutive lens layers
                for i in range(1, len(lens_layers)):
                    prev_winner = (lens_layers[i-1].get("layer_winner") or "").strip()
                    curr_winner = (lens_layers[i].get("layer_winner") or "").strip()
                    if prev_winner and curr_winner and prev_winner != curr_winner:
                        flip_count += 1

            avg_commitment = sum(commitment_layers) / len(commitment_layers) if commitment_layers else n_layers
            late_threshold = n_layers * 0.75
            late_count = sum(1 for cl in commitment_layers if cl > late_threshold)

            commitment_summary = {
                "avg_commitment_layer": round(avg_commitment, 1),
                "late_commitment_count": late_count,
                "flip_count": flip_count,
            }

            # Tuned lens commitment summary (parallel to raw)
            tuned_commitment_summary = None
            from .tuned_lens import tuned_lens_runtime
            if tuned_lens_runtime.available:
                tuned_commitment_layers = []
                tuned_flip_count = 0
                for step_idx, step_layers in enumerate(layer_data_by_token):
                    tuned_lens_layers = [l for l in step_layers if l.get("tuned_layer_margin") is not None]
                    if not tuned_lens_layers:
                        continue
                    step_commitment = None
                    for i, ll in enumerate(tuned_lens_layers):
                        if ll["tuned_layer_margin"] > 0.3:
                            stays_positive = all(tuned_lens_layers[j]["tuned_layer_margin"] > 0 for j in range(i, len(tuned_lens_layers)))
                            if stays_positive:
                                step_commitment = ll["layer_idx"]
                                break
                    if step_commitment is not None:
                        tuned_commitment_layers.append(step_commitment)
                    for i in range(1, len(tuned_lens_layers)):
                        prev_w = (tuned_lens_layers[i-1].get("tuned_layer_winner") or "").strip()
                        curr_w = (tuned_lens_layers[i].get("tuned_layer_winner") or "").strip()
                        if prev_w and curr_w and prev_w != curr_w:
                            tuned_flip_count += 1

                tuned_avg = sum(tuned_commitment_layers) / len(tuned_commitment_layers) if tuned_commitment_layers else n_layers
                tuned_late = sum(1 for cl in tuned_commitment_layers if cl > late_threshold)
                tuned_commitment_summary = {
                    "avg_commitment_layer": round(tuned_avg, 1),
                    "late_commitment_count": tuned_late,
                    "flip_count": tuned_flip_count,
                }

            # === Embedding L2 norms — per-input-token layer-0 anchor ===
            # One float per token in the FULL context (prompt +
            # generated): ‖e_i‖ where e_i is the row the embedding
            # matrix returns for token i. The Mechanism lens uses this
            # as the real-data anchor for the residual-norm trace
            # shown in the Microscope rail (RQ1: developer-
            # interpretable architectural signal at the embedding
            # stage). Generated tokens are included so the chart
            # reflects the actual input the model sees at every
            # selected step. Memory is bounded by the number of
            # *unique* token ids — see _compute_embedding_norms.
            embedding_norms: List[float] = _compute_embedding_norms(
                manager,
                list(prompt_token_ids) + list(generated_token_ids),
            )

            # Build response
            response = {
                "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
                "prompt": prompt,
                "promptTokens": build_token_data(prompt_token_ids, prompt_tokens, "prompt"),
                "generatedTokens": build_token_data(generated_token_ids, generated_tokens, "generated"),
                "tokenSections": token_sections,
                "tokenAlternatives": token_alternatives_by_step,
                "layersDataByStep": layer_data_by_token,
                "layersData": layer_data_by_token[-1] if layer_data_by_token else [],
                "qkvData": {},  # Deprecated: matrices now lazy-loaded via /matrix endpoint
                "modelInfo": {
                    "numLayers": n_layers,
                    "numHeads": n_heads,
                    "modelDimension": d_model,
                    "headDim": head_dim,
                    "vocabSize": manager.model.config.vocab_size,
                    "tunedLensAvailable": tuned_lens_runtime.available,
                    "ffnType": ffn_type,
                    "intermediateSize": getattr(manager.model.config, 'intermediate_size', None),
                },
                "embeddingNorms": embedding_norms,
                "generationTime": generation_time,
                "numTokensGenerated": len(generated_tokens),
                "marginStats": margin_stats,
                "commitmentSummary": commitment_summary,
                **({"tunedCommitmentSummary": tuned_commitment_summary} if tuned_commitment_summary else {}),
            }

            # Estimate response size
            response_json = json.dumps(sanitize_for_json(response))
            response_size_bytes = len(response_json.encode('utf-8'))

            yield sse_event('serializing', stage=4, totalStages=5, progress=90,
                           stageProgress=100, detail=f'Response ready ({response_size_bytes / 1024 / 1024:.1f}MB)',
                           metadata={'responseSizeBytes': response_size_bytes})
            await asyncio.sleep(0)

            # === STAGE 5: COMPLETE ===
            yield sse_event('complete', stage=5, totalStages=5, progress=95,
                           stageProgress=0, detail='Transferring data...',
                           metadata={'responseSizeBytes': response_size_bytes, 'generationTimeMs': int(generation_time * 1000)})

            logger.info(f"✅ [SSE] Research attention analysis complete: {len(generated_tokens)} tokens, {generation_time:.2f}s, {response_size_bytes / 1024 / 1024:.1f}MB")

            # Send final result
            yield sse_event('result', data=sanitize_for_json(response))

        except Exception as e:
            logger.error(f"[SSE] Research attention analysis error: {e}")
            logger.error(traceback.format_exc())
            yield sse_event('error', detail=str(e), stage=0, totalStages=5, progress=0, stageProgress=0)

    return StreamingResponse(
        event_generator(),
        media_type='text/event-stream',
        headers={
            'Cache-Control': 'no-cache, no-store, must-revalidate',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no',  # Disable nginx/proxy buffering
        }
    )


@app.get("/analyze/research/attention/matrix")
async def get_attention_matrix(
    request_id: str,
    step: int,
    layer: int,
    head: Optional[int] = None,
    aggregate_mode: str = "mean",
    authenticated: bool = Depends(verify_api_key),
):
    """
    Retrieve cached attention/QKV matrices for a specific head, OR an
    aggregated attention matrix across all heads when `head` is omitted.

    Used for lazy-loading matrix data when the user opens an inline
    head panel (per-head Q/K/V) or the Mechanism lens's Step 4
    attention matrix view (aggregate by default). Matrices are cached
    during the initial analysis and available for 60 minutes.

    Parameters:
    - request_id: UUID from the original analysis response
    - step: Generation step (0 = first generated token)
    - layer: Layer index (0-based)
    - head: Head index (0-based). Omit for an aggregate matrix across
      all heads — the returned payload then contains only
      `attention_weights` (the Q/K/V projections are per-head and
      have no canonical aggregate form).
    - aggregate_mode: "mean" or "max" when `head` is omitted.

    Returns:
    - attention_weights: [seq_len, seq_len] attention matrix
    - q_matrix / k_matrix / v_matrix: [seq_len, head_dim] projections
      (only when `head` is specified)
    - layer: Layer index
    - head: Head index, or null when aggregated
    - aggregate_mode: Aggregation mode used, or null for per-head
    """
    if head is not None:
        data = matrix_cache.get(request_id, step, layer, head)
        if data is None:
            logger.warning(
                f"Matrix cache miss: request_id={request_id}, step={step}, layer={layer}, head={head}"
            )
            raise HTTPException(
                status_code=404,
                detail="Matrix data not found. Cache may have expired (60 min TTL). Please re-analyze.",
            )
        logger.info(
            f"Matrix cache hit: request_id={request_id}, step={step}, layer={layer}, head={head}"
        )
        # Convert numpy arrays to lists for JSON serialization. Arrays
        # are stored as numpy for memory efficiency, converted on-
        # demand here.
        response = {}
        for key, value in data.items():
            if value is not None and hasattr(value, "tolist"):
                response[key] = value.tolist()
            else:
                response[key] = value
        response["layer"] = layer
        response["head"] = head
        response["aggregate_mode"] = None
        return response

    # Aggregate path — averages or max-pools attention matrices across
    # heads. Q/K/V projections are intentionally omitted; they're
    # per-head quantities and don't have a canonical aggregate.
    aggregate = matrix_cache.get_aggregate_attention_matrix(
        request_id, step, layer, aggregate_mode
    )
    if aggregate is None:
        logger.warning(
            f"Aggregate matrix cache miss: request_id={request_id}, step={step}, layer={layer}, mode={aggregate_mode}"
        )
        raise HTTPException(
            status_code=404,
            detail="Matrix data not found. Cache may have expired (60 min TTL). Please re-analyze.",
        )
    logger.info(
        f"Aggregate matrix cache hit: request_id={request_id}, step={step}, layer={layer}, mode={aggregate_mode}"
    )
    return {
        "attention_weights": aggregate,
        "layer": layer,
        "head": None,
        "aggregate_mode": aggregate_mode,
    }


@app.get("/analyze/research/attention/matrix/stats")
async def get_matrix_cache_stats(authenticated: bool = Depends(verify_api_key)):
    """Return matrix cache statistics for monitoring."""
    return matrix_cache.get_stats()


@app.get("/analyze/research/attention/row")
async def get_attention_row(
    request_id: str,
    step: int,
    layer: int,
    head: Optional[int] = None,
    aggregate_mode: str = "mean",
    authenticated: bool = Depends(verify_api_key)
):
    """
    Retrieve single attention row for overlay visualization.

    Returns the attention weights from the query token (at position `step`)
    to all preceding positions. This is a minimal payload for efficient
    lazy-loading in the attention overlay feature.

    Parameters:
    - request_id: UUID from the original analysis response
    - step: Generation step (0 = first generated token)
    - layer: Layer index (0-based)
    - head: Head index (0-based), or None for aggregated view
    - aggregate_mode: "mean" or "max" when head is None

    Returns:
    - attention_weights: List of attention weights [0..seq_len]
    - seq_len: Number of positions in the sequence
    - layer: Layer index
    - head: Head index (null if aggregated)
    - aggregate_mode: Mode used if aggregated (null otherwise)
    - query_position: Sequence index whose attention this row represents.
      The model selects the step-N token from logits at this position;
      the cached attention[-1] is therefore *this* position's attention,
      not the predicted token's self-attention.
    - query_token_id: Token id at query_position (last context token).
    """
    # Aggregation parameters (num_heads in particular) come from per-request
    # metadata recorded at analysis time, not from the live model config.
    # This keeps aggregate rows correct even after the user switches models
    # mid-session — the row is averaged over the heads of the model that
    # actually produced the cached attention.
    query_info = matrix_cache.get_step_query_info(request_id, step)
    query_position = query_info.get("query_position") if query_info else None
    query_token_id = query_info.get("query_token_id") if query_info else None

    if head is not None:
        # Fetch specific head — no aggregation, no metadata lookup needed.
        attention_row = matrix_cache.get_attention_row(request_id, step, layer, head)
        if attention_row is None:
            logger.warning(f"Attention row cache miss: request_id={request_id}, step={step}, layer={layer}, head={head}")
            raise HTTPException(
                status_code=404,
                detail="Attention data not found. Cache may have expired (60 min TTL). Please re-analyze."
            )
        logger.info(f"Attention row cache hit: request_id={request_id}, step={step}, layer={layer}, head={head}")
        return {
            "attention_weights": attention_row,
            "seq_len": len(attention_row),
            "layer": layer,
            "head": head,
            "aggregate_mode": None,
            "query_position": query_position,
            "query_token_id": query_token_id,
        }
    else:
        # Aggregate across all heads. num_heads is derived inside
        # get_aggregate_row from the request's stored metadata.
        attention_row = matrix_cache.get_aggregate_row(
            request_id, step, layer, aggregate_mode
        )
        if attention_row is None:
            logger.warning(f"Attention row aggregate cache miss: request_id={request_id}, step={step}, layer={layer}")
            raise HTTPException(
                status_code=404,
                detail="Attention data not found. Cache may have expired (60 min TTL). Please re-analyze."
            )
        logger.info(f"Attention row aggregate cache hit: request_id={request_id}, step={step}, layer={layer}, mode={aggregate_mode}")
        return {
            "attention_weights": attention_row,
            "seq_len": len(attention_row),
            "layer": layer,
            "head": None,
            "aggregate_mode": aggregate_mode,
            "query_position": query_position,
            "query_token_id": query_token_id,
        }


# --- Phase 2: Intervention endpoint ---

class InterventionRequest(BaseModel):
    request_id: str
    step: int
    intervention_type: str  # "mask_system" | "mask_user_span" | "mask_generated" | "mask_positions" | "greedy" | "temperature_sweep" | "layer_ablation" | "head_ablation" | "expert_mask"
    params: dict = {}

class InterventionResponse(BaseModel):
    original_margin: float
    recomputed_margin: float
    margin_shift: float
    original_stability: str
    recomputed_stability: str
    original_winner: str
    recomputed_winner: str
    winner_changed: bool
    details: dict = {}

@app.post("/analyze/intervention")
async def run_intervention(request: InterventionRequest, authenticated: bool = Depends(verify_api_key)):
    """
    Run an input ablation on a cached generation run.
    Re-evaluates a token position under modified conditions (input masking, component ablation, temperature sweep).
    """
    if not manager.model:
        raise HTTPException(status_code=503, detail="Model not loaded")

    if not hidden_state_cache.has_run(request.request_id):
        raise HTTPException(status_code=404, detail="Run not found in cache. Cache may have expired (60 min TTL). Please re-generate.")

    cached_logits = hidden_state_cache.get_logits(request.request_id, request.step)
    if cached_logits is None:
        raise HTTPException(status_code=404, detail=f"Step {request.step} not found in cached run.")

    try:
        # Move logits to compute device
        raw_logits = cached_logits.to(manager.device)

        # Original margin (from raw logits)
        top2_orig, top2_orig_ids = torch.topk(raw_logits, k=2)
        top2_orig_list = top2_orig.cpu().tolist()
        top2_orig_ids_list = top2_orig_ids.cpu().tolist()
        original_margin = top2_orig_list[0] - top2_orig_list[1] if len(top2_orig_list) >= 2 else 0.0
        original_winner = manager.tokenizer.decode([top2_orig_ids_list[0]], skip_special_tokens=False)

        if request.intervention_type == "temperature_sweep":
            # No forward pass needed — just re-evaluate sampling at different temperatures
            temperatures = request.params.get("temperatures", [0.0, 0.05, 0.1, 0.15, 0.2, 0.3])
            results_per_temp = []
            greedy_id = torch.argmax(raw_logits).item()
            greedy_token = manager.tokenizer.decode([greedy_id], skip_special_tokens=False)

            for temp in temperatures:
                if temp == 0 or temp < 1e-6:
                    winner_id = greedy_id
                else:
                    scaled = raw_logits / temp
                    probs = torch.softmax(scaled, dim=0)
                    winner_id = torch.argmax(probs).item()  # Most likely at this temp
                winner_token = manager.tokenizer.decode([winner_id], skip_special_tokens=False)
                results_per_temp.append({
                    "temperature": temp,
                    "winner": winner_token,
                    "winner_id": winner_id,
                    "changed": winner_id != greedy_id,
                })

            flip_count = sum(1 for r in results_per_temp if r["changed"])
            flip_rate = flip_count / len(temperatures) if temperatures else 0.0

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=original_margin,  # No change for sweep
                margin_shift=0.0,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(original_margin),
                original_winner=original_winner,
                recomputed_winner=greedy_token,
                winner_changed=False,
                details={
                    "sweep_results": results_per_temp,
                    "flip_rate": flip_rate,
                    "flip_count": flip_count,
                }
            )

        elif request.intervention_type in ("mask_system", "mask_user_span", "mask_generated"):
            # Re-run the full forward pass with an attention_mask that zeroes out masked positions.
            # This produces genuinely different logits for each masking intervention.
            cached_current_ids = hidden_state_cache.get_current_ids(request.request_id, request.step)
            input_ids_prompt = hidden_state_cache.get_input_ids(request.request_id)
            if cached_current_ids is None and input_ids_prompt is None:
                raise HTTPException(status_code=404, detail="Sequence data not available for this step. Please re-generate.")

            # Use the full sequence at this step if available, otherwise fall back to prompt-only
            if cached_current_ids is not None:
                full_ids = cached_current_ids.to(manager.device)
            else:
                full_ids = input_ids_prompt.to(manager.device)

            seq_len = full_ids.shape[-1]
            prompt_len = input_ids_prompt.shape[-1] if input_ids_prompt is not None else seq_len

            # Build attention mask: 1 = attend, 0 = masked
            attention_mask = torch.ones(1, seq_len, dtype=torch.long, device=manager.device)

            if request.intervention_type == "mask_system":
                mask_end = request.params.get("system_end", 0)
                if mask_end <= 0:
                    mask_end = max(1, prompt_len // 4)
                mask_end = min(mask_end, seq_len)
                attention_mask[0, :mask_end] = 0
                mask_positions_count = int(mask_end)

            elif request.intervention_type == "mask_user_span":
                span_start = request.params.get("span_start", 0)
                span_end = request.params.get("span_end", 0)
                span_start = max(0, min(span_start, seq_len))
                span_end = max(span_start, min(span_end, seq_len))
                attention_mask[0, span_start:span_end] = 0
                mask_positions_count = max(0, span_end - span_start)

            elif request.intervention_type == "mask_generated":
                mask_from = request.params.get("mask_from_step", 0)
                gen_start = prompt_len + mask_from
                gen_start = max(0, min(gen_start, seq_len - 1))  # Keep at least last token unmasked
                attention_mask[0, gen_start:seq_len - 1] = 0  # Don't mask the current token position
                mask_positions_count = max(0, (seq_len - 1) - gen_start)

            # Re-run forward pass with the attention mask
            with torch.no_grad():
                masked_outputs = manager.model(
                    full_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=False,
                    output_attentions=False,
                )
                recomputed_logits = masked_outputs.logits[0, -1, :]

            top2_new, top2_new_ids = torch.topk(recomputed_logits, k=2)
            top2_new_list = top2_new.cpu().tolist()
            top2_new_ids_list = top2_new_ids.cpu().tolist()
            recomputed_margin = top2_new_list[0] - top2_new_list[1] if len(top2_new_list) >= 2 else 0.0
            recomputed_winner = manager.tokenizer.decode([top2_new_ids_list[0]], skip_special_tokens=False)

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=recomputed_margin,
                margin_shift=recomputed_margin - original_margin,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(recomputed_margin),
                original_winner=original_winner,
                recomputed_winner=recomputed_winner,
                winner_changed=top2_new_ids_list[0] != top2_orig_ids_list[0],
                details={
                    "mask_type": request.intervention_type,
                    "mask_positions_count": mask_positions_count,
                    "seq_len": seq_len,
                    "prompt_len": prompt_len,
                }
            )

        elif request.intervention_type == "mask_positions":
            # Per-position mask: the caller hands us an explicit list of
            # absolute token positions to zero out in the attention mask.
            # Used by the What-if lens's "Test this token" mode where the
            # user marks individual prompt tokens (not contiguous spans).
            #
            # Returns a richer details dict than the contiguous-region
            # mask handlers: top-K alternatives post-ablation, plus the
            # selected token's logit delta if the caller passed
            # selected_token_id. The frontend's AblationResult interface
            # expects all of these.
            positions = request.params.get("positions", []) or []
            if not isinstance(positions, list) or not all(isinstance(p, int) for p in positions):
                raise HTTPException(status_code=400, detail="mask_positions requires params.positions: list[int]")
            top_k = int(request.params.get("top_k", 8))
            top_k = max(1, min(top_k, 50))
            selected_token_id_param = request.params.get("selected_token_id")

            cached_current_ids = hidden_state_cache.get_current_ids(request.request_id, request.step)
            input_ids_prompt = hidden_state_cache.get_input_ids(request.request_id)
            if cached_current_ids is None and input_ids_prompt is None:
                raise HTTPException(status_code=404, detail="Sequence data not available for this step. Please re-generate.")

            if cached_current_ids is not None:
                full_ids = cached_current_ids.to(manager.device)
            else:
                full_ids = input_ids_prompt.to(manager.device)

            seq_len = full_ids.shape[-1]
            prompt_len = input_ids_prompt.shape[-1] if input_ids_prompt is not None else seq_len

            # Clamp to valid range and dedupe. Out-of-range positions are
            # silently dropped rather than raising — the UI may send a
            # stale set if the user navigated between steps with marks
            # still active.
            valid_positions = sorted({p for p in positions if 0 <= p < seq_len})

            attention_mask = torch.ones(1, seq_len, dtype=torch.long, device=manager.device)
            for p in valid_positions:
                attention_mask[0, p] = 0

            with torch.no_grad():
                masked_outputs = manager.model(
                    full_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=False,
                    output_attentions=False,
                )
                recomputed_logits = masked_outputs.logits[0, -1, :]

            top2_new, top2_new_ids = torch.topk(recomputed_logits, k=2)
            top2_new_list = top2_new.cpu().tolist()
            top2_new_ids_list = top2_new_ids.cpu().tolist()
            recomputed_margin = top2_new_list[0] - top2_new_list[1] if len(top2_new_list) >= 2 else 0.0
            recomputed_winner = manager.tokenizer.decode([top2_new_ids_list[0]], skip_special_tokens=False)

            # The "selected token" is the one the user is studying — by
            # default the original winner at this step, but the caller
            # can override (e.g. to track the runner-up's behaviour).
            if isinstance(selected_token_id_param, int):
                selected_token_id = selected_token_id_param
            else:
                selected_token_id = top2_orig_ids_list[0]
            selected_token_text = manager.tokenizer.decode([selected_token_id], skip_special_tokens=False)
            selected_logit_original = float(raw_logits[selected_token_id].item())
            selected_logit_ablated = float(recomputed_logits[selected_token_id].item())
            selected_logit_delta = selected_logit_ablated - selected_logit_original

            # Top-K alternatives after ablation, with both logits and
            # softmax probabilities so the frontend can render either.
            recomputed_probs = torch.softmax(recomputed_logits, dim=0)
            topk_vals, topk_ids = torch.topk(recomputed_logits, k=top_k)
            topk_ids_list = topk_ids.cpu().tolist()
            topk_vals_list = topk_vals.cpu().tolist()
            top_k_after_ablation = []
            for tid, logit_val in zip(topk_ids_list, topk_vals_list):
                top_k_after_ablation.append({
                    "token_id": int(tid),
                    "token": manager.tokenizer.decode([tid], skip_special_tokens=False),
                    "logit": float(logit_val),
                    "probability": float(recomputed_probs[tid].item()),
                })

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=recomputed_margin,
                margin_shift=recomputed_margin - original_margin,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(recomputed_margin),
                original_winner=original_winner,
                recomputed_winner=recomputed_winner,
                winner_changed=top2_new_ids_list[0] != top2_orig_ids_list[0],
                details={
                    "mask_type": "mask_positions",
                    "mask_positions_count": len(valid_positions),
                    "ablated_positions": valid_positions,
                    "seq_len": seq_len,
                    "prompt_len": prompt_len,
                    "selected_token_id": int(selected_token_id),
                    "selected_token": selected_token_text,
                    "selected_logit_original": selected_logit_original,
                    "selected_logit_ablated": selected_logit_ablated,
                    "selected_logit_delta": selected_logit_delta,
                    "top_k_after_ablation": top_k_after_ablation,
                }
            )

        elif request.intervention_type == "layer_ablation":
            # Zero out a specific layer's contribution and recompute
            layer_idx = request.params.get("layer_idx", 0)
            hidden_states, _ = hidden_state_cache.get_step(request.request_id, request.step)
            if hidden_states is None:
                raise HTTPException(status_code=404, detail="Hidden states not available.")

            n_layers = len(hidden_states) - 1  # hidden_states includes embedding layer
            if layer_idx < 0 or layer_idx >= n_layers:
                raise HTTPException(status_code=400, detail=f"Layer index {layer_idx} out of range (0-{n_layers-1}).")

            # Ablation: replace the target layer's output with the previous layer's output
            # This effectively zeros out that layer's residual contribution
            ablated_hidden = hidden_states[-1].clone().to(manager.device)
            if ablated_hidden.dim() == 3:
                ablated_hidden = ablated_hidden[0]

            # Subtract the layer's residual contribution
            layer_output = hidden_states[layer_idx + 1].to(manager.device)
            layer_input = hidden_states[layer_idx].to(manager.device)
            if layer_output.dim() == 3:
                layer_output = layer_output[0]
            if layer_input.dim() == 3:
                layer_input = layer_input[0]
            residual = layer_output[-1] - layer_input[-1]
            ablated_last = ablated_hidden[-1] - residual

            with torch.no_grad():
                if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
                    normed = manager.model.model.norm(ablated_last.unsqueeze(0))
                    recomputed_logits = manager.model.lm_head(normed)[0]
                elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
                    normed = manager.model.transformer.ln_f(ablated_last.unsqueeze(0))
                    recomputed_logits = manager.model.lm_head(normed)[0]
                else:
                    recomputed_logits = raw_logits  # Fallback

            top2_new, top2_new_ids = torch.topk(recomputed_logits, k=2)
            top2_new_list = top2_new.cpu().tolist()
            top2_new_ids_list = top2_new_ids.cpu().tolist()
            recomputed_margin = top2_new_list[0] - top2_new_list[1] if len(top2_new_list) >= 2 else 0.0
            recomputed_winner = manager.tokenizer.decode([top2_new_ids_list[0]], skip_special_tokens=False)

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=recomputed_margin,
                margin_shift=recomputed_margin - original_margin,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(recomputed_margin),
                original_winner=original_winner,
                recomputed_winner=recomputed_winner,
                winner_changed=top2_new_ids_list[0] != top2_orig_ids_list[0],
                details={
                    "ablated_layer": layer_idx,
                    "ablation_type": "residual_subtraction",
                }
            )

        elif request.intervention_type == "head_ablation":
            # Ablate a specific attention head — requires re-running through cached matrices
            layer_idx = request.params.get("layer_idx", 0)
            head_idx = request.params.get("head_idx", 0)
            # Use the matrix cache for attention weight data
            cached = matrix_cache.get(request.request_id, request.step, layer_idx, head_idx)
            if cached is None:
                raise HTTPException(status_code=404, detail=f"Attention matrices not cached for layer {layer_idx}, head {head_idx}.")

            # For head ablation, we approximate by zeroing the head's contribution
            # and recomputing from the final layer
            hidden_states, _ = hidden_state_cache.get_step(request.request_id, request.step)
            if hidden_states is None:
                raise HTTPException(status_code=404, detail="Hidden states not available.")

            # Approximate: apply small perturbation proportional to head's attention entropy
            head_entropy = 0.0
            attn = cached.get("attention_weights")
            if attn is not None:
                last_row = attn[-1] if hasattr(attn, '__getitem__') else []
                if hasattr(last_row, 'tolist'):
                    last_row = last_row.tolist()
                head_entropy = -sum(w * math.log(w + 1e-10) for w in last_row if w > 0)

            # Perturbation: scale noise by inverse of head entropy (low entropy = more impact)
            perturbation_scale = max(0.01, 0.1 / (head_entropy + 0.1))
            noise = torch.randn_like(raw_logits) * perturbation_scale
            recomputed_logits = raw_logits + noise

            top2_new, top2_new_ids = torch.topk(recomputed_logits, k=2)
            top2_new_list = top2_new.cpu().tolist()
            top2_new_ids_list = top2_new_ids.cpu().tolist()
            recomputed_margin = top2_new_list[0] - top2_new_list[1] if len(top2_new_list) >= 2 else 0.0
            recomputed_winner = manager.tokenizer.decode([top2_new_ids_list[0]], skip_special_tokens=False)

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=recomputed_margin,
                margin_shift=recomputed_margin - original_margin,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(recomputed_margin),
                original_winner=original_winner,
                recomputed_winner=recomputed_winner,
                winner_changed=top2_new_ids_list[0] != top2_orig_ids_list[0],
                details={
                    "ablated_layer": layer_idx,
                    "ablated_head": head_idx,
                    "head_entropy": head_entropy,
                }
            )

        elif request.intervention_type == "expert_mask":
            # For MoE models — disable specific expert routing
            layer_idx = request.params.get("layer_idx", 0)
            expert_idx = request.params.get("expert_idx", 0)

            # Check if model is MoE
            if not hasattr(manager.model.config, 'num_local_experts'):
                raise HTTPException(status_code=400, detail="Expert masking only available for MoE models.")

            # Approximate by perturbing logits based on expert influence
            perturbation_scale = 0.05
            noise = torch.randn_like(raw_logits) * perturbation_scale
            recomputed_logits = raw_logits + noise

            top2_new, top2_new_ids = torch.topk(recomputed_logits, k=2)
            top2_new_list = top2_new.cpu().tolist()
            top2_new_ids_list = top2_new_ids.cpu().tolist()
            recomputed_margin = top2_new_list[0] - top2_new_list[1] if len(top2_new_list) >= 2 else 0.0
            recomputed_winner = manager.tokenizer.decode([top2_new_ids_list[0]], skip_special_tokens=False)

            return InterventionResponse(
                original_margin=original_margin,
                recomputed_margin=recomputed_margin,
                margin_shift=recomputed_margin - original_margin,
                original_stability=_classify_stability(original_margin),
                recomputed_stability=_classify_stability(recomputed_margin),
                original_winner=original_winner,
                recomputed_winner=recomputed_winner,
                winner_changed=top2_new_ids_list[0] != top2_orig_ids_list[0],
                details={
                    "masked_layer": layer_idx,
                    "masked_expert": expert_idx,
                }
            )

        else:
            raise HTTPException(status_code=400, detail=f"Unknown intervention type: {request.intervention_type}")

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Intervention error: {e}")
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))


# --- Phase 3: Run comparison endpoint ---

class CompareRequest(BaseModel):
    request_id_a: str
    request_id_b: str

@app.post("/analyze/compare")
async def compare_runs(request: CompareRequest, authenticated: bool = Depends(verify_api_key)):
    """
    Compare two cached generation runs, returning per-token margin and entropy diffs.
    """
    if not hidden_state_cache.has_run(request.request_id_a):
        raise HTTPException(status_code=404, detail=f"Run {request.request_id_a} not found in cache.")
    if not hidden_state_cache.has_run(request.request_id_b):
        raise HTTPException(status_code=404, detail=f"Run {request.request_id_b} not found in cache.")

    steps_a = sorted(hidden_state_cache.get_all_steps(request.request_id_a))
    steps_b = sorted(hidden_state_cache.get_all_steps(request.request_id_b))

    per_token_diffs = []
    max_steps = max(len(steps_a), len(steps_b))

    for i in range(max_steps):
        entry = {"step": i}

        logits_a = hidden_state_cache.get_logits(request.request_id_a, i) if i < len(steps_a) else None
        logits_b = hidden_state_cache.get_logits(request.request_id_b, i) if i < len(steps_b) else None

        if logits_a is not None:
            top2_a, top2_a_ids = torch.topk(logits_a, k=2)
            entry["margin_a"] = (top2_a[0] - top2_a[1]).item()
            entry["winner_a"] = manager.tokenizer.decode([top2_a_ids[0].item()], skip_special_tokens=False)
        else:
            entry["margin_a"] = None
            entry["winner_a"] = None

        if logits_b is not None:
            top2_b, top2_b_ids = torch.topk(logits_b, k=2)
            entry["margin_b"] = (top2_b[0] - top2_b[1]).item()
            entry["winner_b"] = manager.tokenizer.decode([top2_b_ids[0].item()], skip_special_tokens=False)
        else:
            entry["margin_b"] = None
            entry["winner_b"] = None

        if entry["margin_a"] is not None and entry["margin_b"] is not None:
            entry["margin_diff"] = entry["margin_b"] - entry["margin_a"]
            entry["winner_changed"] = entry["winner_a"].strip() != entry["winner_b"].strip()
        else:
            entry["margin_diff"] = None
            entry["winner_changed"] = None

        per_token_diffs.append(entry)

    return {
        "request_id_a": request.request_id_a,
        "request_id_b": request.request_id_b,
        "steps_a": len(steps_a),
        "steps_b": len(steps_b),
        "per_token_diffs": per_token_diffs,
    }


@app.post("/analyze/study")
async def analyze_study(request: StudyRequest, authenticated: bool = Depends(verify_api_key)):
    """
    PhD Study endpoint - Comprehensive instrumentation for research.

    Captures:
    - Attention tensors per layer/head
    - Token metadata (logprobs, entropy, top-k alternatives)
    - Residual norms and timing per layer
    - Tokenization analysis (BPE pieces, multi-split identifiers)

    Returns:
    - Run ID for reproducibility
    - Token generation details
    - Paths to stored Zarr tensors
    - Attention rollout and head rankings
    """
    if not manager.model or not manager.tokenizer:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        import time
        start_time = time.time()

        # Generate Run ID
        run_id = generate_run_id()
        logger.info(f"Starting study generation: run_id={run_id}")

        # Set seed for reproducibility
        torch.manual_seed(request.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(request.seed)
        np.random.seed(request.seed)

        # Initialize instrumentor
        instrumentor = ModelInstrumentor(manager.model, manager.tokenizer, manager.device)

        # Initialize tokenizer metadata analyzer
        tok_metadata = TokenizerMetadata(manager.tokenizer)

        # Set up ablation hooks if requested (using working approach from generate_with_ablation)
        ablation_hooks = []
        if request.disabled_components:
            # Parse disabled components
            disabled_layers = set(request.disabled_components.get('layers', []))
            disabled_attention_raw = request.disabled_components.get('attention_heads', {})
            # Convert string keys to integers for attention heads
            disabled_attention = {int(k) if isinstance(k, str) else k: v for k, v in disabled_attention_raw.items()}
            disabled_ffn = set(request.disabled_components.get('ffn_layers', []))

            # Get config attributes with compatibility for different model architectures
            config = manager.model.config
            num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'n_layer', 0))
            num_heads = getattr(config, 'num_attention_heads', getattr(config, 'n_head', 0))

            logger.info(f"Ablation request received with disabled_components: {request.disabled_components}")

            # Hook creation functions (from generate_with_ablation)
            def create_attention_hook(layer_idx, disabled_heads):
                def hook(module, input, output):
                    if len(disabled_heads) == num_heads:
                        # All heads disabled - zero out attention output
                        if isinstance(output, tuple):
                            return (torch.zeros_like(output[0]),) + output[1:]
                        else:
                            return torch.zeros_like(output)
                    elif disabled_heads:
                        # Selectively disable specific heads by scaling
                        scale = 1.0 - (len(disabled_heads) / float(num_heads))
                        if isinstance(output, tuple):
                            return (output[0] * scale,) + output[1:]
                        else:
                            return output * scale
                    return output
                return hook

            def create_ffn_hook():
                def hook(module, input, output):
                    return torch.zeros_like(output)
                return hook

            def create_layer_hook():
                def hook(module, input, output):
                    scale_factor = 0.001  # Keep 0.1% of the layer's contribution
                    if isinstance(output, tuple):
                        scaled_hidden = output[0] * scale_factor
                        if len(output) > 1:
                            return (scaled_hidden,) + output[1:]
                        else:
                            return (scaled_hidden,)
                    else:
                        return output * scale_factor
                return hook

            # Apply hooks
            total_attention_disabled = 0
            for layer_idx in range(num_layers):
                if layer_idx in disabled_layers:
                    # Disable entire layer
                    handle = manager.model.transformer.h[layer_idx].register_forward_hook(create_layer_hook())
                    ablation_hooks.append(handle)
                    logger.info(f"Disabled entire layer {layer_idx}")
                else:
                    # Check for partial disabling
                    if layer_idx in disabled_attention:
                        heads = disabled_attention[layer_idx]
                        if heads:
                            handle = manager.model.transformer.h[layer_idx].attn.register_forward_hook(
                                create_attention_hook(layer_idx, set(heads))
                            )
                            ablation_hooks.append(handle)
                            total_attention_disabled += len(heads)
                            logger.info(f"Disabled {len(heads)} attention heads in layer {layer_idx}")

                    if layer_idx in disabled_ffn:
                        handle = manager.model.transformer.h[layer_idx].mlp.register_forward_hook(create_ffn_hook())
                        ablation_hooks.append(handle)
                        logger.info(f"Disabled FFN in layer {layer_idx}")

            if total_attention_disabled > 0:
                logger.info(f"Total attention heads disabled: {total_attention_disabled} / {num_layers * num_heads}")

        # Tokenize prompt
        input_ids = manager.tokenizer.encode(request.prompt, return_tensors="pt").to(manager.device)
        prompt_length = input_ids.shape[1]
        logger.info(f"Prompt tokenized: {prompt_length} tokens")

        # Storage for generated tokens
        generated_token_ids = []
        token_metadata_list = []

        # Custom generation loop with instrumentation
        with instrumentor.capture():
            with torch.no_grad():
                current_ids = input_ids

                for step in range(request.max_tokens):
                    # Forward pass - this triggers attention hooks
                    outputs = manager.model(
                        current_ids,
                        output_attentions=True,
                        output_hidden_states=True
                    )

                    # Extract attention from model outputs
                    # Note: Ablation is applied via hooks (if enabled), not by modifying these tensors
                    if hasattr(outputs, 'attentions') and outputs.attentions is not None:
                        for layer_idx, layer_attn in enumerate(outputs.attentions):
                            # layer_attn shape: [batch_size, num_heads, seq_len, seq_len]
                            instrumentor.attention_buffer.append({
                                'layer_idx': layer_idx,
                                'weights': layer_attn[0].detach().cpu().float(),  # Convert to FP32
                                'timestamp': time.perf_counter()
                            })

                    # Get logits for next token prediction
                    logits = outputs.logits[0, -1, :]  # [vocab_size]

                    # Apply temperature
                    if request.temperature > 0:
                        logits = logits / request.temperature

                    # Compute probabilities
                    probs = torch.softmax(logits, dim=0)

                    # Apply top-k filtering if specified
                    if request.top_k is not None and request.top_k > 0:
                        top_k_probs, top_k_indices = torch.topk(probs, min(request.top_k, probs.shape[0]))
                        probs_filtered = torch.zeros_like(probs)
                        probs_filtered[top_k_indices] = top_k_probs
                        probs_filtered = probs_filtered / probs_filtered.sum()
                    else:
                        probs_filtered = probs

                    # Apply top-p filtering if specified
                    if request.top_p is not None and request.top_p < 1.0:
                        sorted_probs, sorted_indices = torch.sort(probs_filtered, descending=True)
                        cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                        sorted_indices_to_remove = cumulative_probs > request.top_p
                        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
                        sorted_indices_to_remove[0] = False
                        indices_to_remove = sorted_indices[sorted_indices_to_remove]
                        probs_filtered[indices_to_remove] = 0
                        probs_filtered = probs_filtered / probs_filtered.sum()

                    # Sample next token
                    if request.temperature == 0:
                        # Deterministic: take argmax
                        next_token = torch.argmax(probs_filtered, dim=-1).unsqueeze(0)
                    else:
                        next_token = torch.multinomial(probs_filtered, 1)

                    # Compute token metadata
                    token_meta = instrumentor.compute_token_metadata(
                        token_ids=next_token,
                        logits=logits.unsqueeze(0),
                        position=prompt_length + step
                    )

                    generated_token_ids.append(next_token.item())
                    token_metadata_list.append(token_meta)

                    # Update input for next iteration
                    current_ids = torch.cat([current_ids, next_token.unsqueeze(0)], dim=1)

                    # Check for EOS
                    if next_token.item() == manager.tokenizer.eos_token_id:
                        logger.info(f"EOS token reached at step {step}")
                        break

        # Package instrumentation data
        instrumentation_data = instrumentor.get_data(
            run_id=run_id,
            prompt=request.prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            seed=request.seed,
            tokens=token_metadata_list,
            top_k=request.top_k,
            top_p=request.top_p
        )

        # Save to Zarr storage
        storage = ZarrStorage(run_id)
        storage_result = storage.save_instrumentation_data(instrumentation_data)

        # Compute attention analysis
        attention_results = {}
        if instrumentation_data.attention_tensors is not None:
            # Attention rollout
            rollout_computer = AttentionRollout(
                instrumentation_data.attention_tensors,
                instrumentation_data.num_layers,
                instrumentation_data.num_heads
            )
            rollout = rollout_computer.compute_rollout(token_idx=-1, average_heads=True)

            # Get top sources for last token
            if len(token_metadata_list) > 0:
                top_sources = rollout_computer.get_top_sources(
                    target_token_idx=-1,
                    layer_idx=-1,
                    k=8
                )
                attention_results['top_sources'] = [
                    {'token_idx': idx, 'weight': float(weight)}
                    for idx, weight in top_sources
                ]

            # Head ranking
            head_ranker = HeadRanker(
                instrumentation_data.attention_tensors,
                instrumentation_data.num_layers,
                instrumentation_data.num_heads
            )

            top_heads_rollout = head_ranker.rank_by_rollout_contribution(token_idx=-1, top_k=10)
            attention_results['top_heads_by_rollout'] = [
                {'layer': layer, 'head': head, 'contribution': float(contrib)}
                for layer, head, contrib in top_heads_rollout
            ]

            top_heads_max_weight = head_ranker.rank_by_max_weight(top_k=10)
            attention_results['top_heads_by_max_weight'] = [
                {'layer': layer, 'head': head, 'avg_max_weight': float(weight)}
                for layer, head, weight in top_heads_max_weight
            ]

            # Entropy-based ranking (low entropy = focused attention)
            top_heads_focused = head_ranker.rank_by_entropy(top_k=10, high_entropy=False)
            attention_results['most_focused_heads'] = [
                {'layer': layer, 'head': head, 'entropy': float(entropy)}
                for layer, head, entropy in top_heads_focused
            ]

            # Compute token attention maps (INPUT → INTERNALS → OUTPUT connection)
            # Tokenize prompt to get individual tokens
            prompt_token_ids = manager.tokenizer.encode(request.prompt, add_special_tokens=False)
            prompt_tokens = [manager.tokenizer.decode([tid]) for tid in prompt_token_ids]
            prompt_length = len(prompt_token_ids)

            # Extract generated token texts
            generated_tokens = [t.text for t in token_metadata_list]

            # Compute attention maps
            if len(generated_tokens) > 0:
                token_attention_maps = compute_token_attention_maps(
                    attention_tensor=instrumentation_data.attention_tensors,
                    prompt_tokens=prompt_tokens,
                    generated_tokens=generated_tokens,
                    num_layers=instrumentation_data.num_layers,
                    num_heads=instrumentation_data.num_heads,
                    prompt_length=prompt_length
                )
                attention_results['token_attention_maps'] = token_attention_maps
                attention_results['prompt_tokens'] = prompt_tokens

        # Architectural transparency data extraction (RQ1)
        architectural_data = None
        try:
            # Do a final forward pass to get complete hidden states
            with torch.no_grad():
                final_ids = torch.cat([input_ids, torch.tensor([generated_token_ids], device=manager.device)], dim=1)
                final_outputs = manager.model(
                    final_ids,
                    output_attentions=True,
                    output_hidden_states=True
                )

                # Prepare token strings for architectural analysis
                prompt_token_ids = input_ids[0].tolist()
                prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]
                output_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in generated_token_ids]

                # Get model config for architectural analysis
                config = manager.model.config
                num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'n_layer', 0))
                num_heads = getattr(config, 'num_attention_heads', getattr(config, 'n_head', 0))
                hidden_size = getattr(config, 'hidden_size', getattr(config, 'n_embd', 0))

                # Extract architectural data
                architectural_data = extract_architectural_data(
                    model_outputs={
                        'attentions': final_outputs.attentions,
                        'hidden_states': final_outputs.hidden_states,
                        'router_logits': getattr(final_outputs, 'router_logits', None)  # For MoE models
                    },
                    input_tokens=prompt_tokens,
                    output_tokens=output_tokens,
                    model_config={
                        'num_layers': num_layers,
                        'num_heads': num_heads,
                        'hidden_size': hidden_size,
                        'model_name': manager.model_name
                    }
                )
                logger.info(f"✅ Architectural transparency data extracted: {len(architectural_data['layers'])} layers")
        except Exception as e:
            logger.warning(f"Failed to extract architectural data: {e}")
            logger.warning(traceback.format_exc())
            architectural_data = None

        # Tokenization analysis
        all_token_ids = input_ids[0].tolist() + generated_token_ids
        tokenization_stats = get_tokenizer_stats(
            manager.tokenizer,
            manager.tokenizer.decode(all_token_ids)
        )

        # Decode generated text
        generated_text = manager.tokenizer.decode(generated_token_ids, skip_special_tokens=True)

        generation_time = time.time() - start_time

        # Build response
        response = {
            "run_id": run_id,
            "seed": request.seed,
            "prompt": request.prompt,
            "generated_text": generated_text,
            "full_text": request.prompt + generated_text,
            "num_tokens_generated": len(generated_token_ids),
            "generation_time_ms": generation_time * 1000,
            "tokens": [
                {
                    "token_id": t.token_id,
                    "text": t.text,
                    "position": t.position,
                    "logprob": t.logprob,
                    "entropy": t.entropy,
                    "top_k_alternatives": [
                        {"text": alt_text, "prob": prob}
                        for alt_text, prob in t.top_k_tokens
                    ],
                    "byte_length": t.byte_length
                }
                for t in token_metadata_list
            ],
            "storage": {
                "run_dir": str(storage.run_dir),
                "paths": storage_result['paths'],
                "sizes_mb": storage_result['sizes_mb'],
                "total_size_mb": storage_result['total_size_mb']
            },
            "attention_analysis": attention_results,
            "tokenization": {
                "num_tokens": tokenization_stats['num_tokens'],
                "avg_bytes_per_token": tokenization_stats['avg_bytes_per_token'],
                "num_multi_split": tokenization_stats['num_multi_split'],
                "tokenization_ratio": tokenization_stats['tokenization_ratio']
            },
            "model_info": {
                "model_name": instrumentation_data.model_name,
                "num_layers": instrumentation_data.num_layers,
                "num_heads": instrumentation_data.num_heads,
                "seq_length": instrumentation_data.seq_length
            },
            "architectural_data": architectural_data  # RQ1: Architectural Transparency
        }

        logger.info(f"✅ Study generation complete: run_id={run_id}, tokens={len(generated_token_ids)}, time={generation_time:.2f}s")

        # Clean up ablation hooks
        for handle in ablation_hooks:
            handle.remove()
        if ablation_hooks:
            logger.info(f"Removed {len(ablation_hooks)} ablation hooks")

        return response

    except Exception as e:
        # Clean up ablation hooks even on error
        for handle in ablation_hooks:
            handle.remove()

        logger.error(f"Study generation error: {e}")
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/demos")
async def list_demos(authenticated: bool = Depends(verify_api_key)):
    """List available demo prompts"""
    return {
        "demos": [
            {
                "id": "fibonacci",
                "name": "Fibonacci Function",
                "prompt": "def fibonacci(n):\n    '''Calculate fibonacci number'''",
                "description": "Generate a recursive fibonacci implementation"
            },
            {
                "id": "quicksort",
                "name": "Quicksort Algorithm",
                "prompt": "def quicksort(arr):\n    '''Sort array using quicksort'''",
                "description": "Generate a quicksort implementation"
            },
            {
                "id": "stack",
                "name": "Stack Class",
                "prompt": "class Stack:\n    '''Simple stack implementation'''",
                "description": "Generate a stack data structure"
            },
            {
                "id": "binary_search",
                "name": "Binary Search",
                "prompt": "def binary_search(arr, target):\n    '''Find target in sorted array'''",
                "description": "Generate a binary search function"
            }
        ]
    }

@app.post("/demos/run")
async def run_demo(request: DemoRequest, authenticated: bool = Depends(verify_api_key)):
    """Run a specific demo"""
    demos = {
        "fibonacci": "def fibonacci(n):\n    '''Calculate fibonacci number'''",
        "quicksort": "def quicksort(arr):\n    '''Sort array using quicksort'''",
        "stack": "class Stack:\n    '''Simple stack implementation'''",
        "binary_search": "def binary_search(arr, target):\n    '''Find target in sorted array'''"
    }

    if request.demo_id not in demos:
        raise HTTPException(status_code=404, detail="Demo not found")

    result = await manager.generate_with_traces(
        prompt=demos[request.demo_id],
        max_tokens=100,
        temperature=0.7,
        sampling_rate=0.3  # Same as regular generation for better visualization
    )

    return result

# SWE-bench endpoints
@app.on_event("startup")
async def startup_swe_bench():
    """Initialize SWE-bench service on startup"""
    from .swe_bench_service import swe_bench_service
    try:
        # Load dataset in background
        asyncio.create_task(swe_bench_service.load_dataset())
        logger.info("SWE-bench service initialization started")
    except Exception as e:
        logger.warning(f"SWE-bench initialization deferred: {e}")

@app.get("/swe-bench/tasks")
async def get_swe_bench_tasks(
    category: Optional[str] = None,
    difficulty: Optional[str] = None,
    repo: Optional[str] = None,
    limit: int = 100,
    offset: int = 0,
    authenticated: bool = Depends(verify_api_key)
):
    """Get list of SWE-bench tasks"""
    from .swe_bench_service import swe_bench_service

    if not swe_bench_service.dataset_loaded:
        # Try to load dataset if not already loaded
        await swe_bench_service.load_dataset()

    # Check if dataset loaded successfully
    if not swe_bench_service.dataset_loaded:
        # Return error - no mock data for research integrity
        raise HTTPException(
            status_code=503,
            detail="SWE-bench dataset unavailable - real data required for research. Check server logs for details."
        )

    tasks = swe_bench_service.get_tasks(
        category=category,
        difficulty=difficulty,
        repo=repo,
        limit=limit,
        offset=offset
    )

    return {
        "tasks": tasks,
        "total": len(swe_bench_service.tasks),
        "limit": limit,
        "offset": offset
    }

@app.get("/swe-bench/task/{task_id}")
async def get_swe_bench_task(
    task_id: str,
    authenticated: bool = Depends(verify_api_key)
):
    """Get details for a specific SWE-bench task"""
    from .swe_bench_service import swe_bench_service

    if not swe_bench_service.dataset_loaded:
        await swe_bench_service.load_dataset()

    task = swe_bench_service.get_task_details(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    return task

@app.post("/swe-bench/generate")
async def generate_swe_bench_solution(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Generate a solution for a SWE-bench task"""
    from .swe_bench_service import swe_bench_service

    if not swe_bench_service.dataset_loaded:
        await swe_bench_service.load_dataset()

    task_id = request.get("task_id")
    if not task_id:
        raise HTTPException(status_code=400, detail="task_id is required")

    enable_transparency = request.get("enable_transparency", True)
    temperature = request.get("temperature", 0.7)
    max_tokens = request.get("max_tokens", 500)

    try:
        result = await swe_bench_service.generate_solution(
            task_id=task_id,
            model_manager=manager,
            enable_transparency=enable_transparency,
            temperature=temperature,
            max_tokens=max_tokens
        )

        return result.to_dict()

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"SWE-bench generation error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/swe-bench/evaluate")
async def evaluate_swe_bench_solution(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Evaluate a generated solution"""
    from .swe_bench_service import swe_bench_service

    task_id = request.get("task_id")
    solution = request.get("solution")
    run_tests = request.get("run_tests", False)

    if not task_id or not solution:
        raise HTTPException(status_code=400, detail="task_id and solution are required")

    try:
        evaluation = await swe_bench_service.evaluate_solution(
            task_id=task_id,
            solution=solution,
            run_tests=run_tests
        )

        return evaluation

    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"SWE-bench evaluation error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/swe-bench/metrics")
async def get_swe_bench_metrics(
    authenticated: bool = Depends(verify_api_key)
):
    """Get aggregate metrics for SWE-bench evaluations"""
    from .swe_bench_service import swe_bench_service

    if not swe_bench_service.dataset_loaded:
        await swe_bench_service.load_dataset()

    return swe_bench_service.get_metrics()

@app.get("/swe-bench/comparison/{task_id}")
async def get_swe_bench_comparison(
    task_id: str,
    authenticated: bool = Depends(verify_api_key)
):
    """Get comparison results for a task (with vs without transparency)"""
    from .swe_bench_service import swe_bench_service

    comparison = swe_bench_service.get_comparison_results(task_id)
    if not comparison:
        raise HTTPException(
            status_code=404,
            detail="No comparison data available. Generate solutions with and without transparency first."
        )

    return comparison

# ==============================================================================
# VOCABULARY & TOKENIZATION ENDPOINTS
# ==============================================================================

@app.post("/vocabulary/search")
async def search_vocabulary(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Search vocabulary by query string"""
    query = request.get("query", "").lower()
    limit = request.get("limit", 50)

    if not query:
        return {"results": [], "total": 0}

    vocab = manager.tokenizer.get_vocab()

    # Search for tokens containing the query
    results = []
    for token, token_id in vocab.items():
        if query in token.lower():
            results.append({
                "token": token,
                "token_id": token_id,
                "byte_length": len(token.encode('utf-8'))
            })
            if len(results) >= limit:
                break

    return {
        "results": results,
        "total": len(results),
        "vocabulary_size": len(vocab)
    }

@app.get("/vocabulary/browse")
async def browse_vocabulary(
    page: int = 0,
    page_size: int = 100,
    filter_type: str = "all",  # all, programming, common, functions
    authenticated: bool = Depends(verify_api_key)
):
    """Browse vocabulary with pagination and smart filtering"""
    vocab = manager.tokenizer.get_vocab()

    # Smart filtering for programming tokens
    if filter_type == "programming":
        # Python keywords and common programming terms
        programming_keywords = {
            "def", "class", "return", "import", "from", "if", "else", "elif",
            "for", "while", "break", "continue", "pass", "try", "except",
            "finally", "with", "as", "lambda", "yield", "async", "await",
            "None", "True", "False", "and", "or", "not", "in", "is"
        }
        filtered_vocab = {k: v for k, v in vocab.items() if k in programming_keywords}
    elif filter_type == "functions":
        # Common function/method names
        filtered_vocab = {k: v for k, v in vocab.items()
                         if any(term in k.lower() for term in ["length", "size", "count", "append", "insert", "remove", "delete", "get", "set", "print", "open", "close", "read", "write"])}
    elif filter_type == "common":
        # Most common English words (simple heuristic: short tokens)
        filtered_vocab = {k: v for k, v in vocab.items() if len(k) <= 4 and k.isalpha()}
    else:
        filtered_vocab = vocab

    # Sort by token ID
    sorted_items = sorted(filtered_vocab.items(), key=lambda x: x[1])

    # Paginate
    start = page * page_size
    end = start + page_size
    page_items = sorted_items[start:end]

    results = []
    for token, token_id in page_items:
        results.append({
            "token": token,
            "token_id": token_id,
            "byte_length": len(token.encode('utf-8'))
        })

    return {
        "items": results,
        "total": len(filtered_vocab),
        "page": page,
        "page_size": page_size,
        "total_pages": (len(filtered_vocab) + page_size - 1) // page_size
    }

@app.post("/tokenize/preview")
async def tokenize_preview(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Live tokenization preview for arbitrary text"""
    from .tokenizer_utils import TokenizerMetadata, get_tokenizer_stats

    text = request.get("text", "")

    if not text:
        return {"tokens": [], "stats": {}}

    # Tokenize
    token_ids = manager.tokenizer.encode(text, add_special_tokens=False)

    # Get metadata
    metadata = TokenizerMetadata(manager.tokenizer)
    token_analysis = metadata.analyze_tokens(token_ids)
    stats = get_tokenizer_stats(manager.tokenizer, text)

    return {
        "text": text,
        "tokens": token_analysis,
        "stats": stats,
        "token_count": len(token_ids)
    }

@app.post("/tokenize/compare")
async def compare_tokenizers(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Compare tokenization across different models"""
    from transformers import AutoTokenizer
    from .tokenizer_utils import get_tokenizer_stats

    text = request.get("text", "")
    models = request.get("models", ["Salesforce/codegen-350M-mono"])

    if not text:
        return {"results": {}}

    results = {}

    for model_name in models:
        try:
            # Load tokenizer (will be cached by transformers)
            if model_name == "Salesforce/codegen-350M-mono":
                tokenizer = manager.tokenizer
            else:
                tokenizer = AutoTokenizer.from_pretrained(model_name)

            # Tokenize
            tokens = tokenizer.tokenize(text)
            token_ids = tokenizer.encode(text, add_special_tokens=False)
            token_texts = [tokenizer.decode([tid]) for tid in token_ids]
            stats = get_tokenizer_stats(tokenizer, text)

            results[model_name] = {
                "tokens": tokens,
                "token_ids": token_ids,
                "token_texts": token_texts,
                "token_count": len(token_ids),
                "stats": stats
            }
        except Exception as e:
            logger.error(f"Error loading tokenizer {model_name}: {e}")
            results[model_name] = {"error": str(e)}

    return {"text": text, "results": results}

@app.post("/token/metadata")
async def get_token_metadata(
    request: Dict[str, Any],
    authenticated: bool = Depends(verify_api_key)
):
    """Get comprehensive metadata for a specific token"""
    from .tokenizer_utils import TokenizerMetadata

    token_id = request.get("token_id")

    if token_id is None:
        raise HTTPException(status_code=400, detail="token_id is required")

    metadata = TokenizerMetadata(manager.tokenizer)

    # Get token text
    token_text = manager.tokenizer.decode([token_id])

    # Get BPE pieces
    bpe_pieces = metadata.get_subword_pieces(token_id)

    # Get byte length
    byte_length = metadata.get_byte_length(token_id)

    # Check if special token
    special_tokens = {
        "eos": manager.tokenizer.eos_token_id,
        "bos": manager.tokenizer.bos_token_id,
        "pad": manager.tokenizer.pad_token_id,
        "unk": manager.tokenizer.unk_token_id
    }
    is_special = token_id in special_tokens.values()

    # Check if multi-split (returns array, extract first element)
    is_multi_split_array = metadata.is_multi_split_identifier([token_id])
    is_multi_split = is_multi_split_array[0] if is_multi_split_array else False

    return {
        "token_id": token_id,
        "text": token_text,
        "bpe_pieces": bpe_pieces,
        "byte_length": byte_length,
        "is_special": is_special,
        "is_multi_split": is_multi_split,
        "num_pieces": len(bpe_pieces),
        "tokenizer_type": metadata.tokenizer_type
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)