//! Model Diversity and Correlation Analysis
//!
//! This example computes the prediction correlation matrix for all available models
//! and identifies the optimal ensemble composition based on diversity and Sharpe ratio.
//!
//! Analysis includes:
//! - 6x6 correlation matrix (DQN, PPO, TFT, MAMBA-2, Liquid, TLOB)
//! - Diversity score (average pairwise correlation)
//! - Ensemble size testing (3, 4, 5, 6 models)
//! - Sharpe ratio vs latency tradeoff
//! - Optimal composition recommendation

use anyhow::Result;
use ml::ensemble::coordinator_extended::{ExtendedEnsembleCoordinator, EnsembleConfig};
use ml::ModelPrediction;
use std::collections::HashMap;
use std::time::Instant;
use tracing::{info, Level};
use tracing_subscriber::FmtSubscriber;

/// Model characteristics based on training analysis
#[derive(Clone)]
struct ModelCharacteristics {
    name: String,
    sharpe: f64,           // Expected Sharpe ratio
    correlation: f64,      // Correlation with market signal
    latency_us: f64,       // Inference latency in microseconds
}

impl ModelCharacteristics {
    fn new(name: &str, sharpe: f64, correlation: f64, latency_us: f64) -> Self {
        Self {
            name: name.to_string(),
            sharpe,
            correlation,
            latency_us,
        }
    }

    /// Generate prediction based on market signal
    fn predict(&self, market_signal: f64, noise: f64) -> f64 {
        market_signal * self.correlation + noise
    }

    /// Get confidence from Sharpe ratio
    fn confidence(&self) -> f64 {
        ((self.sharpe / 2.5).max(0.5).min(1.0) + 0.2).min(1.0)
    }
}

/// Calculate Pearson correlation coefficient
fn pearson_correlation(x: &[f64], y: &[f64]) -> f64 {
    let n = x.len().min(y.len());
    if n < 2 {
        return 0.0;
    }

    let mean_x: f64 = x.iter().sum::<f64>() / n as f64;
    let mean_y: f64 = y.iter().sum::<f64>() / n as f64;

    let mut numerator = 0.0;
    let mut sum_sq_x = 0.0;
    let mut sum_sq_y = 0.0;

    for i in 0..n {
        let dx = x[i] - mean_x;
        let dy = y[i] - mean_y;
        numerator += dx * dy;
        sum_sq_x += dx * dx;
        sum_sq_y += dy * dy;
    }

    let denominator = (sum_sq_x * sum_sq_y).sqrt();
    if denominator < 1e-10 {
        0.0
    } else {
        numerator / denominator
    }
}

/// Calculate Sharpe ratio from returns
fn calculate_sharpe_ratio(returns: &[f64]) -> f64 {
    if returns.len() < 2 {
        return 0.0;
    }

    let mean_return = returns.iter().sum::<f64>() / returns.len() as f64;
    let variance = returns
        .iter()
        .map(|r| (r - mean_return).powi(2))
        .sum::<f64>()
        / returns.len() as f64;
    let std_dev = variance.sqrt();

    if std_dev < 1e-10 {
        0.0
    } else {
        let annualization_factor = (252.0 * 6.5 * 60.0_f64).sqrt();
        (mean_return / std_dev) * annualization_factor
    }
}

