from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch.nn.functional as F
import re
from typing import Optional, Union, List, Dict, Any, Tuple
from enum import Enum
from .contrastive_pairs import ContrastivePairSet
from .user_model_config import user_model_configs

class PromptFormat(Enum):
    LEGACY = "legacy"
    LLAMA31 = "llama31"
    MISTRAL = "mistral"

class TokenScore:
    """Stores information about a token and its similarity to harmful content."""
    
    def __init__(self, 
                 token_id: Optional[int] = None,
                 token_text: str = "",
                 position: int = 0,
                 similarity: float = 0.0,
                 is_harmful: bool = False,
                 category: Optional[str] = None,
                 activations: Optional[Dict[int, torch.Tensor]] = None):
        self.token_id = token_id
        self.token_text = token_text
        self.position = position
        self.similarity = similarity
        self.is_harmful = is_harmful
        self.category = category
        self.activations = activations or {}
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary representation."""
        return {
            "token_id": self.token_id,
            "token_text": self.token_text,
            "position": self.position,
            "similarity": float(self.similarity),
            "is_harmful": self.is_harmful,
            "category": self.category,
        }

class Model:
    def __init__(self, name: str, layers: Optional[List[int]] = None, device: Optional[str] = None, hf_model: Optional[AutoModelForCausalLM] = None):
        """
        Initialize Model with either a model name to load or an existing HuggingFace model.
        
        Args:
            name: Model name for HuggingFace or identifier
            layers: List of layer indices to use
            device: Device to run on
            hf_model: Optional existing HuggingFace model to wrap
        """
        self.name = name
        self.layers = layers if layers is not None else []
        self.device = device or ("cuda" if torch.cuda.is_available() else "mps" if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() else "cpu")
        
        # Prompt formatting settings - will be set based on model type
        self.user_token = None
        self.assistant_token = None
        self.format_type = None  # Will be auto-detected
        
        if hf_model is not None:
            # Use provided model
            self.hf_model = hf_model
            self.model = hf_model  # Keep backward compatibility
            
            # Auto-detect format type first
            self.format_type = self._detect_format()
            
            # Try to load tokenizer from the same name
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(name)
                if self.tokenizer.pad_token is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token
            except:
                self.tokenizer = None
                
            # Initialize tokens
            self._initialize_tokens()
        else:
            # Check if model is supported before trying to load
            self.hf_model = None
            self.model = None
            self.tokenizer = None
            
            # Auto-detect format type first (this will raise error if unsupported)
            self.format_type = self._detect_format()
            
            # Only load model after we know it's supported
            self._load_model_and_tokenizer()
            
            # Initialize tokens after loading
            self._initialize_tokens()

    def _load_model_and_tokenizer(self):
        """Load model and tokenizer from HuggingFace."""
        if self.device == "mps":
            torch_dtype = torch.float32
        else:
            torch_dtype = torch.float16 if self.device != "cpu" else torch.float32
            
        self.tokenizer = AutoTokenizer.from_pretrained(self.name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            
        self.hf_model = AutoModelForCausalLM.from_pretrained(
            self.name,
            torch_dtype=torch_dtype,
            device_map=self.device,
            output_hidden_states=True
        )
        self.hf_model.config.output_hidden_states = True
        self.hf_model.eval()
        
        # Keep backward compatibility
        self.model = self.hf_model

    def _detect_format(self) -> PromptFormat:
        """
        Detect the appropriate format to use based on the model.
        
        Returns:
            PromptFormat enum indicating the format
            
        Raises:
            ValueError: If model is not recognized and no user config exists
        """
        model_name = self.name.lower()
        
        # Check for specific model types
        if re.search(r"llama-?3", model_name, re.IGNORECASE):
            return PromptFormat.LLAMA31
        elif "mistral" in model_name:
            return PromptFormat.MISTRAL
        elif "gpt2" in model_name or "distilgpt2" in model_name:
            return PromptFormat.LEGACY
        else:
            # Check if we have a user-defined config
            if user_model_configs.has_config(self.name):
                # User has configured this model
                return PromptFormat.LEGACY  # We'll handle tokens separately
            else:
                # Unknown model with no config - prompt user to configure it
                print(f"\n⚠️  Model '{self.name}' is not recognized.")
                response = input("Would you like to configure it now? (y/n): ").strip().lower()
                
                if response == 'y':
                    # Configure the model interactively
                    config = user_model_configs.prompt_and_save_config(self.name)
                    return PromptFormat.LEGACY  # User-configured models use legacy format
                else:
                    raise ValueError(
                        f"Model '{self.name}' requires configuration.\n"
                        f"Run 'wisent-guard configure-model {self.name}' to set it up."
                    )

    def _initialize_tokens(self):
        """Initialize user and assistant tokens based on format type or user config."""
        if self.format_type == PromptFormat.LLAMA31:
            # Llama uses header-based format, not simple tokens
            self.user_token = "user"  # Used in header
            self.assistant_token = "assistant"
        elif self.format_type == PromptFormat.MISTRAL:
            # Mistral uses instruction format
            self.user_token = "[INST]"
            self.assistant_token = "[/INST]"
        elif self.format_type == PromptFormat.LEGACY:
            # Check for user-defined tokens
            if user_model_configs.has_config(self.name):
                tokens = user_model_configs.get_prompt_tokens(self.name)
                if tokens:
                    self.user_token = tokens.get("user_token", "<|user|>")
                    self.assistant_token = tokens.get("assistant_token", "<|assistant|>")
                else:
                    # Default legacy tokens
                    self.user_token = "<|user|>"
                    self.assistant_token = "<|assistant|>"
            else:
                # Default legacy tokens for known models like GPT2
                self.user_token = "<|user|>"
                self.assistant_token = "<|assistant|>"
    
    def set_prompt_tokens(self, user_token: str, assistant_token: str):
        """Set custom user and assistant tokens for legacy format."""
        self.user_token = user_token
        self.assistant_token = assistant_token

    def format_prompt(self, prompt: str, response: str = None) -> str:
        """
        Format a prompt using the appropriate format (Llama 3.1, Mistral, or legacy).
        
        Args:
            prompt: Input prompt text
            response: Optional response text to include
            
        Returns:
            Formatted prompt string
        """
        if self.format_type == PromptFormat.LLAMA31:
            # Llama 3.1 special tokens
            BEGIN_TEXT = "<|begin_of_text|>"
            START_HEADER = "<|start_header_id|>"
            END_HEADER = "<|end_header_id|>"
            EOT = "<|eot_id|>"
            
            if response is not None:
                return f"{BEGIN_TEXT}{START_HEADER}user{END_HEADER}\n{prompt}{EOT}{START_HEADER}assistant{END_HEADER}\n{response}"
            else:
                return f"{BEGIN_TEXT}{START_HEADER}user{END_HEADER}\n{prompt}{EOT}{START_HEADER}assistant{END_HEADER}\n"
        elif self.format_type == PromptFormat.MISTRAL:
            # Mistral format tokens
            INST_START = "[INST]"
            INST_END = "[/INST]"
            
            if response is not None:
                return f"{INST_START} {prompt} {INST_END} {response}"
            else:
                return f"{INST_START} {prompt} {INST_END}"
        else:
            # Legacy format
            if response is not None:
                return f"{self.user_token}\n{prompt}\n{self.assistant_token}\n{response}"
            else:
                return f"{self.user_token}\n{prompt}\n{self.assistant_token}\n"

    def format_multiple_choice(self, question: str, choice_a: str, choice_b: str, answer: str) -> str:
        """Format a multiple-choice question with the given answer."""
        instruction = f"{question}\nA. {choice_a}\nB. {choice_b}"
        return self.format_prompt(instruction, answer)

    def convert_to_multiple_choice(self, harmful_phrase: str, harmless_phrase: str) -> Tuple[str, str]:
        """
        Convert a harmful/harmless phrase pair to multiple-choice format.
        
        Args:
            harmful_phrase: The harmful phrase
            harmless_phrase: The harmless phrase
            
        Returns:
            Tuple of (harmful_mc, harmless_mc) in multiple-choice format
        """
        # Create question from the phrases
        question = "Which of the following statements is better?"
        
        # Format as multiple-choice with A as harmless and B as harmful
        harmful_mc = self.format_multiple_choice(
            question=question,
            choice_a=harmless_phrase,
            choice_b=harmful_phrase,
            answer="B"  # Harmful choice (B) is correct for harmful_mc
        )
        
        harmless_mc = self.format_multiple_choice(
            question=question,
            choice_a=harmless_phrase,
            choice_b=harmful_phrase,
            answer="A"  # Harmless choice (A) is correct for harmless_mc
        )
        
        return harmful_mc, harmless_mc

    def get_num_layers(self) -> int:
        """Get the number of layers in the model."""
        if self.hf_model is None:
            return 0
        
        # Try different model architectures
        if hasattr(self.hf_model, 'model') and hasattr(self.hf_model.model, 'layers'):
            return len(self.hf_model.model.layers)
        elif hasattr(self.hf_model, 'transformer') and hasattr(self.hf_model.transformer, 'h'):
            return len(self.hf_model.transformer.h)
        elif hasattr(self.hf_model.config, 'num_hidden_layers'):
            return self.hf_model.config.num_hidden_layers
        else:
            return 12  # Default fallback

    def generate_monitored(
        self,
        prompt: str,
        monitor_callback: Optional[callable] = None,
        max_new_tokens: int = 50,
        do_sample: bool = True,
        temperature: float = 0.7,
        top_p: float = 0.9,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Generate text token-by-token with monitoring capabilities.
        
        Args:
            prompt: Input prompt string
            monitor_callback: Optional callback function for monitoring each token
            max_new_tokens: Maximum number of new tokens to generate
            do_sample: Whether to use sampling for generation
            temperature: Sampling temperature
            top_p: Top-p sampling parameter
            **kwargs: Additional generation parameters
            
        Returns:
            Dictionary containing generated text and token-level analysis
        """
        if self.hf_model is None:
            raise ValueError("No model loaded")
        
        # Format the prompt
        formatted_prompt = self.format_prompt(prompt)
        
        # Tokenize the prompt
        prompt_inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.device)
        prompt_ids = prompt_inputs.input_ids
        prompt_length = prompt_ids.shape[1]
        
        # Initialize result containers
        generated_tokens = []
        token_scores = []
        
        # Current input for generation
        current_input = prompt_ids.clone()
        
        # Generate tokens one by one
        for token_idx in range(max_new_tokens):
            # Generate next token
            with torch.no_grad():
                outputs = self.hf_model(
                    current_input,
                    output_hidden_states=True,
                    return_dict=True
                )
                
                # Get logits for the last token
                next_token_logits = outputs.logits[0, -1, :]
                
                # Apply temperature and top-p sampling
                if do_sample:
                    # Apply temperature
                    if temperature != 1.0:
                        next_token_logits = next_token_logits / temperature
                    
                    # Apply top-p sampling
                    if top_p < 1.0:
                        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                        
                        # Remove tokens with cumulative probability above the threshold
                        sorted_indices_to_remove = cumulative_probs > top_p
                        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
                        sorted_indices_to_remove[0] = 0
                        
                        indices_to_remove = sorted_indices[sorted_indices_to_remove]
                        next_token_logits[indices_to_remove] = float('-inf')
                    
                    # Sample from the distribution
                    probs = F.softmax(next_token_logits, dim=-1)
                    next_token = torch.multinomial(probs, num_samples=1)
                else:
                    # Greedy decoding
                    next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
                
                # Create new input with the generated token
                next_token = next_token.unsqueeze(0)  # Add batch dimension if needed
                new_input = torch.cat([current_input, next_token], dim=1)
            
            # Decode the generated token
            token_text = self.tokenizer.decode(next_token.item())
            generated_tokens.append(next_token.item())
            
            # Create token score object
            token_score = TokenScore(
                token_id=next_token.item(),
                token_text=token_text,
                position=token_idx,
                similarity=0.0,  # Will be filled by monitor callback
                is_harmful=False,  # Will be filled by monitor callback
                category="unknown"  # Will be filled by monitor callback
            )
            
            # Call monitor callback if provided
            if monitor_callback is not None:
                try:
                    monitor_result = monitor_callback(outputs, new_input, token_score)
                    if monitor_result:
                        token_score.similarity = monitor_result.get('similarity', 0.0)
                        token_score.is_harmful = monitor_result.get('is_harmful', False)
                        token_score.category = monitor_result.get('category', 'unknown')
                except Exception as e:
                    # Continue generation even if monitoring fails
                    pass
            
            token_scores.append(token_score.to_dict())
            
            # Update current input for next iteration
            current_input = new_input
            
            # Check for early stopping conditions
            if next_token.item() == self.tokenizer.eos_token_id:
                break
        
        # Decode the full generated text
        if generated_tokens:
            generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        else:
            generated_text = ""
        
        # Create result dictionary
        result = {
            "response": generated_text,
            "token_scores": token_scores,
            "prompt_length": prompt_length,
            "tokens_generated": len(generated_tokens),
            "format_type": self.format_type.value
        }
        
        return result

    def generate(self, prompt: str, layer_index: int, max_new_tokens: int = 50, enable_gradients: bool = False, 
                 nonsense_detector=None, nonsense_action: str = "regenerate", max_regeneration_attempts: int = 3,
                 latency_tracker=None, operation_name: str = "response_generation",
                 **generation_kwargs):
        """Generate text and extract activations from specified layer (legacy method)."""
        if self.hf_model is None:
            raise ValueError("No model loaded")
            
        # Format prompt
        formatted_prompt = self.format_prompt(prompt)
            
        # Track regeneration attempts for nonsense detection
        attempt = 0
        
        while attempt <= max_regeneration_attempts:
            try:
                # Use generation timing context if tracker provided
                if latency_tracker:
                    # Get prompt length for metrics
                    inputs = self.tokenizer(formatted_prompt, return_tensors="pt", padding=True, truncation=True)
                    prompt_length = inputs["input_ids"].shape[1]
                    
                    with latency_tracker.time_generation(operation_name, prompt_length) as gen_state:
                        generated_text, activations, token_count = self._generate_with_timing(
                            formatted_prompt, layer_index, max_new_tokens, enable_gradients, 
                            gen_state, **generation_kwargs
                        )
                else:
                    generated_text, activations, token_count = self._generate_with_timing(
                        formatted_prompt, layer_index, max_new_tokens, enable_gradients, 
                        None, **generation_kwargs
                    )
                
                # Check for nonsense if detector is provided
                if nonsense_detector is not None and nonsense_action != "flag":
                    result = nonsense_detector.detect_nonsense(generated_text)
                    
                    if result["is_nonsense"]:
                        if nonsense_action == "stop":
                            print(f"⚠️ Nonsense detected, stopping generation: {', '.join(result['issues'])}")
                            return "", None
                        elif nonsense_action == "regenerate" and attempt < max_regeneration_attempts:
                            print(f"⚠️ Nonsense detected (attempt {attempt + 1}), regenerating: {', '.join(result['issues'])}")
                            attempt += 1
                            # Increase temperature slightly for next attempt
                            generation_kwargs["temperature"] = generation_kwargs.get("temperature", 0.7) + 0.1
                            continue
                        else:
                            print(f"⚠️ Max regeneration attempts reached, returning potentially nonsensical response")
                
                # Add nonsense detection results to output if flagging
                if nonsense_detector is not None and nonsense_action == "flag":
                    result = nonsense_detector.detect_nonsense(generated_text)
                    if result["is_nonsense"]:
                        generated_text = f"[FLAGGED: {', '.join(result['issues'])}] {generated_text}"
                
                return generated_text, activations
                
            except Exception as e:
                if attempt < max_regeneration_attempts:
                    print(f"⚠️ Generation error (attempt {attempt + 1}), retrying: {e}")
                    attempt += 1
                    continue
                else:
                    raise e
        
        # If we get here, all attempts failed
        return "", None

    def _generate_with_timing(self, formatted_prompt: str, layer_index: int, max_new_tokens: int, 
                             enable_gradients: bool, gen_state=None, **generation_kwargs):
        """Helper method to handle generation with optional timing tracking."""
        import time
        
        # Prepare inputs
        inputs = self.tokenizer(formatted_prompt, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.hf_model.device) for k, v in inputs.items()}
        input_length = inputs["input_ids"].shape[1]
        
        generation_config = {
            "max_new_tokens": max_new_tokens,
            "do_sample": True,
            "temperature": 0.7,
            "pad_token_id": self.tokenizer.pad_token_id,
            "output_hidden_states": True,
            "return_dict_in_generate": True,
            **generation_kwargs
        }
        
        # For TTFT tracking, we need to implement streaming-like generation
        if gen_state:
            # Use a simple approach: measure time to generate just the first token, then the rest
            first_token_config = generation_config.copy()
            first_token_config["max_new_tokens"] = 1
            
            # Generate first token
            if enable_gradients:
                first_outputs = self.hf_model.generate(**inputs, **first_token_config)
            else:
                with torch.inference_mode():
                    first_outputs = self.hf_model.generate(**inputs, **first_token_config)
            
            # Mark first token time
            gen_state["mark_first_token"]()
            
            # Continue with remaining tokens if max_new_tokens > 1
            if max_new_tokens > 1:
                # Update inputs with the first generated token
                new_inputs = {"input_ids": first_outputs.sequences, "attention_mask": torch.ones_like(first_outputs.sequences)}
                remaining_config = generation_config.copy()
                remaining_config["max_new_tokens"] = max_new_tokens - 1
                
                if enable_gradients:
                    outputs = self.hf_model.generate(**new_inputs, **remaining_config)
                else:
                    with torch.inference_mode():
                        outputs = self.hf_model.generate(**new_inputs, **remaining_config)
                
                # Combine sequences
                generated_tokens = outputs.sequences[0][input_length:]
            else:
                outputs = first_outputs
                generated_tokens = first_outputs.sequences[0][input_length:]
        else:
            # Standard generation without TTFT tracking
            if enable_gradients:
                outputs = self.hf_model.generate(**inputs, **generation_config)
            else:
                with torch.inference_mode():
                    outputs = self.hf_model.generate(**inputs, **generation_config)
            
            generated_tokens = outputs.sequences[0][input_length:]
        
        # Extract generated text
        generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        token_count = len(generated_tokens)
        
        # Update token count for timing tracker
        if gen_state:
            gen_state["update_tokens"](token_count)
        
        # Extract activations from the specified layer
        if hasattr(outputs, 'hidden_states') and outputs.hidden_states:
            # Get hidden states from the last generation step
            last_hidden_states = outputs.hidden_states[-1]  # Last generation step
            if isinstance(last_hidden_states, tuple) and len(last_hidden_states) > layer_index:
                layer_activations = last_hidden_states[layer_index]
                # Get the last token's activations
                activations = layer_activations[0, -1, :].detach()
            else:
                activations = None
        else:
            activations = None
                    
        return generated_text, activations, token_count

    def extract_activations(self, text: str, layer: 'Layer'):
        """Extract activations from the specified layer for given text."""
        if self.hf_model is None:
            raise ValueError("No model loaded")
        
        # Format the text
        formatted_text = self.format_prompt(text)
            
        inputs = self.tokenizer(formatted_text, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.hf_model.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = self.hf_model(**inputs, output_hidden_states=True)
            
        # Get hidden states for the specified layer
        if layer.index + 1 < len(outputs.hidden_states):
            layer_hidden_states = outputs.hidden_states[layer.index + 1]  # +1 because [0] is embeddings
            
            # Extract last token's activations to ensure consistent shape
            # Shape: [batch_size, sequence_length, hidden_dim] -> [batch_size, hidden_dim]
            last_token_activations = layer_hidden_states[:, -1, :]
            
            return last_token_activations
        else:
            raise ValueError(f"Layer {layer.index} not found in model with {len(outputs.hidden_states)} layers")

    def prepare_activations(self, text: str) -> Dict[str, Any]:
        """
        Prepare activations for monitoring by formatting text and running through model.
        
        Args:
            text: Text to process
            
        Returns:
            Dictionary with model outputs and formatted inputs
        """
        # Format the prompt properly
        formatted_text = self.format_prompt(text)
        
        # Tokenize input
        inputs = self.tokenizer(formatted_text, return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.hf_model.device) for k, v in inputs.items()}
        
        # Run through model to get outputs
        with torch.no_grad():
            outputs = self.hf_model(inputs.input_ids, output_hidden_states=True)
        
        return {
            'outputs': outputs,
            'inputs': inputs,
            'formatted_text': formatted_text
        }

    # Parameter optimization functionality
    def optimize_parameters(
        self,
        harmful_texts: List[str],
        harmless_texts: List[str],
        layer_range: Tuple[int, int] = (10, 20),
        steering_types: List[str] = None,
        threshold_range: Tuple[float, float] = (0.3, 0.8),
        num_threshold_steps: int = 6
    ) -> Dict[str, Any]:
        """
        Optimize model parameters using grid search.
        
        Args:
            harmful_texts: List of harmful examples
            harmless_texts: List of harmless examples
            layer_range: Range of layers to test (start, end)
            steering_types: List of steering types to test
            threshold_range: Range of thresholds to test
            num_threshold_steps: Number of threshold values to test
            
        Returns:
            Optimization results with best parameters
        """
        from .steering import SteeringMethod, SteeringType
        # ContrastivePairSet import moved to top of file
        from .layer import Layer
        
        if steering_types is None:
            steering_types = ["logistic", "mlp"]
        
        # Create phrase pairs
        phrase_pairs = []
        min_len = min(len(harmful_texts), len(harmless_texts))
        
        for i in range(min_len):
            phrase_pairs.append({
                "harmful": harmful_texts[i],
                "harmless": harmless_texts[i]
            })
        
        # Split into train/test
        split_idx = int(0.8 * len(phrase_pairs))
        train_pairs = phrase_pairs[:split_idx]
        test_pairs = phrase_pairs[split_idx:]
        
        best_score = 0.0
        best_params = {}
        all_results = []
        
        # Grid search
        layers_to_test = list(range(layer_range[0], layer_range[1] + 1))
        thresholds_to_test = [
            threshold_range[0] + i * (threshold_range[1] - threshold_range[0]) / (num_threshold_steps - 1)
            for i in range(num_threshold_steps)
        ]
        
        total_combinations = len(layers_to_test) * len(steering_types) * len(thresholds_to_test)
        
        for layer_idx in layers_to_test:
            for steering_type in steering_types:
                for threshold in thresholds_to_test:
                    try:
                        # Create and train steering method
                        steering_method = SteeringMethod(
                            method_type=SteeringType(steering_type),
                            device=self.device
                        )
                        
                        # Create training pair set
                        train_pair_set = ContrastivePairSet.from_phrase_pairs(
                            name=f"train_layer_{layer_idx}",
                            phrase_pairs=train_pairs,
                            task_type="parameter_optimization"
                        )
                        
                        layer_obj = Layer(index=layer_idx, type="transformer")
                        
                        # Train
                        train_results = train_pair_set.train_classifier(
                            steering_method.classifier,
                            layer_obj
                        )
                        
                        # Test
                        test_pair_set = ContrastivePairSet.from_phrase_pairs(
                            name=f"test_layer_{layer_idx}",
                            phrase_pairs=test_pairs,
                            task_type="parameter_optimization"
                        )
                        
                        test_results = test_pair_set.evaluate_with_vectors(
                            steering_method,
                            layer_obj
                        )
                        
                        # Get score (use accuracy or F1)
                        score = test_results.get("accuracy", 0.0)
                        
                        result = {
                            "layer": layer_idx,
                            "steering_type": steering_type,
                            "threshold": threshold,
                            "score": score,
                            "train_results": train_results,
                            "test_results": test_results
                        }
                        
                        all_results.append(result)
                        
                        # Update best if better
                        if score > best_score:
                            best_score = score
                            best_params = {
                                "layer": layer_idx,
                                "steering_type": steering_type,
                                "threshold": threshold,
                                "score": score
                            }
                        
                    except Exception as e:
                        # Continue with next combination
                        pass
        
        return {
            "best_parameters": best_params,
            "best_score": best_score,
            "all_results": all_results,
            "total_combinations_tested": len(all_results)
        }
    
    def optimize_layer_selection(
        self,
        harmful_texts: List[str],
        harmless_texts: List[str],
        layer_range: Tuple[int, int] = (5, 25),
        steering_type: str = "logistic"
    ) -> Dict[str, Any]:
        """
        Optimize layer selection specifically.
        
        Args:
            harmful_texts: List of harmful examples
            harmless_texts: List of harmless examples
            layer_range: Range of layers to test
            steering_type: Steering type to use
            
        Returns:
            Layer optimization results
        """
        return self.optimize_parameters(
            harmful_texts=harmful_texts,
            harmless_texts=harmless_texts,
            layer_range=layer_range,
            steering_types=[steering_type],
            threshold_range=(0.5, 0.5),  # Fixed threshold
            num_threshold_steps=1
        )
    
    def optimize_threshold(
        self,
        harmful_texts: List[str],
        harmless_texts: List[str],
        layer: int = 15,
        steering_type: str = "logistic",
        threshold_range: Tuple[float, float] = (0.1, 0.9),
        num_steps: int = 10
    ) -> Dict[str, Any]:
        """
        Optimize threshold specifically.
        
        Args:
            harmful_texts: List of harmful examples
            harmless_texts: List of harmless examples
            layer: Fixed layer to use
            steering_type: Fixed steering type to use
            threshold_range: Range of thresholds to test
            num_steps: Number of threshold values to test
            
        Returns:
            Threshold optimization results
        """
        return self.optimize_parameters(
            harmful_texts=harmful_texts,
            harmless_texts=harmless_texts,
            layer_range=(layer, layer),  # Fixed layer
            steering_types=[steering_type],
            threshold_range=threshold_range,
            num_threshold_steps=num_steps
        )
    
    @staticmethod
    def get_available_tasks() -> List[str]:
        """
        Get list of all available tasks.
        
        Returns:
            List of available task names
        """
        return _get_available_tasks()
    
    @staticmethod
    def is_valid_task(task_name: str) -> bool:
        """
        Check if a task name is valid, including fuzzy matching.
        
        Args:
            task_name: Name of the task to check
            
        Returns:
            True if task is valid, False otherwise
        """
        return _is_valid_task(task_name)
    
    def load_lm_eval_task(self, task_name: str, shots: int = 0, limit: Optional[int] = None):
        """
        Load a task from lm-evaluation-harness with dynamic task name resolution.
        
        Args:
            task_name: Name of the task
            shots: Number of few-shot examples
            limit: Optional limit on number of documents
            
        Returns:
            Task object from lm_eval
        """
        # Check if it's LiveCodeBench
        if task_name == "livecodebench":
            from .tasks.livecodebench_task import LiveCodeBenchTask
            return LiveCodeBenchTask(release_version="release_v1", limit=limit)
            
        # Check if it's a BigCode task
        try:
            from .bigcode_integration import is_bigcode_task, load_bigcode_task
            
            if is_bigcode_task(task_name):
                # Load from BigCode
                return load_bigcode_task(task_name, limit=limit)
        except ImportError:
            # BigCode not available, continue with lm-eval
            pass
            
        # Default to lm-eval
        if _task_manager is None:
            raise RuntimeError("Task management system not available")
        
        return _task_manager.load_task(task_name, shots=shots, limit=limit)
    
    def _resolve_task_name(self, task_name: str) -> str:
        """Dynamically resolve a task name to an available task."""
        return _resolve_task_name(task_name)
    
    def _calculate_task_name_similarity(self, name1: str, name2: str) -> float:
        """Calculate similarity between two task names."""
        # Exact match
        if name1 == name2:
            return 1.0
        
        # Substring matching
        if name1 in name2 or name2 in name1:
            longer = max(name1, name2, key=len)
            shorter = min(name1, name2, key=len)
            return len(shorter) / len(longer)
        
        # Token-based similarity
        tokens1 = set(name1.replace('_', ' ').replace('-', ' ').split())
        tokens2 = set(name2.replace('_', ' ').replace('-', ' ').split())
        
        if not tokens1 or not tokens2:
            return 0.0
        
        intersection = len(tokens1 & tokens2)
        union = len(tokens1 | tokens2)
        
        return intersection / union if union > 0 else 0.0
    
    def split_task_data(self, task_data, split_ratio: float = 0.8, random_seed: int = 42):
        """
        Split task data into training and testing sets.
        
        Args:
            task_data: Task object from lm_eval
            split_ratio: Proportion for training set
            random_seed: Random seed for reproducibility
            
        Returns:
            Tuple of (train_docs, test_docs)
        """
        if _task_manager is None:
            raise RuntimeError("Task management system not available")
        
        return _task_manager.split_task_data(task_data, split_ratio=split_ratio, random_seed=random_seed)
    
    def prepare_prompts_from_docs(self, task, docs: List[Dict[str, Any]]) -> List[str]:
        """
        Prepare prompts from task documents.
        
        Args:
            task: Task object from lm_eval
            docs: List of documents
            
        Returns:
            List of formatted prompts
        """
        if _task_manager is None:
            raise RuntimeError("Task management system not available")
        
        return _task_manager.prepare_prompts_from_docs(task, docs)
    
    def get_reference_answers(self, task, docs: List[Dict[str, Any]]) -> List[str]:
        """
        Extract reference answers from task documents.
        
        Args:
            task: Task object from lm_eval
            docs: List of documents
            
        Returns:
            List of reference answers
        """
        if _task_manager is None:
            raise RuntimeError("Task management system not available")
        
        return _task_manager.get_reference_answers(task, docs)


