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.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import asyncio
import json
import os
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
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__)

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", "codegen-350m")
        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
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=self.dtype,
                low_cpu_mem_usage=True,
                trust_remote_code=True
            ).to(self.device)
            
            # Load tokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self.tokenizer.pad_token = self.tokenizer.eos_token

            # 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")

        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().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().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"""
    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),
        "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 all available models this backend can serve.

    Returns model metadata including availability based on current hardware.
    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"]
    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():
        # Determine if model is available on current hardware
        is_available = True
        if config["requires_gpu"] and not has_gpu:
            is_available = False
        elif has_gpu 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"],
            "available": is_available
        })

    return {"models": models}


@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"""
    from .model_config import list_all_models, SUPPORTED_MODELS

    # Get current device type
    device_type = "cpu"
    if torch.cuda.is_available():
        device_type = "cuda"
    elif torch.backends.mps.is_available():
        device_type = "mps"

    all_models = list_all_models()

    # Filter models based on hardware capabilities
    available_models = []
    for model in all_models:
        model_config = SUPPORTED_MODELS.get(model['id'])

        # Check if model requires GPU but we're on CPU
        if model_config and model_config['requires_gpu'] and device_type == "cpu":
            # Skip GPU-only models when on CPU
            continue

        # Model is available on this hardware
        model['available'] = True
        model['is_current'] = (model['id'] == manager.model_id)
        available_models.append(model)

    return {"models": available_models}

