foxhunt/ml/examples/ensemble_visualization.rs

//! Ensemble Visualization Tool
//!
//! Generates visualizations for 6-model ensemble:
//! - Weight evolution over time
//! - Correlation heatmap
//! - Performance attribution

use anyhow::Result;
use ml::ensemble::coordinator_extended::{
    EnsembleConfig, ExtendedEnsembleCoordinator, WeightSnapshot,
};
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;

/// Generate CSV data for weight evolution plot
pub fn export_weight_evolution_csv(snapshots: &[WeightSnapshot], output_path: &str) -> Result<()> {
    let mut file = File::create(output_path)?;

    // Header
    writeln!(
        file,
        "Timestamp,DQN,PPO,TFT,MAMBA-2,Liquid,TLOB,Ensemble_Sharpe"
    )?;

    // Data rows
    for snapshot in snapshots {
        write!(file, "{}", snapshot.timestamp)?;
        for model in &["DQN", "PPO", "TFT", "MAMBA-2", "Liquid", "TLOB"] {
            let weight = snapshot.weights.get(*model).copied().unwrap_or(0.0);
            write!(file, ",{:.6}", weight)?;
        }
        writeln!(file, ",{:.6}", snapshot.ensemble_sharpe)?;
    }

    println!("✅ Weight evolution data exported to: {}", output_path);
    println!("   Import into Excel/Google Sheets for visualization");

    Ok(())
}

/// Generate CSV data for correlation heatmap
pub fn export_correlation_heatmap_csv(
    heatmap_data: &[(String, String, f64)],
    output_path: &str,
) -> Result<()> {
    let mut file = File::create(output_path)?;

    let models = vec!["DQN", "PPO", "TFT", "MAMBA-2", "Liquid", "TLOB"];

    // Create correlation matrix
    let mut matrix: HashMap<(String, String), f64> = HashMap::new();

    for (model_a, model_b, corr) in heatmap_data {
        matrix.insert((model_a.clone(), model_b.clone()), *corr);
    }

    // Header
    write!(file, "Model")?;
    for model in &models {
        write!(file, ",{}", model)?;
    }
    writeln!(file)?;

    // Data rows
    for model_a in &models {
        write!(file, "{}", model_a)?;
        for model_b in &models {
            if model_a == model_b {
                write!(file, ",1.000")?;
            } else {
                let corr = matrix
                    .get(&(model_a.to_string(), model_b.to_string()))
                    .copied()
                    .unwrap_or(0.0);
                write!(file, ",{:.3}", corr)?;
            }
        }
        writeln!(file)?;
    }

    println!("✅ Correlation heatmap exported to: {}", output_path);
    println!("   Import into Excel/Google Sheets for visualization");

    Ok(())
}

