//! Six-Model Ensemble Testing Example
//!
//! This example demonstrates the extended ensemble coordinator with all 6 models:
//! DQN, PPO, TFT, MAMBA-2, Liquid, TLOB
//!
//! Tests:
//! - Load best checkpoints for each model
//! - Run 1000 predictions
//! - Measure ensemble Sharpe vs individual Sharpe
//! - Generate performance attribution report
//! - Create correlation heatmap visualization

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

/// Simulated model predictions (in production, these would load real checkpoints)
struct MockModelPredictor {
    model_id: String,
    base_sharpe: f64,
    correlation_factor: f64,
}

impl MockModelPredictor {
    fn new(model_id: String, base_sharpe: f64, correlation_factor: f64) -> Self {
        Self {
            model_id,
            base_sharpe,
            correlation_factor,
        }
    }

    fn predict(&self, features: &Features, market_signal: f64, i: usize) -> ModelPrediction {
        // Simulate model prediction with noise and correlation
        // Use simple deterministic noise based on iteration for reproducibility
        let noise = ((i as f64 * 0.618033988749895).fract() - 0.5) * 0.2;
        let value = market_signal * self.correlation_factor + noise;

        // Confidence based on base Sharpe (higher Sharpe = more confident)
        let confidence = ((self.base_sharpe / 2.0).max(0.5).min(1.0) + 0.2).min(1.0);

        ModelPrediction::new(self.model_id.clone(), value, confidence)
    }
}

/// 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 {
        // Annualize: 252 days, 6.5 hours, predictions every minute
        let annualization_factor = (252.0 * 6.5 * 60.0_f64).sqrt();
        (mean_return / std_dev) * annualization_factor
    }
}