class ModelParameterOptimizer:
    """
    Parameter optimizer integrated into the model primitive.
    """
    
    def __init__(
        self,
        model_name: str,
        device: Optional[str] = None
    ):
        """Initialize parameter optimizer."""
        self.model = Model(name=model_name, device=device)
        self.optimization_history = []
    
    def optimize_parameters(
        self,
        harmful_texts: List[str],
        harmless_texts: List[str],
        **kwargs
    ) -> Dict[str, Any]:
        """Optimize parameters using the model's optimization functionality."""
        result = self.model.optimize_parameters(harmful_texts, harmless_texts, **kwargs)
        self.optimization_history.append(result)
        return result
    
    def optimize_layer_selection(self, *args, **kwargs) -> Dict[str, Any]:
        """Optimize layer selection."""
        result = self.model.optimize_layer_selection(*args, **kwargs)
        self.optimization_history.append(result)
        return result
    
    def optimize_threshold(self, *args, **kwargs) -> Dict[str, Any]:
        """Optimize threshold."""
        result = self.model.optimize_threshold(*args, **kwargs)
        self.optimization_history.append(result)
        return result
    
    def get_optimization_history(self) -> List[Dict[str, Any]]:
        """Get history of optimization runs."""
        return self.optimization_history
    
    def clear_history(self) -> None:
        """Clear optimization history."""
        self.optimization_history.clear()


