foxhunt/crates/ml/examples/test_ensemble.rs

//! Test Ensemble Coordinator with Mock Models
//!
//! This example demonstrates the ensemble coordinator functionality by:
//! 1. Registering 3 models (DQN, PPO, TFT) with equal weights
//! 2. Running ensemble predictions on 100 sample data points
//! 3. Analyzing ensemble signals, per-model votes, and disagreement rate
//!
//! ## Usage
//!
//! ```bash
//! cargo run -p ml --example test_ensemble --release
//! ```
//!
//! ## Expected Output
//!
//! - Ensemble predictions with Buy/Sell/Hold signals
//! - Per-model vote breakdown
//! - Confidence scores and disagreement analysis
//! - Summary statistics for ensemble behavior

use anyhow::Result;
use ml::ensemble::{EnsembleCoordinator, TradingAction};
use ml::Features;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("=== Ensemble Coordinator Test ===\n");

    // Create ensemble coordinator
    let coordinator = EnsembleCoordinator::new();

    // Register 3 models with equal weights
    println!("Registering models...");
    coordinator.register_model("DQN".to_string(), 0.33).await?;
    coordinator.register_model("PPO".to_string(), 0.33).await?;
    coordinator.register_model("TFT".to_string(), 0.34).await?;
    println!("✓ Registered 3 models: DQN, PPO, TFT\n");

    // Generate 100 sample feature vectors
    println!("Running ensemble predictions on 100 samples...\n");

    let mut results = EnsembleResults::new();

    for i in 0..100 {
        // Generate diverse feature values to test different market conditions
        let feature_vals = generate_sample_features(i);

        let features = Features::new(
            feature_vals,
            vec![
                "close_price".to_string(),
                "volume".to_string(),
                "rsi".to_string(),
                "macd".to_string(),
                "volatility".to_string(),
            ],
        );

        // Make ensemble prediction
        let decision = coordinator.predict(&features).await?;

        // Track results
        results.record_decision(&decision);

        // Print first 5 predictions for debugging
        if i < 5 {
            println!("Sample {}:", i + 1);
            println!("  Action: {:?}", decision.action);
            println!("  Signal: {:.3}", decision.signal);
            println!("  Confidence: {:.3}", decision.confidence);
            println!("  Disagreement: {:.1}%", decision.disagreement_rate * 100.0);
            println!("  Model Votes:");
            for (model_id, vote) in &decision.model_votes {
                println!(
                    "    {}: signal={:.3}, confidence={:.3}, weight={:.3}",
                    model_id, vote.signal, vote.confidence, vote.weight
                );
            }
            println!();
        }
    }

    // Print summary
    println!("\n=== Ensemble Summary (100 predictions) ===\n");

    println!("Action Distribution:");
    println!(
        "  Buy:  {} ({:.1}%)",
        results.buy_count,
        results.buy_count as f64 / 100.0 * 100.0
    );
    println!(
        "  Sell: {} ({:.1}%)",
        results.sell_count,
        results.sell_count as f64 / 100.0 * 100.0
    );
    println!(
        "  Hold: {} ({:.1}%)",
        results.hold_count,
        results.hold_count as f64 / 100.0 * 100.0
    );

    println!("\nConfidence Statistics:");
    println!("  Average: {:.3}", results.avg_confidence());
    println!("  Min: {:.3}", results.min_confidence);
    println!("  Max: {:.3}", results.max_confidence);

    println!("\nDisagreement Analysis:");
    println!("  Average: {:.1}%", results.avg_disagreement() * 100.0);
    println!("  Max: {:.1}%", results.max_disagreement * 100.0);
    println!(
        "  High Disagreement (>50%): {} predictions",
        results.high_disagreement_count
    );

    println!("\nSignal Statistics:");
    println!("  Average: {:.3}", results.avg_signal());
    println!("  Min: {:.3}", results.min_signal);
    println!("  Max: {:.3}", results.max_signal);

    println!("\n✓ Ensemble test completed successfully!");

    Ok(())
}

/// Generate sample features with varying market conditions
fn generate_sample_features(index: usize) -> Vec<f64> {
    // Create diverse market scenarios
    let phase = (index as f64 * 0.1).sin();

    // Feature 1: Close price (normalized)
    let close = 0.5 + phase * 0.3;

    // Feature 2: Volume (normalized)
    let volume = 0.6 + (index as f64 * 0.05).cos() * 0.2;

    // Feature 3: RSI (0-1 range)
    let rsi = 0.5 + phase * 0.4;

    // Feature 4: MACD (normalized)
    let macd = phase * 0.5;

    // Feature 5: Volatility (0-1 range)
    let volatility = 0.3 + (index as f64 * 0.08).sin().abs() * 0.3;

    vec![close, volume, rsi, macd, volatility]
}

/// Track ensemble prediction results
struct EnsembleResults {
    buy_count: usize,
    sell_count: usize,
    hold_count: usize,

    confidence_sum: f64,
    min_confidence: f64,
    max_confidence: f64,

    disagreement_sum: f64,
    max_disagreement: f64,
    high_disagreement_count: usize,

    signal_sum: f64,
    min_signal: f64,
    max_signal: f64,

    total_predictions: usize,
}

impl EnsembleResults {
    fn new() -> Self {
        Self {
            buy_count: 0,
            sell_count: 0,
            hold_count: 0,
            confidence_sum: 0.0,
            min_confidence: 1.0,
            max_confidence: 0.0,
            disagreement_sum: 0.0,
            max_disagreement: 0.0,
            high_disagreement_count: 0,
            signal_sum: 0.0,
            min_signal: 1.0,
            max_signal: -1.0,
            total_predictions: 0,
        }
    }

    fn record_decision(&mut self, decision: &ml::ensemble::EnsembleDecision) {
        // Track action
        match decision.action {
            TradingAction::Buy => self.buy_count += 1,
            TradingAction::Sell => self.sell_count += 1,
            TradingAction::Hold => self.hold_count += 1,
        }

        // Track confidence
        self.confidence_sum += decision.confidence;
        self.min_confidence = self.min_confidence.min(decision.confidence);
        self.max_confidence = self.max_confidence.max(decision.confidence);

        // Track disagreement
        self.disagreement_sum += decision.disagreement_rate;
        self.max_disagreement = self.max_disagreement.max(decision.disagreement_rate);
        if decision.is_high_disagreement() {
            self.high_disagreement_count += 1;
        }

        // Track signal
        self.signal_sum += decision.signal;
        self.min_signal = self.min_signal.min(decision.signal);
        self.max_signal = self.max_signal.max(decision.signal);

        self.total_predictions += 1;
    }

    fn avg_confidence(&self) -> f64 {
        if self.total_predictions > 0 {
            self.confidence_sum / self.total_predictions as f64
        } else {
            0.0
        }
    }

    fn avg_disagreement(&self) -> f64 {
        if self.total_predictions > 0 {
            self.disagreement_sum / self.total_predictions as f64
        } else {
            0.0
        }
    }

    fn avg_signal(&self) -> f64 {
        if self.total_predictions > 0 {
            self.signal_sum / self.total_predictions as f64
        } else {
            0.0
        }
    }
}