"""Hyperparameter Sweep Infrastructure

Grid/random search over key parameters with automatic evaluation.
No more hand-tuning one parameter at a time — let the machine find the best config.
"""
import numpy as np
import pandas as pd
from itertools import product
from typing import Dict, List, Optional, Callable, Any, Tuple
import json
from dataclasses import dataclass
import warnings
warnings.filterwarnings('ignore')


@dataclass
class SweepConfig:
    """Configuration for a hyperparameter sweep"""
    param_grid: Dict[str, List[Any]]
    metric: str = 'sharpe_ratio'
    metric_direction: str = 'maximize'
    n_trials: Optional[int] = None  # For random search
    random_seed: int = 42
    

def grid_search(param_grid: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
    """
    Generate all combinations from parameter grid.
    
    Example:
        param_grid = {
            'learning_rate': [1e-4, 1e-3, 1e-2],
            'hidden_size': [64, 128, 256],
            'dropout': [0.1, 0.2, 0.3]
        }
        → 3 × 3 × 3 = 27 combinations
    
    WARNING: Grid search is exponential in parameters. 
    Use random search for high-dimensional spaces.
    """
    keys = list(param_grid.keys())
    values = list(param_grid.values())
    
    combinations = []
    for combo in product(*values):
        combinations.append(dict(zip(keys, combo)))
    
    return combinations


def random_search(param_grid: Dict[str, List[Any]], 
                  n_trials: int,
                  random_seed: int = 42) -> List[Dict[str, Any]]:
    """
    Random search over parameter grid.
    
    Often more efficient than grid search (Bergstra & Bengio, 2012):
    Random search finds good hyperparameters faster than grid search
    in high-dimensional spaces.
    """
    np.random.seed(random_seed)
    
    combinations = []
    for _ in range(n_trials):
        config = {}
        for key, values in param_grid.items():
            config[key] = np.random.choice(values)
        combinations.append(config)
    
    return combinations


def latin_hypercube_sampling(param_ranges: Dict[str, Tuple[float, float]],
                              n_trials: int,
                              discrete_params: Optional[Dict[str, List]] = None,
                              random_seed: int = 42) -> List[Dict[str, Any]]:
    """
    Latin Hypercube Sampling for efficient space coverage.
    
    Divides each dimension into n equal strata and samples once from each.
    Ensures better coverage of the parameter space than random.
    
    Args:
        param_ranges: {param_name: (min, max)} for continuous params
        n_trials: Number of samples
        discrete_params: {param_name: [values]} for discrete params
    """
    np.random.seed(random_seed)
    
    n_continuous = len(param_ranges)
    n_total = n_continuous + (len(discrete_params) if discrete_params else 0)
    
    # Generate LHS samples for continuous params
    samples = np.zeros((n_trials, n_continuous))
    for i in range(n_continuous):
        # Divide [0,1] into n intervals
        intervals = np.linspace(0, 1, n_trials + 1)
        # Sample uniformly within each interval
        points = intervals[:-1] + np.random.uniform(0, 1/n_trials, n_trials)
        # Shuffle
        np.random.shuffle(points)
        samples[:, i] = points
    
    # Convert to parameter values
    combinations = []
    param_names = list(param_ranges.keys())
    
    for j in range(n_trials):
        config = {}
        for i, name in enumerate(param_names):
            low, high = param_ranges[name]
            config[name] = low + samples[j, i] * (high - low)
        
        # Add discrete params
        if discrete_params:
            for name, values in discrete_params.items():
                config[name] = np.random.choice(values)
        
        combinations.append(config)
    
    return combinations


class HyperparameterTuner:
    """
    Hyperparameter tuner with multiple search strategies.
    
    Usage:
        tuner = HyperparameterTuner(strategy='random')
        best_config, results = tuner.search(
            param_grid,
            train_fn=train_and_evaluate,
            n_trials=50
        )
    """
    
    def __init__(self, strategy: str = 'random'):
        self.strategy = strategy
        self.results = []
    
    def search(self, 
               param_grid: Dict[str, List[Any]],
               train_fn: Callable[[Dict], Dict[str, float]],
               n_trials: Optional[int] = None,
               metric: str = 'sharpe_ratio',
               direction: str = 'maximize',
               verbose: bool = True) -> Tuple[Dict, pd.DataFrame]:
        """
        Run hyperparameter search.
        
        Args:
            param_grid: Parameter grid
            train_fn: Function(params) -> dict of metrics
            n_trials: Number of trials (for random/LHS)
            metric: Metric to optimize
            direction: 'maximize' or 'minimize'
        
        Returns:
            best_config: Best hyperparameter configuration
            results_df: DataFrame of all trials
        """
        # Generate configurations
        if self.strategy == 'grid':
            configs = grid_search(param_grid)
        elif self.strategy == 'random':
            configs = random_search(param_grid, n_trials or 20)
        elif self.strategy == 'lhs':
            # Separate continuous and discrete
            continuous = {k: v for k, v in param_grid.items() 
                         if isinstance(v, tuple) and len(v) == 2}
            discrete = {k: v for k, v in param_grid.items() 
                       if k not in continuous}
            configs = latin_hypercube_sampling(continuous, n_trials or 20, discrete)
        else:
            raise ValueError(f"Unknown strategy: {self.strategy}")
        
        print(f"Running {len(configs)} trials with {self.strategy} search...")
        
        # Evaluate each configuration
        results = []
        for i, config in enumerate(configs):
            if verbose:
                print(f"\nTrial {i+1}/{len(configs)}: {config}")
            
            try:
                metrics = train_fn(config)
                
                result = {
                    'trial': i,
                    'status': 'success',
                    'config': config,
                    **metrics
                }
                
                if verbose:
                    print(f"  → {metric} = {metrics.get(metric, 'N/A')}")
                
            except Exception as e:
                result = {
                    'trial': i,
                    'status': 'failed',
                    'error': str(e),
                    'config': config
                }
                if verbose:
                    print(f"  → FAILED: {e}")
            
            results.append(result)
        
        # Find best configuration
        valid_results = [r for r in results if r.get('status') == 'success']
        
        if not valid_results:
            print("WARNING: All trials failed!")
            return {}, pd.DataFrame(results)
        
        if direction == 'maximize':
            best_result = max(valid_results, key=lambda r: r.get(metric, -np.inf))
        else:
            best_result = min(valid_results, key=lambda r: r.get(metric, np.inf))
        
        best_config = best_result['config']
        
        # Create results DataFrame
        results_df = pd.DataFrame(results)
        
        # Flatten config columns
        if 'config' in results_df.columns:
            config_df = pd.json_normalize(results_df['config'].tolist())
            config_df.columns = [f'param_{c}' for c in config_df.columns]
            results_df = pd.concat([results_df.drop('config', axis=1), config_df], axis=1)
        
        print(f"\n{'='*60}")
        print(f"BEST CONFIGURATION:")
        print(f"  {metric}: {best_result.get(metric):.4f}")
        for k, v in best_config.items():
            print(f"  {k}: {v}")
        print(f"{'='*60}")
        
        return best_config, results_df
    
    def analyze_importance(self, results_df: pd.DataFrame,
                          metric: str) -> pd.DataFrame:
        """
        Analyze which hyperparameters matter most.
        
        Uses correlation between each parameter and the metric.
        """
        param_cols = [c for c in results_df.columns if c.startswith('param_')]
        
        if not param_cols:
            return pd.DataFrame()
        
        importance = []
        for col in param_cols:
            param_name = col.replace('param_', '')
            
            # Calculate correlation with metric
            valid = results_df.dropna(subset=[col, metric])
            if len(valid) > 3:
                corr = np.corrcoef(valid[col].values, valid[metric].values)[0, 1]
                if not np.isnan(corr):
                    importance.append({
                        'parameter': param_name,
                        'correlation': corr,
                        'abs_correlation': abs(corr),
                        'importance_rank': abs(corr)
                    })
        
        importance_df = pd.DataFrame(importance)
        importance_df = importance_df.sort_values('abs_correlation', ascending=False)
        importance_df['importance_rank'] = range(1, len(importance_df) + 1)
        
        return importance_df


def create_alpha_model_sweep() -> Dict:
    """
    Pre-configured sweep for AlphaForge alpha model.
    
    Key parameters to tune:
    - lookback_window: How much history to use
    - lstm_hidden_size: Model capacity
    - lstm_layers: Depth
    - dropout: Regularization
    - learning_rate: Optimization
    - ensemble_weights: How to combine models
    """
    return {
        'lookback_window': [30, 60, 90, 120],
        'lstm_hidden_size': [64, 128, 256],
        'lstm_num_layers': [1, 2, 3],
        'lstm_dropout': [0.1, 0.2, 0.3],
        'transformer_d_model': [64, 128],
        'transformer_nhead': [2, 4],
        'transformer_num_layers': [1, 2],
        'learning_rate': [1e-5, 5e-5, 1e-4, 5e-4],
        'batch_size': [32, 64, 128],
        'xgb_max_depth': [4, 6, 8],
        'xgb_n_estimators': [100, 200, 500],
        'ensemble_lstm_weight': [0.2, 0.3, 0.4],
        'ensemble_transformer_weight': [0.2, 0.3, 0.4],
        'ensemble_xgboost_weight': [0.2, 0.4, 0.5]
    }


def create_portfolio_sweep() -> Dict:
    """Pre-configured sweep for portfolio optimizer"""
    return {
        'max_weight': [0.15, 0.20, 0.25, 0.30],
        'risk_aversion': [0.5, 1.0, 2.0, 3.0],
        'turnover_penalty': [0.0005, 0.001, 0.002],
        'rebalance_freq': [1, 3, 5, 10, 21],
        'risk_free_rate': [0.02, 0.03, 0.04, 0.05]
    }


def create_mtl_sweep() -> Dict:
    """Pre-configured sweep for Multi-Task Learning model"""
    return {
        'hidden_dim': [64, 128, 256],
        'n_lstm_layers': [1, 2, 3],
        'dropout': [0.1, 0.15, 0.2, 0.3],
        'learning_rate': [1e-5, 5e-5, 1e-4],
        'weight_return': [0.5, 1.0, 2.0],
        'weight_volatility': [0.25, 0.5, 1.0],
        'weight_portfolio': [1.0, 2.0, 3.0],
        'weight_direction': [0.1, 0.3, 0.5],
        'max_grad_norm': [0.1, 0.5, 1.0]
    }


def example_sweep():
    """Example of running a hyperparameter sweep"""
    # Define a simple objective function
    def mock_train(config):
        # Simulate training with different parameters
        lr = config.get('learning_rate', 1e-4)
        hidden = config.get('hidden_size', 128)
        dropout = config.get('dropout', 0.2)
        
        # Mock metric: Sharpe ratio (simulate a surface)
        # Best around lr=5e-5, hidden=128, dropout=0.15
        sharpe = 0.5 + np.exp(-((np.log10(lr) - (-4.3))**2) * 10) * 0.5
        sharpe += np.exp(-((hidden - 128)**2) / 5000) * 0.3
        sharpe += (0.2 - abs(dropout - 0.15)) * 0.2
        sharpe += np.random.randn() * 0.1  # Noise
        
        return {
            'sharpe_ratio': sharpe,
            'ic': sharpe * 0.3,
            'max_drawdown': -0.15 + np.random.rand() * 0.1
        }
    
    # Parameter grid
    param_grid = {
        'learning_rate': [1e-5, 5e-5, 1e-4, 5e-4],
        'hidden_size': [64, 128, 256],
        'dropout': [0.1, 0.2, 0.3]
    }
    
    # Run random search
    tuner = HyperparameterTuner(strategy='random')
    best_config, results = tuner.search(
        param_grid,
        mock_train,
        n_trials=20,
        metric='sharpe_ratio',
        direction='maximize'
    )
    
    # Analyze importance
    importance = tuner.analyze_importance(results, 'sharpe_ratio')
    print("\nParameter Importance:")
    print(importance.to_string())
    
    return best_config, results


if __name__ == '__main__':
    best_config, results = example_sweep()