"""Search space generation for MethodOptimizer."""
from __future__ import annotations

import logging
from typing import Any, Dict, Iterator, List, Optional, Tuple

from wisent.core.activations import ExtractionStrategy
from wisent.core.cli.optimization.core.method_optimizer_config import OptimizationConfig
from wisent.core.steering_methods.registry import SteeringMethodRegistry

logger = logging.getLogger(__name__)


    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"]
            steering_strategies = ["constant"]
        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", "max_pooling", "continuation_token"]
            prompt_strats = custom_prompt_strategies or ["chat_template", "direct_completion", "multiple_choice", "role_playing", "instruction_following"]
            steering_strategies = ["constant", "initial_only", "diminishing", "increasing", "gaussian"]
        
        # 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": ExtractionStrategy.CHAT_LAST,
            "mean_pooling": ExtractionStrategy.CHAT_MEAN,
            "first_token": ExtractionStrategy.CHAT_FIRST,
            "max_pooling": ExtractionStrategy.CHAT_MAX_NORM,
            "continuation_token": ExtractionStrategy.CHAT_FIRST,  # First answer token
        }
        
        prompt_strat_map = {
            "chat_template": ExtractionStrategy.CHAT_LAST,
            "direct_completion": ExtractionStrategy.CHAT_LAST,
            "multiple_choice": ExtractionStrategy.MC_BALANCED,
            "role_playing": ExtractionStrategy.ROLE_PLAY,
            "instruction_following": ExtractionStrategy.CHAT_LAST,
        }
        
        # 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 steering_strat in steering_strategies:
                            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, ExtractionStrategy.CHAT_LAST
                                    ),
                                    prompt_strategy=prompt_strat_map.get(
                                        prompt_strat_name, ExtractionStrategy.CHAT_LAST
                                    ),
                                    strength=strength,
                                    strategy=steering_strat,
                                    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 ALL layers from 0 to num_layers-1
        return list(range(num_layers))
    
    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 == "tecza":
            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 == "tetno":
            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 == "grom":
            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
        """
        # Store activations on CPU - GROM and other methods expect CPU tensors for training
        collector = ActivationCollector(model=self.model, store_device="cpu")

        updated_pairs = []
        for pair in pairs.pairs:
            updated_pair = collector.collect(
                pair,
                strategy=config.token_aggregation,
                layers=config.layers,
                normalize=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,
        )