/// Test ensemble with specific model combination
async fn test_ensemble_combination(
    models: &[&ModelCharacteristics],
    num_predictions: usize,
) -> Result<(f64, f64, f64)> {
    let config = EnsembleConfig {
        adaptive_weighting: true,
        min_correlation_threshold: 0.7,
        diversity_adjustment_factor: 0.2,
        performance_window_size: 1000,
        min_weight: 0.05,
        max_weight: 0.40,
    };

    let coordinator = ExtendedEnsembleCoordinator::new(config);

    // Register models with equal weights
    let weight = 1.0 / models.len() as f64;
    for model in models {
        coordinator
            .register_model(model.name.clone(), weight)
            .await?;
    }

    let mut ensemble_returns = Vec::new();
    let mut total_latency_us = 0.0;

    let start = Instant::now();

    for i in 0..num_predictions {
        // Generate market signal
        let market_signal = ((i as f64 * 0.314159265359).fract() - 0.48) * 0.02;

        // Generate noise for each model
        let noise_base = (i as f64 * 0.618033988749895).fract() - 0.5;

        // Create predictions
        let predictions: Vec<ModelPrediction> = models
            .iter()
            .enumerate()
            .map(|(idx, model)| {
                let noise = (noise_base + idx as f64 * 0.1) * 0.2;
                let value = model.predict(market_signal, noise);
                ModelPrediction::new(model.name.clone(), value, model.confidence())
            })
            .collect();

        // Make ensemble decision
        let decision = coordinator.predict(predictions.clone()).await?;

        // Simulate trading outcome
        let ensemble_return = if decision.signal > 0.1 {
            market_signal * 0.95
        } else if decision.signal < -0.1 {
            -market_signal * 0.95
        } else {
            0.0
        };

        ensemble_returns.push(ensemble_return);

        // Record individual outcomes
        for pred in predictions {
            let model_return = if pred.value > 0.1 {
                market_signal * 0.9
            } else if pred.value < -0.1 {
                -market_signal * 0.9
            } else {
                0.0
            };

            coordinator
                .record_outcome(&pred.model_id, model_return)
                .await?;
        }

        // Accumulate latency
        for model in models {
            total_latency_us += model.latency_us;
        }
    }

    let _elapsed = start.elapsed();
    let avg_latency_us = total_latency_us / num_predictions as f64;
    let sharpe = calculate_sharpe_ratio(&ensemble_returns);

    // Get diversity metrics
    let diversity = coordinator.get_diversity_metrics().await;

    Ok((sharpe, avg_latency_us, diversity.avg_correlation))
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    info!("🔬 MODEL DIVERSITY AND CORRELATION ANALYSIS");
    info!("=" .repeat(80));
    info!("");

    // Define model characteristics based on training results
    // Source: Agent 78 (DQN), Agent 54 (PPO), checkpoint analysis reports
    let models = vec![
        ModelCharacteristics::new("DQN", 2.31, 0.80, 15.0),      // Agent 78 epoch 30: Sharpe 2.31
        ModelCharacteristics::new("PPO", 1.85, 0.75, 18.0),      // PPO epoch 380: Sharpe ~1.85
        ModelCharacteristics::new("TFT", 1.45, 0.60, 25.0),      // Estimated: Moderate Sharpe
        ModelCharacteristics::new("MAMBA-2", 1.92, 0.70, 20.0),  // Estimated: Good Sharpe
        ModelCharacteristics::new("Liquid", 1.38, 0.50, 12.0),   // Estimated: Lower Sharpe, fast
        ModelCharacteristics::new("TLOB", 1.56, 0.55, 8.0),      // Estimated: Moderate, very fast
    ];

    // Step 1: Compute 6x6 correlation matrix
    info!("📊 STEP 1: Computing 6x6 Correlation Matrix");
    info!("-" .repeat(80));

    let num_samples = 1000;
    let mut prediction_history: HashMap<String, Vec<f64>> = HashMap::new();

    for i in 0..num_samples {
        let market_signal = ((i as f64 * 0.314159265359).fract() - 0.48) * 0.02;
        let noise_base = (i as f64 * 0.618033988749895).fract() - 0.5;

        for (idx, model) in models.iter().enumerate() {
            let noise = (noise_base + idx as f64 * 0.1) * 0.2;
            let value = model.predict(market_signal, noise);

            prediction_history
                .entry(model.name.clone())
                .or_insert_with(Vec::new)
                .push(value);
        }
    }

    // Calculate correlation matrix
    info!("");
    info!("Correlation Matrix (6x6):");
    info!("");

    // Print header
    print!("          ");
    for model in &models {
        print!("{:>8} ", model.name);
    }
    println!();

    // Print matrix
    let mut total_correlation = 0.0;
    let mut correlation_count = 0;

    for i in 0..models.len() {
        print!("{:<10}", models[i].name);
        for j in 0..models.len() {
            if i == j {
                print!("   1.000 ");
            } else {
                let hist_i = prediction_history.get(&models[i].name).unwrap();
                let hist_j = prediction_history.get(&models[j].name).unwrap();
                let corr = pearson_correlation(hist_i, hist_j);
                print!("{:>8.3} ", corr);

                if i < j {
                    total_correlation += corr.abs();
                    correlation_count += 1;
                }
            }
        }
        println!();
    }

    let avg_correlation = total_correlation / correlation_count as f64;

    info!("");
    info!("Average Pairwise Correlation: {:.3}", avg_correlation);
    info!("Diversity Score: {:.3} (1 - avg_corr)", 1.0 - avg_correlation);
    info!("");

    // Step 2: Test different ensemble sizes
    info!("📈 STEP 2: Testing Ensemble Sizes (3, 4, 5, 6 models)");
    info!("-" .repeat(80));
    info!("");

    let test_predictions = 1000;

    // Test 3-model ensembles (all combinations)
    info!("Testing 3-Model Ensembles:");
    let three_model_combos = vec![
        vec![&models[0], &models[1], &models[2]], // DQN, PPO, TFT
        vec![&models[0], &models[1], &models[3]], // DQN, PPO, MAMBA-2
        vec![&models[0], &models[2], &models[4]], // DQN, TFT, Liquid
        vec![&models[1], &models[3], &models[5]], // PPO, MAMBA-2, TLOB
    ];

    let mut best_3_sharpe = 0.0;
    let mut best_3_combo = String::new();
    let mut best_3_latency = 0.0;

    for combo in &three_model_combos {
        let names: Vec<String> = combo.iter().map(|m| m.name.clone()).collect();
        let (sharpe, latency, diversity) = test_ensemble_combination(combo, test_predictions).await?;

        info!(
            "  {} | Sharpe: {:>6.3} | Latency: {:>5.0}μs | Diversity: {:.3}",
            names.join(", "),
            sharpe,
            latency,
            1.0 - diversity
        );

        if sharpe > best_3_sharpe {
            best_3_sharpe = sharpe;
            best_3_combo = names.join(", ");
            best_3_latency = latency;
        }
    }

    info!("");

    // Test 4-model ensemble
    info!("Testing 4-Model Ensembles:");
    let four_model_combos = vec![
        vec![&models[0], &models[1], &models[3], &models[5]], // DQN, PPO, MAMBA-2, TLOB
        vec![&models[0], &models[1], &models[2], &models[4]], // DQN, PPO, TFT, Liquid
    ];

    let mut best_4_sharpe = 0.0;
    let mut best_4_combo = String::new();
    let mut best_4_latency = 0.0;

    for combo in &four_model_combos {
        let names: Vec<String> = combo.iter().map(|m| m.name.clone()).collect();
        let (sharpe, latency, diversity) = test_ensemble_combination(combo, test_predictions).await?;

        info!(
            "  {} | Sharpe: {:>6.3} | Latency: {:>5.0}μs | Diversity: {:.3}",
            names.join(", "),
            sharpe,
            latency,
            1.0 - diversity
        );

        if sharpe > best_4_sharpe {
            best_4_sharpe = sharpe;
            best_4_combo = names.join(", ");
            best_4_latency = latency;
        }
    }

    info!("");

    // Test 5-model ensemble
    info!("Testing 5-Model Ensemble:");
    let five_model_combo = vec![&models[0], &models[1], &models[2], &models[3], &models[5]];
    let names: Vec<String> = five_model_combo.iter().map(|m| m.name.clone()).collect();
    let (sharpe_5, latency_5, diversity_5) =
        test_ensemble_combination(&five_model_combo, test_predictions).await?;

    info!(
        "  {} | Sharpe: {:>6.3} | Latency: {:>5.0}μs | Diversity: {:.3}",
        names.join(", "),
        sharpe_5,
        latency_5,
        1.0 - diversity_5
    );

    info!("");

    // Test 6-model ensemble
    info!("Testing 6-Model Ensemble:");
    let six_model_combo: Vec<&ModelCharacteristics> = models.iter().collect();
    let names: Vec<String> = six_model_combo.iter().map(|m| m.name.clone()).collect();
    let (sharpe_6, latency_6, diversity_6) =
        test_ensemble_combination(&six_model_combo, test_predictions).await?;

    info!(
        "  {} | Sharpe: {:>6.3} | Latency: {:>5.0}μs | Diversity: {:.3}",
        names.join(", "),
        sharpe_6,
        latency_6,
        1.0 - diversity_6
    );

    info!("");

    // Step 3: Sharpe vs Latency Tradeoff Analysis
    info!("⚖️  STEP 3: Sharpe Ratio vs Latency Tradeoff");
    info!("-" .repeat(80));
    info!("");

    let latency_budget_us = 50.0;

    info!("Latency Budget: {:.0}μs (HFT requirement)", latency_budget_us);
    info!("");

    let results = vec![
        ("3-model (best)", best_3_sharpe, best_3_latency),
        ("4-model (best)", best_4_sharpe, best_4_latency),
        ("5-model", sharpe_5, latency_5),
        ("6-model", sharpe_6, latency_6),
    ];

    info!("{:<20} {:>10} {:>12} {:>15}", "Ensemble", "Sharpe", "Latency (μs)", "Within Budget?");
    info!("-" .repeat(60));

    for (name, sharpe, latency) in &results {
        let within_budget = if *latency <= latency_budget_us {
            "✅ YES"
        } else {
            "❌ NO"
        };

        info!(
            "{:<20} {:>10.3} {:>12.0} {:>15}",
            name, sharpe, latency, within_budget
        );
    }

    info!("");

    // Step 4: Optimal Composition Recommendation
    info!("🎯 STEP 4: Optimal Ensemble Composition");
    info!("-" .repeat(80));
    info!("");

    // Find best combination within latency budget
    let mut best_sharpe = 0.0;
    let mut best_ensemble = String::new();
    let mut best_latency = 0.0;
    let mut best_size = 0;

    for (name, sharpe, latency) in &results {
        if *latency <= latency_budget_us && *sharpe > best_sharpe {
            best_sharpe = *sharpe;
            best_ensemble = name.to_string();
            best_latency = *latency;

            if name.contains("3-model") {
                best_size = 3;
            } else if name.contains("4-model") {
                best_size = 4;
            } else if name.contains("5-model") {
                best_size = 5;
            } else {
                best_size = 6;
            }
        }
    }

    if best_size > 0 {
        info!("✅ RECOMMENDED ENSEMBLE");
        info!("");
        info!("  Configuration: {}", best_ensemble);
        info!("  Models: {}", if best_size == 3 {
            best_3_combo.as_str()
        } else if best_size == 4 {
            best_4_combo.as_str()
        } else if best_size == 5 {
            "DQN, PPO, TFT, MAMBA-2, TLOB"
        } else {
            "All 6 models"
        });
        info!("  Expected Sharpe: {:.3}", best_sharpe);
        info!("  Average Latency: {:.0}μs", best_latency);
        info!("  Latency Budget: {:.0}μs ({}% utilized)", latency_budget_us, (best_latency / latency_budget_us * 100.0));
        info!("");

        if best_sharpe > 1.5 {
            info!("  🎉 Excellent! Sharpe > 1.5 (production-ready)");
        } else if best_sharpe > 1.2 {
            info!("  ✅ Good! Sharpe > 1.2 (viable for production)");
        } else {
            info!("  ⚠️  Moderate Sharpe. Consider further optimization.");
        }
    } else {
        info!("⚠️  WARNING: No ensemble meets latency budget of {:.0}μs", latency_budget_us);
        info!("");
        info!("Recommendations:");
        info!("  1. Increase latency budget to {:.0}μs", best_3_latency);
        info!("  2. Optimize model inference (GPU acceleration)");
        info!("  3. Use sequential ensemble (fast models first)");
    }

    info!("");

    // Step 5: Model-Specific Recommendations
    info!("📋 STEP 5: Model-Specific Recommendations");
    info!("-" .repeat(80));
    info!("");

    info!("Individual Model Performance:");
    info!("");

    let mut sorted_models = models.clone();
    sorted_models.sort_by(|a, b| b.sharpe.partial_cmp(&a.sharpe).unwrap());

    for model in &sorted_models {
        let rank = if model.sharpe > 2.0 {
            "⭐⭐⭐ (Excellent)"
        } else if model.sharpe > 1.5 {
            "⭐⭐ (Good)"
        } else {
            "⭐ (Moderate)"
        };

        info!(
            "  {:<12} Sharpe: {:>6.3} | Latency: {:>5.0}μs | Correlation: {:.2} | {}",
            model.name, model.sharpe, model.latency_us, model.correlation, rank
        );
    }

    info!("");
    info!("Model Selection Criteria:");
    info!("  ✅ Include: DQN (highest Sharpe: 2.31)");
    info!("  ✅ Include: MAMBA-2 (good Sharpe: 1.92, diverse)");
    info!("  ✅ Include: PPO (solid Sharpe: 1.85, stable)");
    info!("  ⚠️  Consider: TLOB (fastest: 8μs, moderate Sharpe)");
    info!("  ⚠️  Consider: TFT (diverse: 0.60 correlation)");
    info!("  ⚠️  Optional: Liquid (diversity benefit, low correlation)");

    info!("");
    info!("=" .repeat(80));
    info!("✅ ANALYSIS COMPLETE");
    info!("");

    // Summary metrics
    info!("📊 SUMMARY METRICS");
    info!("");
    info!("  Average Model Correlation: {:.3}", avg_correlation);
    info!("  Best 3-Model Sharpe: {:.3}", best_3_sharpe);
    info!("  Best 4-Model Sharpe: {:.3}", best_4_sharpe);
    info!("  5-Model Sharpe: {:.3}", sharpe_5);
    info!("  6-Model Sharpe: {:.3}", sharpe_6);
    info!("  Optimal Ensemble Size: {} models", best_size);
    info!("");

    // Improvement vs best individual
    let best_individual_sharpe = sorted_models[0].sharpe;
    let improvement = (best_sharpe / best_individual_sharpe - 1.0) * 100.0;

    if improvement > 15.0 {
        info!("  🎉 Ensemble improvement: +{:.1}% over best individual", improvement);
        info!("     (Exceeds 15% target - excellent diversification benefit)");
    } else if improvement > 10.0 {
        info!("  ✅ Ensemble improvement: +{:.1}% over best individual", improvement);
        info!("     (Good diversification benefit, approaching 15% target)");
    } else {
        info!("  ⚠️  Ensemble improvement: +{:.1}% over best individual", improvement);
        info!("     (Below 10% target - consider adjusting model weights)");
    }

    info!("");

    Ok(())
}