"""
Universal Steering Method Optimizer.

This module provides a foolproof way to optimize ANY steering method's parameters.
It uses the universal `train(pair_set)` interface that all methods must implement,
ensuring compatibility with current and future steering methods.

Key design principles:
1. Uses `method.train(pair_set)` - the universal interface all methods implement
2. Automatically extracts method-specific parameters from the registry
3. Search spaces are defined per-method and automatically iterated
4. Evaluation is decoupled from training - any evaluator can be used
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type

import torch

from wisent.core.activations.activations_collector import ActivationCollector
from wisent.core.activations.core.atoms import ActivationAggregationStrategy, LayerActivations
from wisent.core.activations.prompt_construction_strategy import PromptConstructionStrategy
from wisent.core.contrastive_pairs.core.pair import ContrastivePair
from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
from wisent.core.models.core.atoms import SteeringPlan, SteeringVector
from wisent.core.steering_methods.core.atoms import BaseSteeringMethod
from wisent.core.steering_methods.registry import SteeringMethodRegistry

logger = logging.getLogger(__name__)


@dataclass
class OptimizationConfig:
    """Configuration for a single optimization trial."""
    
    method_name: str
    """Name of the steering method (caa, prism, pulse, titan)."""
    
    # Activation extraction parameters
    layers: List[str]
    """Layer indices to extract activations from."""
    
    token_aggregation: ActivationAggregationStrategy
    """How to aggregate tokens within a sequence."""
    
    prompt_strategy: PromptConstructionStrategy
    """How to construct prompts for the model."""
    
    # Application parameters
    strength: float = 1.0
    """Steering strength multiplier."""
    
    strategy: str = "constant"
    """Steering application strategy."""
    
    # Method-specific parameters
    method_params: Dict[str, Any] = field(default_factory=dict)
    """Method-specific parameters (num_directions, sensor_layer, etc.)."""
    
    def __hash__(self):
        return hash((
            self.method_name,
            tuple(self.layers),
            self.token_aggregation.value,
            self.prompt_strategy.value,
            self.strength,
            self.strategy,
            tuple(sorted(self.method_params.items())),
        ))
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "method_name": self.method_name,
            "layers": self.layers,
            "token_aggregation": self.token_aggregation.value,
            "prompt_strategy": self.prompt_strategy.value,
            "strength": self.strength,
            "strategy": self.strategy,
            "method_params": self.method_params,
        }


@dataclass
class OptimizationResult:
    """Result of a single optimization trial."""
    
    config: OptimizationConfig
    """The configuration that was tested."""
    
    score: float
    """Primary evaluation score."""
    
    metrics: Dict[str, float] = field(default_factory=dict)
    """Additional metrics (accuracy, f1, etc.)."""
    
    steering_vectors: Optional[LayerActivations] = None
    """Trained steering vectors (optional, for caching)."""
    
    training_time: float = 0.0
    """Time taken to train the method."""
    
    evaluation_time: float = 0.0
    """Time taken to evaluate."""
    
    metadata: Dict[str, Any] = field(default_factory=dict)
    """Additional metadata from training."""


@dataclass
class OptimizationSummary:
    """Summary of optimization run."""
    
    best_result: OptimizationResult
    """Best result found."""
    
    all_results: List[OptimizationResult]
    """All results from the optimization."""
    
    method_name: str
    """Method that was optimized."""
    
    task_name: str
    """Task/benchmark used for evaluation."""
    
    total_time: float = 0.0
    """Total optimization time."""
    
    configs_tested: int = 0
    """Number of configurations tested."""


class MethodOptimizer:
    """
    Universal optimizer for any steering method.
    
    This optimizer works with ANY steering method by using the universal
    `train(pair_set)` interface. It handles:
    
    1. Collecting activations with different extraction parameters
    2. Training methods with different method-specific parameters
    3. Evaluating results with any provided evaluator
    4. Tracking and comparing results across configurations
    
    Example usage:
    
        optimizer = MethodOptimizer(model, method_name="titan")
        
        # Generate search space
        configs = optimizer.generate_search_space(num_layers=16, quick=True)
        
        # Run optimization
        summary = optimizer.optimize(
            train_pairs=train_pairs,
            test_pairs=test_pairs,
            evaluator=evaluator,
            configs=configs,
        )
        
        print(f"Best config: {summary.best_result.config}")
        print(f"Best score: {summary.best_result.score}")
    """
    
    def __init__(
        self,
        model,
        method_name: str,
        device: str = "cpu",
        verbose: bool = True,
    ):
        """
        Initialize the optimizer.
        
        Args:
            model: WisentModel instance
            method_name: Name of steering method to optimize
            device: Device for storing activations
            verbose: Whether to print progress
        """
        self.model = model
        self.method_name = method_name.lower()
        self.device = device
        self.verbose = verbose
        
        # Validate method exists
        if not SteeringMethodRegistry.validate_method(self.method_name):
            available = ", ".join(SteeringMethodRegistry.list_methods())
            raise ValueError(f"Unknown method: {method_name}. Available: {available}")
        
        self.method_definition = SteeringMethodRegistry.get(self.method_name)
    
    def _log(self, msg: str):
        if self.verbose:
            print(msg)
    
    def generate_search_space(
        self,
        num_layers: int,
        quick: bool = False,
        custom_layers: Optional[List[int]] = None,
        custom_strengths: Optional[List[float]] = None,
        custom_token_aggregations: Optional[List[str]] = None,
        custom_prompt_strategies: Optional[List[str]] = None,
        custom_method_params: Optional[Dict[str, List[Any]]] = None,
    ) -> List[OptimizationConfig]:
        """
        Generate search space for optimization.
        
        Args:
            num_layers: Number of layers in the model
            quick: Use reduced search space for faster testing
            custom_*: Override default search values
            
        Returns:
            List of OptimizationConfig to test
        """
        # Default extraction parameter ranges
        if quick:
            layers = custom_layers or self._get_quick_layers(num_layers)
            strengths = custom_strengths or [0.5, 1.0, 1.5]
            token_aggs = custom_token_aggregations or ["last_token"]
            prompt_strats = custom_prompt_strategies or ["chat_template"]
        else:
            layers = custom_layers or self._get_full_layers(num_layers)
            strengths = custom_strengths or [0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0]
            token_aggs = custom_token_aggregations or ["last_token", "mean_pooling", "first_token"]
            prompt_strats = custom_prompt_strategies or ["chat_template", "direct_completion"]
        
        # Get method-specific parameter ranges
        method_param_ranges = self._get_method_param_ranges(quick, custom_method_params)
        
        # Generate all configurations
        configs = []
        
        # Convert to enums
        token_agg_map = {
            "last_token": ActivationAggregationStrategy.LAST_TOKEN,
            "mean_pooling": ActivationAggregationStrategy.MEAN_POOLING,
            "first_token": ActivationAggregationStrategy.FIRST_TOKEN,
            "max_pooling": ActivationAggregationStrategy.MAX_POOLING,
            "continuation_token": ActivationAggregationStrategy.CONTINUATION_TOKEN,
        }
        
        prompt_strat_map = {
            "chat_template": PromptConstructionStrategy.CHAT_TEMPLATE,
            "direct_completion": PromptConstructionStrategy.DIRECT_COMPLETION,
            "multiple_choice": PromptConstructionStrategy.MULTIPLE_CHOICE,
            "role_playing": PromptConstructionStrategy.ROLE_PLAYING,
            "instruction_following": PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
        }
        
        # Generate method param combinations
        method_param_combos = self._generate_param_combinations(method_param_ranges)
        
        for layer in layers:
            for strength in strengths:
                for token_agg_name in token_aggs:
                    for prompt_strat_name in prompt_strats:
                        for method_params in method_param_combos:
                            # Determine which layers to collect activations for
                            # For multi-layer methods, collect all needed layers
                            activation_layers = self._determine_activation_layers(
                                layer, num_layers, method_params
                            )
                            
                            config = OptimizationConfig(
                                method_name=self.method_name,
                                layers=[str(l) for l in activation_layers],
                                token_aggregation=token_agg_map.get(
                                    token_agg_name, ActivationAggregationStrategy.LAST_TOKEN
                                ),
                                prompt_strategy=prompt_strat_map.get(
                                    prompt_strat_name, PromptConstructionStrategy.CHAT_TEMPLATE
                                ),
                                strength=strength,
                                method_params=method_params,
                            )
                            configs.append(config)
        
        return configs
    
    def _get_quick_layers(self, num_layers: int) -> List[int]:
        """Get reduced layer set for quick search."""
        if num_layers <= 12:
            return [num_layers // 3, num_layers // 2, 2 * num_layers // 3]
        elif num_layers <= 24:
            return [num_layers // 4, num_layers // 2, 3 * num_layers // 4]
        else:
            return [num_layers // 4, num_layers // 3, num_layers // 2, 2 * num_layers // 3]
    
    def _get_full_layers(self, num_layers: int) -> List[int]:
        """Get full layer set for comprehensive search."""
        # Test every 2nd layer in the middle-to-late portion
        start = num_layers // 4
        end = num_layers - 2
        return list(range(start, end, 2))
    
    def _get_method_param_ranges(
        self,
        quick: bool,
        custom: Optional[Dict[str, List[Any]]] = None,
    ) -> Dict[str, List[Any]]:
        """Get method-specific parameter ranges."""
        custom = custom or {}
        
        if self.method_name == "caa":
            return {
                "normalize": custom.get("normalize", [True]),
            }
        
        elif self.method_name == "prism":
            if quick:
                return {
                    "num_directions": custom.get("num_directions", [2, 3]),
                    "optimization_steps": custom.get("optimization_steps", [50]),
                    "retain_weight": custom.get("retain_weight", [0.1]),
                    "learning_rate": custom.get("learning_rate", [0.01]),
                    "use_caa_init": custom.get("use_caa_init", [True]),
                }
            return {
                "num_directions": custom.get("num_directions", [1, 2, 3, 5]),
                "optimization_steps": custom.get("optimization_steps", [50, 100]),
                "retain_weight": custom.get("retain_weight", [0.0, 0.1, 0.3]),
                "learning_rate": custom.get("learning_rate", [0.01]),
                "independence_weight": custom.get("independence_weight", [0.05]),
                "use_caa_init": custom.get("use_caa_init", [True]),
            }
        
        elif self.method_name == "pulse":
            if quick:
                return {
                    "sensor_layer": custom.get("sensor_layer", ["auto"]),  # Will be resolved
                    "steering_layers": custom.get("steering_layers", ["range_3"]),
                    "condition_threshold": custom.get("condition_threshold", [0.5]),
                    "gate_temperature": custom.get("gate_temperature", [0.5]),
                    "use_entropy_scaling": custom.get("use_entropy_scaling", [False]),
                    "max_alpha": custom.get("max_alpha", [2.0]),
                }
            return {
                "sensor_layer": custom.get("sensor_layer", ["auto"]),
                "steering_layers": custom.get("steering_layers", ["single", "range_3", "range_5"]),
                "condition_threshold": custom.get("condition_threshold", [0.3, 0.5, 0.7]),
                "gate_temperature": custom.get("gate_temperature", [0.1, 0.5, 1.0]),
                "per_layer_scaling": custom.get("per_layer_scaling", [True, False]),
                "use_entropy_scaling": custom.get("use_entropy_scaling", [True, False]),
                "max_alpha": custom.get("max_alpha", [1.5, 2.0, 3.0]),
            }
        
        elif self.method_name == "titan":
            if quick:
                return {
                    "num_directions": custom.get("num_directions", [3]),
                    "sensor_layer": custom.get("sensor_layer", ["auto"]),
                    "steering_layers": custom.get("steering_layers", ["range_3"]),
                    "gate_hidden_dim": custom.get("gate_hidden_dim", [64]),
                    "intensity_hidden_dim": custom.get("intensity_hidden_dim", [32]),
                    "optimization_steps": custom.get("optimization_steps", [100]),
                    "behavior_weight": custom.get("behavior_weight", [1.0]),
                    "retain_weight": custom.get("retain_weight", [0.2]),
                    "max_alpha": custom.get("max_alpha", [2.0]),
                }
            return {
                "num_directions": custom.get("num_directions", [2, 3, 5]),
                "sensor_layer": custom.get("sensor_layer", ["auto"]),
                "steering_layers": custom.get("steering_layers", ["range_3", "range_5", "all_late"]),
                "gate_hidden_dim": custom.get("gate_hidden_dim", [32, 64, 128]),
                "intensity_hidden_dim": custom.get("intensity_hidden_dim", [16, 32, 64]),
                "optimization_steps": custom.get("optimization_steps", [100, 200]),
                "behavior_weight": custom.get("behavior_weight", [0.5, 1.0]),
                "retain_weight": custom.get("retain_weight", [0.1, 0.2, 0.5]),
                "sparse_weight": custom.get("sparse_weight", [0.0, 0.05]),
                "max_alpha": custom.get("max_alpha", [2.0, 3.0]),
            }
        
        # Default for unknown methods - empty params
        return {}
    
    def _generate_param_combinations(
        self,
        param_ranges: Dict[str, List[Any]],
    ) -> List[Dict[str, Any]]:
        """Generate all combinations of method parameters."""
        if not param_ranges:
            return [{}]
        
        import itertools
        
        keys = list(param_ranges.keys())
        values = [param_ranges[k] for k in keys]
        
        combinations = []
        for combo in itertools.product(*values):
            combinations.append(dict(zip(keys, combo)))
        
        return combinations
    
    def _determine_activation_layers(
        self,
        base_layer: int,
        num_layers: int,
        method_params: Dict[str, Any],
    ) -> List[int]:
        """Determine which layers to collect activations for."""
        # For methods that need multi-layer activations
        steering_layers_config = method_params.get("steering_layers", "single")
        
        if steering_layers_config == "single":
            return [base_layer]
        elif steering_layers_config == "range_3":
            return list(range(max(0, base_layer - 1), min(num_layers, base_layer + 2)))
        elif steering_layers_config == "range_5":
            return list(range(max(0, base_layer - 2), min(num_layers, base_layer + 3)))
        elif steering_layers_config == "all_late":
            start = int(num_layers * 0.75)
            return list(range(start, num_layers - 1))
        elif isinstance(steering_layers_config, list):
            return steering_layers_config
        else:
            return [base_layer]
    
    def collect_activations(
        self,
        pairs: ContrastivePairSet,
        config: OptimizationConfig,
    ) -> ContrastivePairSet:
        """
        Collect activations for a pair set using the given config.
        
        Args:
            pairs: ContrastivePairSet to collect activations for
            config: Configuration specifying layers, aggregation, etc.
            
        Returns:
            ContrastivePairSet with activations populated
        """
        collector = ActivationCollector(model=self.model, store_device=self.device)
        
        updated_pairs = []
        for pair in pairs.pairs:
            updated_pair = collector.collect_for_pair(
                pair,
                layers=config.layers,
                aggregation=config.token_aggregation,
                return_full_sequence=False,
                normalize_layers=False,
            )
            updated_pairs.append(updated_pair)
        
        return ContrastivePairSet(
            name=pairs.name if hasattr(pairs, 'name') else "collected",
            pairs=updated_pairs,
            task_type=pairs.task_type if hasattr(pairs, 'task_type') else None,
        )
    
    def train_method(
        self,
        pair_set: ContrastivePairSet,
        config: OptimizationConfig,
    ) -> Tuple[LayerActivations, Dict[str, Any]]:
        """
        Train a steering method using the universal interface.
        
        Args:
            pair_set: ContrastivePairSet with activations already collected
            config: Configuration with method parameters
            
        Returns:
            Tuple of (LayerActivations, metadata dict)
        """
        # Create method instance with parameters from config
        method_class = self.method_definition.get_method_class()
        
        # Merge default params with config params
        params = self.method_definition.get_default_params()
        params.update(config.method_params)
        
        # Resolve special parameter values
        params = self._resolve_params(params, config)
        
        # Create method instance
        method = method_class(**params)
        
        # Train using universal interface
        steering_vectors = method.train(pair_set)
        
        # Extract any metadata from training
        metadata = {}
        if hasattr(method, "_training_logs"):
            metadata["training_logs"] = method._training_logs
        
        return steering_vectors, metadata
    
    def _resolve_params(
        self,
        params: Dict[str, Any],
        config: OptimizationConfig,
    ) -> Dict[str, Any]:
        """Resolve special parameter values like 'auto'."""
        resolved = dict(params)
        num_layers = self.model.num_layers
        
        # Resolve sensor_layer
        if resolved.get("sensor_layer") == "auto":
            # Use 75% through the network
            resolved["sensor_layer"] = int(num_layers * 0.75)
        
        # Resolve steering_layers from string config to actual layer list
        steering_config = resolved.get("steering_layers")
        if isinstance(steering_config, str):
            base_layer = int(config.layers[0]) if config.layers else num_layers // 2
            if steering_config == "single":
                resolved["steering_layers"] = [base_layer]
            elif steering_config == "range_3":
                resolved["steering_layers"] = list(range(
                    max(0, base_layer - 1),
                    min(num_layers, base_layer + 2)
                ))
            elif steering_config == "range_5":
                resolved["steering_layers"] = list(range(
                    max(0, base_layer - 2),
                    min(num_layers, base_layer + 3)
                ))
            elif steering_config == "all_late":
                start = int(num_layers * 0.75)
                resolved["steering_layers"] = list(range(start, num_layers - 1))
        
        return resolved
    
    def create_steering_plan(
        self,
        steering_vectors: LayerActivations,
        config: OptimizationConfig,
    ) -> SteeringPlan:
        """Create a SteeringPlan from trained vectors and config."""
        raw_map = steering_vectors.to_dict()
        
        # Apply strength scaling
        scaled_map = {}
        for layer, vec in raw_map.items():
            steering_vec = SteeringVector(vector=vec, scale=config.strength)
            scaled_map[layer] = steering_vec
        
        return SteeringPlan(
            layers=scaled_map,
            layers_description=[
                f"{self.method_name.upper()} {config.to_dict()}"
            ],
        )
    
    def evaluate(
        self,
        steering_plan: SteeringPlan,
        test_pairs: ContrastivePairSet,
        evaluator,
        task_name: str,
    ) -> Tuple[float, Dict[str, float]]:
        """
        Evaluate a steering plan.
        
        Args:
            steering_plan: The steering plan to evaluate
            test_pairs: Test pairs for evaluation
            evaluator: Evaluator instance (from EvaluatorRotator)
            task_name: Name of the task for evaluator context
            
        Returns:
            Tuple of (primary_score, metrics_dict)
        """
        correct = 0
        total = 0
        
        for pair in test_pairs.pairs:
            try:
                choices = [
                    pair.negative_response.model_response,
                    pair.positive_response.model_response,
                ]
                expected = pair.positive_response.model_response
                test_code = pair.metadata.get("test_code") if pair.metadata else None
                
                eval_result = evaluator.evaluate(
                    response="",
                    expected=expected,
                    model=self.model,
                    question=pair.prompt,
                    choices=choices,
                    steering_plan=steering_plan,
                    test_code=test_code,
                    task_name=task_name,
                )
                
                if eval_result.ground_truth == "TRUTHFUL":
                    correct += 1
                total += 1
            except Exception as e:
                logger.warning(f"Evaluation failed for pair: {e}")
                total += 1
        
        accuracy = correct / total if total > 0 else 0.0
        
        return accuracy, {
            "accuracy": accuracy,
            "correct": correct,
            "total": total,
        }
    
    def optimize(
        self,
        train_pairs: ContrastivePairSet,
        test_pairs: ContrastivePairSet,
        evaluator,
        task_name: str,
        configs: Optional[List[OptimizationConfig]] = None,
        quick: bool = False,
        max_configs: Optional[int] = None,
        progress_callback: Optional[Callable[[int, int, OptimizationResult], None]] = None,
    ) -> OptimizationSummary:
        """
        Run optimization over configurations.
        
        Args:
            train_pairs: Training pairs for steering vector extraction
            test_pairs: Test pairs for evaluation
            evaluator: Evaluator instance
            task_name: Name of task/benchmark
            configs: Configurations to test (or generate automatically)
            quick: Use quick search space if configs not provided
            max_configs: Maximum configurations to test
            progress_callback: Called after each config with (idx, total, result)
            
        Returns:
            OptimizationSummary with best result and all results
        """
        start_time = time.time()
        
        # Generate configs if not provided
        if configs is None:
            configs = self.generate_search_space(
                num_layers=self.model.num_layers,
                quick=quick,
            )
        
        # Limit configs if requested
        if max_configs and len(configs) > max_configs:
            self._log(f"Limiting search from {len(configs)} to {max_configs} configs")
            configs = configs[:max_configs]
        
        self._log(f"\n{'='*60}")
        self._log(f"Optimizing {self.method_name.upper()} with {len(configs)} configurations")
        self._log(f"{'='*60}\n")
        
        results: List[OptimizationResult] = []
        best_result: Optional[OptimizationResult] = None
        
        for idx, config in enumerate(configs):
            try:
                self._log(f"[{idx+1}/{len(configs)}] Testing: layers={config.layers}, "
                         f"strength={config.strength}, token_agg={config.token_aggregation.value}")
                
                # Collect activations
                t0 = time.time()
                train_with_acts = self.collect_activations(train_pairs, config)
                
                # Train method
                t1 = time.time()
                steering_vectors, metadata = self.train_method(train_with_acts, config)
                training_time = time.time() - t1
                
                # Create steering plan
                steering_plan = self.create_steering_plan(steering_vectors, config)
                
                # Evaluate
                t2 = time.time()
                score, metrics = self.evaluate(
                    steering_plan, test_pairs, evaluator, task_name
                )
                evaluation_time = time.time() - t2
                
                result = OptimizationResult(
                    config=config,
                    score=score,
                    metrics=metrics,
                    steering_vectors=steering_vectors,
                    training_time=training_time,
                    evaluation_time=evaluation_time,
                    metadata=metadata,
                )
                results.append(result)
                
                self._log(f"        Score: {score:.4f} (train: {training_time:.1f}s, eval: {evaluation_time:.1f}s)")
                
                # Track best
                if best_result is None or score > best_result.score:
                    best_result = result
                    self._log(f"        *** New best! ***")
                
                # Callback
                if progress_callback:
                    progress_callback(idx, len(configs), result)
                
            except Exception as e:
                logger.error(f"Config {idx} failed: {e}")
                self._log(f"        FAILED: {e}")
        
        total_time = time.time() - start_time
        
        self._log(f"\n{'='*60}")
        self._log(f"Optimization complete in {total_time:.1f}s")
        self._log(f"Best score: {best_result.score:.4f}" if best_result else "No results")
        self._log(f"{'='*60}\n")
        
        return OptimizationSummary(
            best_result=best_result,
            all_results=results,
            method_name=self.method_name,
            task_name=task_name,
            total_time=total_time,
            configs_tested=len(results),
        )


def optimize_steering_method(
    model,
    method_name: str,
    train_pairs: ContrastivePairSet,
    test_pairs: ContrastivePairSet,
    evaluator,
    task_name: str,
    quick: bool = False,
    max_configs: Optional[int] = None,
    verbose: bool = True,
) -> OptimizationSummary:
    """
    Convenience function to optimize a steering method.
    
    This is the main entry point for optimizing any steering method.
    It handles all the complexity of:
    - Generating appropriate search spaces
    - Collecting activations
    - Training methods
    - Evaluating results
    
    Args:
        model: WisentModel instance
        method_name: Name of method to optimize (caa, prism, pulse, titan)
        train_pairs: Training pairs
        test_pairs: Test pairs
        evaluator: Evaluator instance
        task_name: Name of task/benchmark
        quick: Use reduced search space
        max_configs: Maximum configurations to test
        verbose: Print progress
        
    Returns:
        OptimizationSummary with results
    """
    optimizer = MethodOptimizer(
        model=model,
        method_name=method_name,
        verbose=verbose,
    )
    
    return optimizer.optimize(
        train_pairs=train_pairs,
        test_pairs=test_pairs,
        evaluator=evaluator,
        task_name=task_name,
        quick=quick,
        max_configs=max_configs,
    )
