"""
Example 3: Hyperparameter Optimization
Tune ML model hyperparameters using DPO.
"""

import numpy as np
from dpo import dpo, dpo_optimize
from dpo.core.problem import ContinuousOptimizationProblem


def example_hpo_simple():
    """
    Simple HPO: optimize learning rate and dropout.
    """
    print("="*60)
    print("HPO Example 1: ML Model Hyperparameters")
    print("="*60)
    
    # Simulation of model training
    def train_model(lr, dropout):
        """
        Simulate training with given hyperparameters.
        Real use case: actually train the model.
        """
        # Simulated loss landscape
        loss = (lr - 0.001)**2 + (dropout - 0.3)**2 + np.random.randn() * 0.01
        loss = max(0.01, loss)
        
        accuracy = 0.95 - loss
        accuracy = np.clip(accuracy, 0, 1)
        
        return loss, {'accuracy': accuracy}
    
    # Objective function for DPO
    def objective(params):
        loss, metrics = train_model(params['learning_rate'], params['dropout'])
        return loss, metrics
    
    # Define bounds
    param_bounds = [
        (1e-5, 1e-1),  # learning_rate
        (0.0, 0.9),    # dropout
    ]
    param_names = ['learning_rate', 'dropout']
    
    # Solve
    result = dpo(
        problem_type='continuous',
        objective_fn=objective,
        param_bounds=param_bounds,
        param_names=param_names,
        preset='balanced',
        max_iterations=100,
        population_size=40
    )
    
    best_solution = result['best_solution']
    print(f"\nBest Hyperparameters:")
    print(f"  Learning Rate: {best_solution['learning_rate']:.6f}")
    print(f"  Dropout: {best_solution['dropout']:.4f}")
    print(f"  Best Loss: {result['best_fitness']:.6f}")
    print(f"  Best Accuracy: {result['best_accuracy']:.4f}")
    print(f"  Time: {result['elapsed_time']:.2f}s")
    
    return result


def example_hpo_advanced():
    """
    Advanced HPO: optimize multiple hyperparameters.
    """
    print("\n" + "="*60)
    print("HPO Example 2: Advanced Hyperparameter Tuning")
    print("="*60)
    
    def objective(params):
        """
        Optimize 5 hyperparameters:
        - learning_rate
        - batch_size (will be discretized)
        - dropout
        - l2_regularization
        - momentum
        """
        lr = params.get('learning_rate', 0.001)
        batch_size = params.get('batch_size', 32)
        dropout = params.get('dropout', 0.5)
        l2 = params.get('l2_regularization', 1e-5)
        momentum = params.get('momentum', 0.9)
        
        # Simulated loss function
        loss = (
            0.01 * (lr - 0.001)**2 +
            0.01 * np.log(batch_size / 32)**2 +
            0.02 * (dropout - 0.3)**2 +
            0.03 * np.log(l2) +
            0.01 * (momentum - 0.9)**2 +
            np.random.randn() * 0.005
        )
        
        loss = max(0.001, loss)
        accuracy = 0.98 - loss * 100
        accuracy = np.clip(accuracy, 0, 1)
        
        return loss, {'accuracy': accuracy, 'loss': float(loss)}
    
    param_bounds = [
        (1e-6, 1e-1),           # learning_rate
        (4, 256),               # batch_size (log scale)
        (0.0, 0.9),             # dropout
        (1e-8, 1e-2),           # l2_regularization (log scale)
        (0.0, 0.99),            # momentum
    ]
    
    param_names = [
        'learning_rate',
        'batch_size',
        'dropout',
        'l2_regularization',
        'momentum'
    ]
    
    result = dpo(
        problem_type='continuous',
        objective_fn=objective,
        param_bounds=param_bounds,
        param_names=param_names,
        preset='thorough',
        max_iterations=150,
        population_size=50
    )
    
    best = result['best_solution']
    print(f"\nBest Hyperparameters:")
    for name, value in best.items():
        print(f"  {name}: {value:.6f}")
    
    print(f"\nResults:")
    print(f"  Best Loss: {result['best_fitness']:.6f}")
    print(f"  Best Accuracy: {result['best_accuracy']:.4f}")
    print(f"  Evaluations: {result['total_evaluations']}")
    print(f"  Time: {result['elapsed_time']:.2f}s")
    
    return result


def example_hpo_with_constraints():
    """
    HPO with constraint: balance accuracy and inference time.
    """
    print("\n" + "="*60)
    print("HPO Example 3: Constrained Optimization")
    print("="*60)
    
    def objective(params):
        lr = params['learning_rate']
        model_size = params['model_size']
        
        # Simulated: larger models are more accurate but slower
        accuracy = 0.85 + 0.1 * (1.0 - np.exp(-model_size / 100.0))
        inference_time = 10.0 + model_size * 0.5  # ms
        
        # Loss: favors high accuracy but penalizes slow inference
        loss = (1.0 - accuracy) + 0.01 * max(0, inference_time - 50)
        
        metrics = {
            'accuracy': accuracy,
            'inference_time': inference_time,
            'model_size': model_size
        }
        
        return loss, metrics
    
    param_bounds = [
        (1e-5, 1e-2),   # learning_rate
        (10, 500),      # model_size (neurons, filters, etc.)
    ]
    
    param_names = ['learning_rate', 'model_size']
    
    result = dpo(
        problem_type='continuous',
        objective_fn=objective,
        param_bounds=param_bounds,
        param_names=param_names,
        preset='balanced',
        max_iterations=120,
        population_size=45
    )
    
    best = result['best_solution']
    metrics = result['best_metrics']
    
    print(f"\nOptimal Configuration:")
    print(f"  Learning Rate: {best['learning_rate']:.6f}")
    print(f"  Model Size: {best['model_size']:.0f}")
    print(f"  Accuracy: {metrics.get('accuracy', 0):.4f}")
    print(f"  Inference Time: {metrics.get('inference_time', 0):.2f}ms")
    print(f"  Total Loss: {result['best_fitness']:.6f}")
    
    return result


if __name__ == "__main__":
    result1 = example_hpo_simple()
    result2 = example_hpo_advanced()
    result3 = example_hpo_with_constraints()
    
    print("\n" + "="*60)
    print("All HPO examples complete!")
    print("="*60)