/// Generate Python plotting script
pub fn generate_python_plot_script(output_path: &str) -> Result<()> {
    let script = r#"#!/usr/bin/env python3
"""
Ensemble Visualization Script
Generate weight evolution and correlation heatmap plots
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

def plot_weight_evolution(csv_path='weight_evolution.csv', output_path='weight_evolution.png'):
    """Plot model weight evolution over time"""
    df = pd.read_csv(csv_path)

    models = ['DQN', 'PPO', 'TFT', 'MAMBA-2', 'Liquid', 'TLOB']

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))

    # Plot 1: Stacked area chart of weights
    ax1.stackplot(df.index,
                  *[df[model] for model in models],
                  labels=models,
                  alpha=0.7)
    ax1.set_xlabel('Prediction Number', fontsize=12)
    ax1.set_ylabel('Weight', fontsize=12)
    ax1.set_title('Model Weight Evolution (6-Model Ensemble)', fontsize=14, fontweight='bold')
    ax1.legend(loc='upper left', bbox_to_anchor=(1, 1), fontsize=10)
    ax1.grid(True, alpha=0.3)
    ax1.set_ylim([0, 1])

    # Plot 2: Ensemble Sharpe ratio over time
    ax2.plot(df.index, df['Ensemble_Sharpe'], color='darkblue', linewidth=2, label='Ensemble Sharpe')
    ax2.axhline(y=0, color='red', linestyle='--', alpha=0.5, label='Zero Line')
    ax2.set_xlabel('Prediction Number', fontsize=12)
    ax2.set_ylabel('Sharpe Ratio', fontsize=12)
    ax2.set_title('Ensemble Sharpe Ratio Evolution', fontsize=14, fontweight='bold')
    ax2.legend(loc='best', fontsize=10)
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"✅ Weight evolution plot saved to: {output_path}")
    plt.close()

def plot_correlation_heatmap(csv_path='correlation_heatmap.csv', output_path='correlation_heatmap.png'):
    """Plot correlation heatmap"""
    df = pd.read_csv(csv_path, index_col=0)

    fig, ax = plt.subplots(figsize=(10, 8))

    # Create heatmap with annotations
    sns.heatmap(df,
                annot=True,
                fmt='.3f',
                cmap='RdYlGn',
                center=0,
                vmin=-1,
                vmax=1,
                square=True,
                linewidths=1,
                cbar_kws={'label': 'Correlation Coefficient'},
                ax=ax)

    ax.set_title('Model Prediction Correlation Matrix', fontsize=14, fontweight='bold', pad=20)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"✅ Correlation heatmap saved to: {output_path}")
    plt.close()

def plot_performance_attribution(sharpe_ratios, win_rates, output_path='performance_attribution.png'):
    """Plot performance attribution"""
    models = list(sharpe_ratios.keys())
    sharpes = [sharpe_ratios[m] for m in models]
    wins = [win_rates[m] * 100 for m in models]  # Convert to percentage

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

    # Plot 1: Sharpe ratios
    colors1 = ['green' if s > 0 else 'red' for s in sharpes]
    bars1 = ax1.barh(models, sharpes, color=colors1, alpha=0.7, edgecolor='black')
    ax1.axvline(x=0, color='black', linestyle='-', linewidth=0.8)
    ax1.set_xlabel('Sharpe Ratio', fontsize=12)
    ax1.set_title('Model Sharpe Ratios', fontsize=14, fontweight='bold')
    ax1.grid(True, alpha=0.3, axis='x')

    # Add value labels
    for i, (bar, val) in enumerate(zip(bars1, sharpes)):
        ax1.text(val + 0.05 if val > 0 else val - 0.05, i,
                f'{val:.2f}', va='center', ha='left' if val > 0 else 'right',
                fontsize=10, fontweight='bold')

    # Plot 2: Win rates
    colors2 = ['green' if w >= 50 else 'red' for w in wins]
    bars2 = ax2.barh(models, wins, color=colors2, alpha=0.7, edgecolor='black')
    ax2.axvline(x=50, color='black', linestyle='--', linewidth=0.8, label='50% Baseline')
    ax2.set_xlabel('Win Rate (%)', fontsize=12)
    ax2.set_title('Model Win Rates', fontsize=14, fontweight='bold')
    ax2.set_xlim([0, 100])
    ax2.grid(True, alpha=0.3, axis='x')
    ax2.legend()

    # Add value labels
    for i, (bar, val) in enumerate(zip(bars2, wins)):
        ax2.text(val + 1, i, f'{val:.1f}%', va='center', ha='left',
                fontsize=10, fontweight='bold')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"✅ Performance attribution plot saved to: {output_path}")
    plt.close()

if __name__ == '__main__':
    print("🎨 Generating ensemble visualizations...")
    print()

    # Generate plots
    try:
        plot_weight_evolution()
        plot_correlation_heatmap()

        # Example performance data (replace with actual data from ensemble)
        sharpe_ratios = {
            'DQN': 2.31,
            'PPO': 1.85,
            'MAMBA-2': 1.92,
            'TFT': 1.45,
            'Liquid': 1.38,
            'TLOB': 1.56
        }

        win_rates = {
            'DQN': 0.58,
            'PPO': 0.56,
            'MAMBA-2': 0.57,
            'TFT': 0.54,
            'Liquid': 0.53,
            'TLOB': 0.55
        }

        plot_performance_attribution(sharpe_ratios, win_rates)

        print()
        print("✅ All visualizations generated successfully!")
        print()
        print("📊 Generated files:")
        print("   - weight_evolution.png: Model weight changes over time")
        print("   - correlation_heatmap.png: Model prediction correlations")
        print("   - performance_attribution.png: Individual model performance")

    except FileNotFoundError as e:
        print(f"❌ Error: CSV file not found - {e}")
        print("   Run the six_model_ensemble example first to generate CSV data")
    except Exception as e:
        print(f"❌ Error generating plots: {e}")
"#;

    let mut file = File::create(output_path)?;
    file.write_all(script.as_bytes())?;

    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let mut perms = std::fs::metadata(output_path)?.permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(output_path, perms)?;
    }

    println!("✅ Python visualization script generated: {}", output_path);
    println!("   Run with: python3 {}", output_path);

    Ok(())
}

/// Comprehensive visualization export
pub async fn export_all_visualizations(
    coordinator: &ExtendedEnsembleCoordinator,
    output_dir: &str,
) -> Result<()> {
    // Create output directory
    std::fs::create_dir_all(output_dir)?;

    println!("📊 Exporting ensemble visualizations to: {}", output_dir);
    println!();

    // Export weight evolution
    let weight_history = coordinator.get_weight_history().await;
    let weight_csv_path = format!("{}/weight_evolution.csv", output_dir);
    export_weight_evolution_csv(&weight_history, &weight_csv_path)?;

    // Export correlation heatmap
    let heatmap_data = coordinator.get_correlation_heatmap().await;
    let heatmap_csv_path = format!("{}/correlation_heatmap.csv", output_dir);
    export_correlation_heatmap_csv(&heatmap_data, &heatmap_csv_path)?;

    // Generate Python plotting script
    let plot_script_path = format!("{}/generate_plots.py", output_dir);
    generate_python_plot_script(&plot_script_path)?;

    println!();
    println!("✅ All data exported successfully!");
    println!();
    println!("📊 Next steps:");
    println!("   1. cd {}", output_dir);
    println!("   2. pip install pandas matplotlib seaborn numpy");
    println!("   3. python3 generate_plots.py");
    println!();
    println!("   This will generate:");
    println!("   - weight_evolution.png");
    println!("   - correlation_heatmap.png");
    println!("   - performance_attribution.png");

    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    println!("🎨 Ensemble Visualization Tool");
    println!("=".repeat(80));
    println!();

    // Example usage (normally called from the main ensemble example)
    println!("This tool is designed to be imported and used from the ensemble example.");
    println!();
    println!("Usage:");
    println!("  use ml::examples::ensemble_visualization::export_all_visualizations;");
    println!("  export_all_visualizations(&coordinator, \"./ensemble_viz\").await?;");
    println!();
    println!("To test visualization generation:");
    println!("  cargo run --example six_model_ensemble");

    Ok(())
}