@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"
        )

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

        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()

        # Get parameters
        prompt = request.get("prompt", "def quicksort(arr):")
        max_tokens = request.get("max_tokens", 8)
        temperature = request.get("temperature", 0.7)

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

        # Tokenize and prepare
        inputs = manager.tokenizer(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 = []

        def make_qkv_hook(layer_idx):
            def hook(module, input, output):
                try:
                    # output shape: [batch, seq_len, 3 * hidden_size]
                    # Split into Q, K, V
                    if output.dim() != 3:
                        return  # Skip if unexpected shape
                    batch_size, seq_len, hidden = output.shape
                    expected_hidden = 3 * n_heads * head_dim
                    if hidden != expected_hidden:
                        return  # Skip if dimensions don't match QKV format
                    qkv = output.reshape(batch_size, seq_len, 3, n_heads, head_dim)
                    # Separate Q, K, V: [batch, seq_len, n_heads, head_dim]
                    q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
                    qkv_captures[layer_idx] = {
                        'q': q[0].detach().cpu(),  # Remove batch dim
                        'k': k[0].detach().cpu(),
                        'v': v[0].detach().cpu()
                    }
                except Exception:
                    # Silently skip QKV capture if it fails - it's optional data
                    pass
            return hook

        # Register hooks on all qkv_proj modules (if available)
        # This is model-specific - CodeGen uses different architecture
        try:
            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_qkv_hook(layer_idx))
                        hooks.append(hook)
                    elif hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
                        # GPT-2 style attention
                        hook = layer.attn.c_attn.register_forward_hook(make_qkv_hook(layer_idx))
                        hooks.append(hook)
        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
                )

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

                # Apply temperature and sample
                if temperature > 0:
                    logits = logits / temperature
                probs = torch.softmax(logits, dim=0)

                if temperature == 0:
                    next_token_id = torch.argmax(probs, dim=-1).item()
                else:
                    next_token_id = torch.multinomial(probs, 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-k token alternatives with probabilities
                import math
                top_k = 5  # Get top 5 alternatives
                top_probs, top_indices = torch.topk(probs, k=min(top_k, len(probs)))
                alternatives = []
                for prob, idx in zip(top_probs.tolist(), top_indices.tolist()):
                    token_text = manager.tokenizer.decode([idx], skip_special_tokens=False)
                    alternatives.append({
                        "token": token_text,
                        "token_id": idx,
                        "probability": prob,
                        "log_probability": math.log(prob) if prob > 0 else float('-inf')
                    })
                token_alternatives_by_step.append({
                    "step": step,
                    "selected_token": next_token_text,
                    "selected_token_id": next_token_id,
                    "alternatives": alternatives
                })

                # 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()

                        # 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

                        # Classify pattern
                        pattern_type = None
                        confidence = 0.0

                        # Induction pattern: high attention to previous similar tokens
                        if step > 0 and max_weight > 0.8:
                            pattern_type = "induction"
                            confidence = max_weight
                        # Positional pattern: attention focused on nearby tokens
                        elif entropy < 1.0:
                            pattern_type = "positional"
                            confidence = 1.0 - entropy
                        # Semantic pattern: broader attention with moderate entropy
                        elif 1.0 <= entropy < 2.5:
                            pattern_type = "semantic"
                            confidence = min(1.0, entropy / 2.5)
                        # Previous token pattern: sharp focus on immediate predecessor
                        elif max_weight > 0.9 and head_weights[-2].item() > 0.85:
                            pattern_type = "previous_token"
                            confidence = head_weights[-2].item()

                        # 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]
                        attention_matrix = layer_attn[head_idx].cpu().numpy().tolist()

                        # 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]
                            q_matrix = qkv_captures[layer_idx]['q'][:, head_idx, :].numpy().tolist()
                            k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].numpy().tolist()
                            v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].numpy().tolist()

                        critical_heads.append({
                            "head_idx": head_idx,
                            "entropy": entropy,
                            "max_weight": max_weight,
                            "attention_weights": attention_matrix,  # Full attention matrix for spreadsheet
                            "q_matrix": q_matrix,  # [seq_len, head_dim]
                            "k_matrix": k_matrix,
                            "v_matrix": v_matrix,
                            "pattern": {
                                "type": pattern_type,
                                "confidence": confidence
                            } if pattern_type else None
                        })

                    # 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)

                    # Detect layer-level pattern (percentage-based for any layer count)
                    layer_pattern = None
                    layer_fraction = (layer_idx + 1) / n_layers  # 1-indexed fraction
                    if layer_idx == 0:
                        layer_pattern = {"type": "positional", "confidence": 0.78}
                    elif layer_fraction <= 0.25 and step > 0:
                        layer_pattern = {"type": "previous_token", "confidence": 0.65}
                    elif layer_fraction <= 0.75:
                        layer_pattern = {"type": "induction", "confidence": 0.87}
                    else:
                        layer_pattern = {"type": "semantic", "confidence": 0.92}

                    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

        # 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

        # Build response
        response = {
            "prompt": prompt,
            "promptTokens": [{"text": t, "idx": tid, "bytes": len(t.encode('utf-8')), "type": "prompt"}
                           for tid, t in zip(prompt_token_ids, prompt_tokens)],
            "generatedTokens": [{"text": t, "idx": tid, "bytes": len(t.encode('utf-8')), "type": "generated"}
                              for tid, t in zip(generated_token_ids, generated_tokens)],
            "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": qkv_by_layer_head,
            "modelInfo": {
                "numLayers": n_layers,
                "numHeads": n_heads,
                "modelDimension": d_model,
                "headDim": head_dim,
                "vocabSize": manager.model.config.vocab_size
            },
            "generationTime": generation_time,
            "numTokensGenerated": len(generated_tokens)
        }

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

        return 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))

@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

    # DEBUG LOGGING
    print(f"\n{'='*60}")
    print(f"TOKEN METADATA DEBUG - Token ID: {token_id}")
    print(f"{'='*60}")
    print(f"Token Text: {repr(token_text)}")
    print(f"BPE Pieces: {bpe_pieces}")
    print(f"Num Pieces: {len(bpe_pieces)}")
    print(f"Byte Length: {byte_length}")
    print(f"Is Special: {is_special}")
    print(f"Multi-split Array: {is_multi_split_array}")
    print(f"Multi-split Boolean: {is_multi_split} (type: {type(is_multi_split).__name__})")
    print(f"Tokenizer Type: {metadata.tokenizer_type}")
    print(f"{'='*60}\n")

    result = {
        "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
    }

    print(f"RESPONSE: {result}\n")

    return result

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