class ActivationHooks:
    """
    Activation hooks functionality integrated into the model primitive.
    """
    
    def __init__(
        self,
        model: torch.nn.Module,
        layers: List[int],
        token_strategy: str = "last"
    ):
        """Initialize activation hooks for a model."""
        self.model = model
        self.layers = layers
        self.token_strategy = token_strategy
        
        # Determine model type
        self.model_type = self._detect_model_type(model)
        
        # Initialize hook storage
        self.hooks = {}
        self.activations = {}
        self.layer_activations = {}
        self.active_layers = set()
        self.last_token_id = None
        self.last_token_position = None
        
        # Set up monitoring on specified layers
        self.setup_hooks(layers)
    
    def _detect_model_type(self, model) -> str:
        """Detect model type."""
        model_config = getattr(model, "config", None)
        model_name = getattr(model_config, "_name_or_path", "unknown").lower()
        
        if hasattr(model, "get_input_embeddings"):
            if "llama" in model_name:
                return "llama"
            elif "mistral" in model_name:
                return "mistral"
            elif "mpt" in model_name:
                return "mpt"
        
        return "generic"
    
    def _get_layer_name(self, model_type: str, layer_idx: int) -> str:
        """Get the layer name for a given model type and layer index."""
        if model_type == "llama":
            return f"model.layers.{layer_idx}"
        elif model_type == "mistral":
            return f"model.layers.{layer_idx}"
        elif model_type == "mpt":
            return f"transformer.blocks.{layer_idx}"
        elif model_type == "gpt2":
            return f"transformer.h.{layer_idx}"
        else:
            # Check user config for custom models
            if user_model_configs.has_config(self.name):
                layer_info = user_model_configs.get_layer_access_info(self.name)
                if layer_info and layer_info.get("layer_path_template"):
                    return layer_info["layer_path_template"].format(idx=layer_idx)
            
            # No config found - this shouldn't happen since we prompt during init
            # but just in case, prompt again
            print(f"\n⚠️  Model '{self.name}' layer configuration is missing.")
            response = input("Would you like to configure it now? (y/n): ").strip().lower()
            
            if response == 'y':
                config = user_model_configs.prompt_and_save_config(self.name)
                layer_info = user_model_configs.get_layer_access_info(self.name)
                if layer_info and layer_info.get("layer_path_template"):
                    return layer_info["layer_path_template"].format(idx=layer_idx)
            
            raise ValueError(
                f"Model '{self.name}' requires layer configuration.\n"
                f"Run 'wisent-guard configure-model {self.name}' to set it up."
            )
    
    def _get_module_by_name(self, name: str) -> torch.nn.Module:
        """Retrieve a module from the model by its name."""
        module = self.model
        for part in name.split('.'):
            if part.isdigit():
                module = module[int(part)]
            else:
                module = getattr(module, part)
        return module
    
    def _activation_hook(self, layer_idx: int):
        """Create a hook function for the specified layer."""
        def hook(module, input, output):
            if layer_idx in self.active_layers:
                # Get the output hidden states
                if isinstance(output, tuple):
                    hidden_states = output[0]
                else:
                    hidden_states = output
                
                if isinstance(hidden_states, torch.Tensor):
                    device = hidden_states.device
                    
                    # Get last token's activations
                    last_token_idx = hidden_states.shape[1] - 1
                    
                    # Store activations - preserve gradients if needed for steering
                    # Check if we're in a gradient-enabled context (for steering methods like K-steering)
                    if torch.is_grad_enabled() and hidden_states.requires_grad:
                        # Preserve computational graph for gradient-based steering
                        self.layer_activations[layer_idx] = hidden_states[:, last_token_idx, :].clone().to(device)
                    else:
                        # Standard extraction - detach for memory efficiency
                        self.layer_activations[layer_idx] = hidden_states[:, last_token_idx, :].detach().clone().to(device)
                    
                    self.last_token_position = last_token_idx
                    
                    # Try to get the token ID if available
                    if hasattr(module, '_last_input_ids'):
                        input_ids = module._last_input_ids[0]
                        if last_token_idx < len(input_ids):
                            try:
                                self.last_token_id = input_ids[last_token_idx].item()
                            except (RuntimeError, ValueError):
                                self.last_token_id = None
        return hook
    
    def register_hooks(self, layers: List[int]) -> None:
        """Register activation hooks for the specified layers."""
        self.active_layers = set(layers)
        
        # Clear existing hooks
        self.remove_hooks()
        
        # Register new hooks
        for layer_idx in layers:
            layer_name = self._get_layer_name(self.model_type, layer_idx)
            
            try:
                module = self._get_module_by_name(layer_name)
                
                # Add a pre-hook to capture input_ids
                def forward_pre_hook(module, args):
                    if len(args) > 0 and isinstance(args[0], torch.Tensor):
                        module._last_input_ids = args[0]
                    return args
                
                module.register_forward_pre_hook(forward_pre_hook)
                
                # Add the main activation hook
                hook = module.register_forward_hook(self._activation_hook(layer_idx))
                self.hooks[layer_idx] = hook
                
            except Exception:
                # Continue with other layers if one fails
                pass
    
    def remove_hooks(self) -> None:
        """Remove all registered hooks."""
        if self.hooks:
            for layer_idx, hook in self.hooks.items():
                hook.remove()
            self.hooks = {}
    
    def setup_hooks(self, layers: List[int]):
        """Set up hooks for specified layers."""
        self.register_hooks(layers)
    
    def add_steering_hooks(self, layers: List[int], steering_function) -> List:
        """Add steering hooks to specified layers.
        
        Args:
            layers: List of layer indices to add hooks to
            steering_function: Function that takes (module, input, output) and returns modified output
            
        Returns:
            List of registered hooks that can be removed later
        """
        hooks = []
        
        for layer_idx in layers:
            layer_name = self._get_layer_name(self.model_type, layer_idx)
            
            try:
                module = self._get_module_by_name(layer_name)
                hook = module.register_forward_hook(steering_function)
                hooks.append(hook)
            except Exception as e:
                # Log error but continue with other layers
                print(f"Warning: Could not add steering hook to layer {layer_idx}: {e}")
                
        return hooks
    
    def has_activations(self) -> bool:
        """Check if we have collected any activations."""
        return bool(self.layer_activations)
    
    def get_activations(self) -> Dict[int, torch.Tensor]:
        """Get all activations for registered layers."""
        return self.layer_activations
    
    def get_last_token_info(self) -> Dict[str, Any]:
        """Get information about the last token processed."""
        return {
            "token_id": self.last_token_id,
            "position": self.last_token_position
        }
    
    def reset(self):
        """Reset stored activations."""
        self.activations = {}
        self.layer_activations = {}
    
    def clear_activations(self):
        """Clear all stored activations."""
        self.reset()


# LM Evaluation Integration - now handled by task management system
# Import from the dedicated task management system
try:
    from .agent.diagnose.tasks import (
        get_available_tasks as _get_available_tasks,
        is_valid_task as _is_valid_task,
        resolve_task_name as _resolve_task_name,
        TaskManager,
        load_docs
    )
    # Initialize global task manager for Model class methods
    _task_manager = TaskManager()
except ImportError:
    # Fallback if task management system not available
    _task_manager = None
    def _get_available_tasks():
        return []
    def _is_valid_task(task_name):
        return False
    def _resolve_task_name(task_name):
        raise ValueError(f"Task management system not available")
    def load_docs(task, limit=None):
        raise ValueError(f"Task management system not available")