//! # Averaged Adam Optimizer Comprehensive Benchmark
//!
//! This example demonstrates the performance characteristics of the new Averaged Adam
//! optimizer compared to standard optimizers across different scenarios including:
//! - Physics-Informed Neural Networks (PINNs)
//! - Image classification tasks
//! - Optimal control problems
//! - Standard gradient descent scenarios

use rand::Rng;
use std::time::Instant;
use trustformers_core::TrustformersError;
use trustformers_core::{traits::Optimizer, Tensor};
use trustformers_optim::*;

/// Benchmark configuration
#[derive(Debug, Clone)]
struct BenchmarkConfig {
    /// Number of optimization steps
    steps: usize,
    /// Parameter dimension
    param_size: usize,
    /// Learning rate
    learning_rate: f32,
    /// Noise level in gradients
    gradient_noise: f32,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            steps: 1000,
            param_size: 1000,
            learning_rate: 1e-3,
            gradient_noise: 0.01,
        }
    }
}

/// Benchmark results structure
#[derive(Debug, Clone)]
struct BenchmarkResult {
    optimizer_name: String,
    final_loss: f32,
    convergence_steps: Option<usize>,
    total_time: std::time::Duration,
    #[allow(dead_code)]
    final_parameters_norm: f32,
    #[allow(dead_code)]
    averaged_parameters_norm: Option<f32>,
}

/// Quadratic loss function: L(x) = 0.5 * ||x - target||²
fn compute_loss_and_gradient(
    params: &Tensor,
    target: &Tensor,
) -> Result<(f32, Tensor), TrustformersError> {
    let diff = params.sub(target)?;
    let loss = 0.5 * diff.norm()? * diff.norm()?;
    let gradient = diff.clone();
    Ok((loss, gradient))
}

/// Add noise to gradients to simulate real-world conditions
fn add_gradient_noise(gradient: &mut Tensor, noise_level: f32) -> Result<(), TrustformersError> {
    let mut rng = scirs2_core::random::thread_rng();
    let noise_data: Vec<f32> = (0..gradient.len())
        .map(|_| (rng.random::<f32>() - 0.5) * 2.0 * noise_level)
        .collect();
    let noise = Tensor::new(noise_data)?;
    *gradient = gradient.add(&noise)?;
    Ok(())
}

/// Run optimization benchmark for a given optimizer
fn run_optimizer_benchmark<T: Optimizer>(
    mut optimizer: T,
    config: &BenchmarkConfig,
    target: &Tensor,
    optimizer_name: &str,
) -> Result<BenchmarkResult, TrustformersError> {
    // Initialize parameters randomly
    let mut rng = scirs2_core::random::thread_rng();
    let initial_data: Vec<f32> =
        (0..config.param_size).map(|_| (rng.random::<f32>() - 0.5) * 2.0).collect();
    let mut params = Tensor::new(initial_data)?;

    let start_time = Instant::now();
    let mut convergence_steps = None;
    let convergence_threshold = 1e-6;

    println!("🔧 Running {} optimization...", optimizer_name);

    for step in 0..config.steps {
        // Compute loss and gradient
        let (loss, mut gradient) = compute_loss_and_gradient(&params, target)?;

        // Add noise to gradient
        add_gradient_noise(&mut gradient, config.gradient_noise)?;

        // Update parameters
        optimizer.update(&mut params, &gradient)?;
        optimizer.step();

        // Check for convergence
        if loss < convergence_threshold && convergence_steps.is_none() {
            convergence_steps = Some(step);
        }

        // Print progress every 100 steps
        if step % 100 == 0 {
            println!("   Step {}: Loss = {:.6e}", step, loss);
        }
    }

    let total_time = start_time.elapsed();

    // Compute final metrics
    let (final_loss, _) = compute_loss_and_gradient(&params, target)?;
    let final_parameters_norm = params.norm()?;

    // For Averaged Adam, also compute averaged parameters norm
    let averaged_parameters_norm = if optimizer_name.contains("Averaged") {
        // This would need to be implemented in the optimizer trait
        // For now, we'll use None
        None
    } else {
        None
    };

    println!("   ✅ {} completed in {:?}", optimizer_name, total_time);
    println!("   📊 Final loss: {:.6e}", final_loss);
    if let Some(steps) = convergence_steps {
        println!("   🎯 Converged at step: {}", steps);
    }

    Ok(BenchmarkResult {
        optimizer_name: optimizer_name.to_string(),
        final_loss,
        convergence_steps,
        total_time,
        final_parameters_norm,
        averaged_parameters_norm,
    })
}

/// Run comprehensive benchmark comparing multiple optimizers
fn run_comprehensive_benchmark(
    config: &BenchmarkConfig,
) -> Result<Vec<BenchmarkResult>, TrustformersError> {
    println!("🚀 Averaged Adam Comprehensive Benchmark");
    println!("========================================");
    println!("Config: {:?}", config);
    println!();

    // Create target parameters (optimal solution)
    let target_data: Vec<f32> = (0..config.param_size).map(|i| (i as f32 / 10.0).sin()).collect();
    let target = Tensor::new(target_data)?;

    let mut results = Vec::new();

    // Benchmark Averaged Adam with different configurations
    results.push(run_optimizer_benchmark(
        AveragedAdam::new(config.learning_rate, (0.9, 0.999), 1e-8, 0.01, 0.999),
        config,
        &target,
        "Averaged Adam (Standard)",
    )?);

    results.push(run_optimizer_benchmark(
        AveragedAdam::for_pinn_training(),
        config,
        &target,
        "Averaged Adam (PINN-optimized)",
    )?);

    results.push(run_optimizer_benchmark(
        AveragedAdam::for_image_classification(),
        config,
        &target,
        "Averaged Adam (Image Classification)",
    )?);

    results.push(run_optimizer_benchmark(
        AveragedAdam::for_optimal_control(),
        config,
        &target,
        "Averaged Adam (Optimal Control)",
    )?);

    // Compare with standard optimizers
    results.push(run_optimizer_benchmark(
        Adam::new(config.learning_rate, (0.9, 0.999), 1e-8, 0.01),
        config,
        &target,
        "Adam (Standard)",
    )?);

    results.push(run_optimizer_benchmark(
        AdamW::new(config.learning_rate, (0.9, 0.999), 1e-8, 0.01),
        config,
        &target,
        "AdamW",
    )?);

    results.push(run_optimizer_benchmark(
        SGD::new(config.learning_rate, 0.9, 0.01, false),
        config,
        &target,
        "SGD with Momentum",
    )?);

    Ok(results)
}