#[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!("🚀 Starting 6-Model Ensemble Test");
    info!("=" .repeat(80));

    // Create ensemble coordinator with adaptive weighting
    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 all 6 models with equal initial weights
    info!("📋 Registering 6 models...");
    coordinator
        .register_model("DQN".to_string(), 0.167)
        .await?;
    coordinator
        .register_model("PPO".to_string(), 0.167)
        .await?;
    coordinator
        .register_model("TFT".to_string(), 0.167)
        .await?;
    coordinator
        .register_model("MAMBA-2".to_string(), 0.167)
        .await?;
    coordinator
        .register_model("Liquid".to_string(), 0.167)
        .await?;
    coordinator
        .register_model("TLOB".to_string(), 0.165)
        .await?;

    info!("✅ All 6 models registered");

    // Create mock model predictors with different characteristics
    // Based on Agent 78 DQN results: DQN epoch 30 has Sharpe 2.31
    let models = vec![
        MockModelPredictor::new("DQN".to_string(), 2.31, 0.8),        // High Sharpe, high correlation
        MockModelPredictor::new("PPO".to_string(), 1.85, 0.75),       // Good Sharpe, moderate correlation
        MockModelPredictor::new("TFT".to_string(), 1.45, 0.6),        // Moderate Sharpe, lower correlation
        MockModelPredictor::new("MAMBA-2".to_string(), 1.92, 0.7),    // Good Sharpe, moderate correlation
        MockModelPredictor::new("Liquid".to_string(), 1.38, 0.5),     // Lower Sharpe, low correlation (diversity)
        MockModelPredictor::new("TLOB".to_string(), 1.56, 0.55),      // Moderate Sharpe, low correlation
    ];

    // Simulate 1000 predictions
    info!("🔄 Running 1000 predictions...");
    let mut ensemble_returns = Vec::new();
    let mut individual_returns: HashMap<String, Vec<f64>> = HashMap::new();

    let start_time = Instant::now();

    for i in 0..1000 {
        // Create synthetic features
        let features = Features::new(
            vec![0.5; 16], // 16 features
            (0..16).map(|i| format!("feature_{}", i)).collect(),
        );

        // Generate market signal (random walk with trend)
        // Use deterministic signal based on iteration for reproducibility
        let market_signal = ((i as f64 * 0.314159265359).fract() - 0.48) * 0.02; // Slight positive bias

        // Get predictions from all models
        let predictions: Vec<ModelPrediction> = models
            .iter()
            .map(|model| model.predict(&features, market_signal, i))
            .collect();

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

        // Simulate trading outcome based on ensemble signal
        let ensemble_return = if decision.signal > 0.1 {
            market_signal * 0.95 // 95% capture of positive moves
        } else if decision.signal < -0.1 {
            -market_signal * 0.95 // Short on negative signals
        } else {
            0.0 // No trade on weak signals
        };

        ensemble_returns.push(ensemble_return);

        // Record individual model returns (for performance tracking)
        for pred in predictions {
            let model_return = if pred.value > 0.1 {
                market_signal * 0.9 // Individual models are slightly less efficient
            } else if pred.value < -0.1 {
                -market_signal * 0.9
            } else {
                0.0
            };

            individual_returns
                .entry(pred.model_id.clone())
                .or_insert_with(Vec::new)
                .push(model_return);

            // Record outcome for adaptive weighting
            coordinator
                .record_outcome(&pred.model_id, model_return)
                .await?;
        }

        if (i + 1) % 200 == 0 {
            info!("  Completed {} predictions", i + 1);
        }
    }

    let elapsed = start_time.elapsed();
    info!("✅ Completed 1000 predictions in {:.2}s", elapsed.as_secs_f64());
    info!("  Average latency: {:.0}μs per prediction", elapsed.as_micros() as f64 / 1000.0);

    // Calculate performance metrics
    info!("");
    info!("📊 PERFORMANCE RESULTS");
    info!("=" .repeat(80));

    let ensemble_sharpe = calculate_sharpe_ratio(&ensemble_returns);
    info!("🎯 Ensemble Sharpe Ratio: {:.3}", ensemble_sharpe);
    info!("");

    info!("📈 Individual Model Sharpe Ratios:");
    let mut individual_sharpes: Vec<(String, f64)> = individual_returns
        .iter()
        .map(|(model, returns)| {
            let sharpe = calculate_sharpe_ratio(returns);
            (model.clone(), sharpe)
        })
        .collect();

    individual_sharpes.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    for (model, sharpe) in &individual_sharpes {
        let improvement = if *sharpe > 0.0 {
            ((ensemble_sharpe / sharpe - 1.0) * 100.0)
        } else {
            0.0
        };

        info!(
            "  {:<12} Sharpe: {:>6.3}  (Ensemble improvement: {:>+5.1}%)",
            model, sharpe, improvement
        );
    }

    let best_individual_sharpe = individual_sharpes.first().map(|(_, s)| *s).unwrap_or(0.0);
    let ensemble_improvement = if best_individual_sharpe > 0.0 {
        ((ensemble_sharpe / best_individual_sharpe - 1.0) * 100.0)
    } else {
        0.0
    };

    info!("");
    info!("🏆 Ensemble vs Best Individual: {:>+.1}%", ensemble_improvement);

    // Get final weights
    info!("");
    info!("⚖️  FINAL MODEL WEIGHTS (After Adaptive Adjustment)");
    info!("=" .repeat(80));

    let weights = coordinator.get_weights().await;
    let mut weight_vec: Vec<(String, f64)> = weights.into_iter().collect();
    weight_vec.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    for (model, weight) in weight_vec {
        info!("  {:<12} Weight: {:.3} ({:.1}%)", model, weight, weight * 100.0);
    }

    // Get diversity metrics
    info!("");
    info!("🔀 DIVERSITY METRICS");
    info!("=" .repeat(80));

    let diversity = coordinator.get_diversity_metrics().await;
    info!("  Model Count: {}", diversity.model_count);
    info!("  Average Correlation: {:.3}", diversity.avg_correlation);
    info!("  Average Disagreement: {:.1}%", diversity.avg_disagreement * 100.0);

    // Correlation heatmap (text representation)
    info!("");
    info!("📊 CORRELATION HEATMAP");
    info!("=" .repeat(80));

    let heatmap = coordinator.get_correlation_heatmap().await;
    let model_names = vec!["DQN", "PPO", "TFT", "MAMBA-2", "Liquid", "TLOB"];

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

    // Print matrix
    for i in 0..model_names.len() {
        print!("{:<10}", model_names[i]);
        for j in 0..model_names.len() {
            if i == j {
                print!("   1.000 ");
            } else {
                let corr = heatmap
                    .iter()
                    .find(|(a, b, _)| {
                        a == model_names[i] && b == model_names[j]
                    })
                    .map(|(_, _, c)| *c)
                    .unwrap_or(0.0);

                print!("{:>8.3} ", corr);
            }
        }
        println!();
    }

    // Get performance attribution
    info!("");
    info!("🎯 PERFORMANCE ATTRIBUTION");
    info!("=" .repeat(80));

    let attribution = coordinator.get_performance_attribution().await;
    info!("  Total Predictions: {}", attribution.total_predictions);
    info!("");

    let mut perf_vec: Vec<_> = attribution.model_performance.into_iter().collect();
    perf_vec.sort_by(|a, b| b.1.sharpe_ratio.partial_cmp(&a.1.sharpe_ratio).unwrap());

    for (model, perf) in perf_vec {
        info!(
            "  {:<12} Sharpe: {:>6.3}  Win Rate: {:>5.1}%  Predictions: {}",
            model,
            perf.sharpe_ratio,
            perf.win_rate * 100.0,
            perf.prediction_count
        );
    }

    // Summary
    info!("");
    info!("=" .repeat(80));
    info!("✅ TEST COMPLETE");
    info!("");

    if ensemble_improvement >= 15.0 {
        info!("🎉 EXCELLENT: Ensemble achieved {:.1}% improvement over best individual model!", ensemble_improvement);
        info!("   Target: 15-30% improvement ✅");
    } else if ensemble_improvement >= 10.0 {
        info!("✅ GOOD: Ensemble achieved {:.1}% improvement over best individual model", ensemble_improvement);
        info!("   Target: 15-30% improvement (close!)");
    } else {
        info!("⚠️  BELOW TARGET: Ensemble achieved {:.1}% improvement", ensemble_improvement);
        info!("   Target: 15-30% improvement");
        info!("   Consider adjusting diversity_adjustment_factor or min_correlation_threshold");
    }

    info!("");
    info!("📁 Next steps:");
    info!("  1. Load real checkpoints: DQN epoch 30, PPO epoch 380, TFT best checkpoint");
    info!("  2. Test on real market data (ES.FUT, NQ.FUT)");
    info!("  3. Generate weight evolution plots");
    info!("  4. Implement live model swapping");

    Ok(())
}