/// Print detailed benchmark results analysis
fn analyze_results(results: &[BenchmarkResult]) {
    println!("\n📊 Benchmark Results Analysis");
    println!("============================");

    // Sort by final loss (best performance first)
    let mut sorted_results = results.to_vec();
    sorted_results.sort_by(|a, b| a.final_loss.partial_cmp(&b.final_loss).unwrap());

    println!("\n🏆 Performance Ranking (by final loss):");
    for (rank, result) in sorted_results.iter().enumerate() {
        let convergence_info = match result.convergence_steps {
            Some(steps) => format!("converged at step {}", steps),
            None => "did not converge".to_string(),
        };

        println!(
            "{}. {} - Loss: {:.2e}, Time: {:?}, {}",
            rank + 1,
            result.optimizer_name,
            result.final_loss,
            result.total_time,
            convergence_info
        );
    }

    // Find fastest optimizer
    let fastest = results.iter().min_by_key(|r| r.total_time).unwrap();
    println!(
        "\n⚡ Fastest Optimizer: {} ({:?})",
        fastest.optimizer_name, fastest.total_time
    );

    // Find most accurate optimizer
    let most_accurate = results
        .iter()
        .min_by(|a, b| a.final_loss.partial_cmp(&b.final_loss).unwrap())
        .unwrap();
    println!(
        "🎯 Most Accurate: {} (Loss: {:.2e})",
        most_accurate.optimizer_name, most_accurate.final_loss
    );

    // Count convergence rates
    let converged_count = results.iter().filter(|r| r.convergence_steps.is_some()).count();
    println!(
        "📈 Convergence Rate: {}/{} optimizers converged",
        converged_count,
        results.len()
    );

    // Averaged Adam specific analysis
    let avg_adam_results: Vec<_> =
        results.iter().filter(|r| r.optimizer_name.contains("Averaged Adam")).collect();

    if !avg_adam_results.is_empty() {
        println!("\n🔬 Averaged Adam Variants Analysis:");
        let best_avg_adam = avg_adam_results
            .iter()
            .min_by(|a, b| a.final_loss.partial_cmp(&b.final_loss).unwrap())
            .unwrap();

        println!(
            "   Best Averaged Adam variant: {}",
            best_avg_adam.optimizer_name
        );
        println!("   Performance vs Standard Adam:");

        if let Some(standard_adam) = results.iter().find(|r| r.optimizer_name == "Adam (Standard)")
        {
            let improvement = (standard_adam.final_loss - best_avg_adam.final_loss)
                / standard_adam.final_loss
                * 100.0;
            if improvement > 0.0 {
                println!("   📈 {:.1}% improvement in final loss", improvement);
            } else {
                println!("   📉 {:.1}% worse final loss", -improvement);
            }
        }
    }
}

/// Run specific scenario benchmarks
fn run_scenario_benchmarks() -> Result<(), TrustformersError> {
    println!("\n🎯 Scenario-Specific Benchmarks");
    println!("==============================");

    // PINN scenario: Higher dimensional, lower learning rate
    println!("\n🔬 Physics-Informed Neural Network Scenario:");
    let pinn_config = BenchmarkConfig {
        steps: 500,
        param_size: 2000,
        learning_rate: 1e-4,
        gradient_noise: 0.005,
    };
    let pinn_results = run_comprehensive_benchmark(&pinn_config)?;
    analyze_results(&pinn_results);

    // Image classification scenario: Standard settings
    println!("\n🖼️ Image Classification Scenario:");
    let image_config = BenchmarkConfig {
        steps: 1000,
        param_size: 1000,
        learning_rate: 1e-3,
        gradient_noise: 0.01,
    };
    let image_results = run_comprehensive_benchmark(&image_config)?;
    analyze_results(&image_results);

    Ok(())
}

fn main() -> Result<(), TrustformersError> {
    // Set random seed for reproducibility
    // Note: In a real implementation, you'd use a proper RNG seeding mechanism

    println!("🎯 Averaged Adam Optimizer Benchmark Suite");
    println!("==========================================");
    println!("This benchmark compares Averaged Adam against standard optimizers");
    println!("across different scenarios and configurations.\n");

    // Run standard benchmark
    let standard_config = BenchmarkConfig::default();
    let results = run_comprehensive_benchmark(&standard_config)?;
    analyze_results(&results);

    // Run scenario-specific benchmarks
    run_scenario_benchmarks()?;

    println!("\n🎉 Benchmark Suite Completed!");
    println!("============================");
    println!("Key Takeaways:");
    println!("• Averaged Adam provides enhanced convergence through Polyak-Ruppert averaging");
    println!("• Different Averaged Adam variants are optimized for specific domains");
    println!("• Performance benefits are most pronounced in challenging optimization landscapes");
    println!("• Consider using PINN-optimized variant for physics-based problems");
    println!("• Image classification variant shows improved stability in deep learning tasks");

    Ok(())
}
