//! Feature Importance Analysis for Enhanced Feature Engineering
//!
//! Analyzes the correlation between each of the 36 features and future returns
//! to determine which features have the most predictive power.
//!
//! # Usage
//!
//! ```bash
//! cargo run -p ml --example feature_importance_analysis --release
//! ```

use anyhow::{Context, Result};
use ml::real_data_loader::{OHLCVBar, RealDataLoader};
use std::collections::HashMap;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

// Import the enhanced technical indicators
use ml_training_service::technical_indicators::{IndicatorConfig, TechnicalIndicatorCalculator};

#[tokio::main]
async fn main() -> Result<()> {
    // Setup logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(tracing::Level::INFO)
        .finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🔍 Feature Importance Analysis - Enhanced Feature Engineering");
    info!("Analyzing 36 features vs baseline 16 features");

    // Load real market data
    let data_loader = RealDataLoader::new();
    let mut file_mapping = HashMap::new();
    file_mapping.insert(
        "6E.FUT".to_string(),
        "test_data/real/databento/ml_training/6E_FUT_20240101_20240131.dbn".to_string(),
    );

    info!("📊 Loading market data for 6E.FUT...");
    let bars = data_loader
        .load_ohlcv_data(&file_mapping)
        .await
        .context("Failed to load OHLCV data")?;

    let total_bars = bars.values().map(|v| v.len()).sum::<usize>();
    info!("   Loaded {} bars across {} symbols", total_bars, bars.len());

    // Calculate features for each symbol
    for (symbol, bar_data) in bars.iter() {
        info!("\n📈 Analyzing {} ({} bars)", symbol, bar_data.len());

        if bar_data.len() < 50 {
            warn!("   Skipping {}: insufficient data", symbol);
            continue;
        }

        // Initialize enhanced indicator calculator
        let config = IndicatorConfig::default();
        let mut calculator = TechnicalIndicatorCalculator::new(symbol.clone(), config);

        // Collect all features and returns
        let mut feature_matrix = Vec::new();
        let mut returns = Vec::new();

        info!("   Computing features and returns...");
        for (i, bar) in bar_data.iter().enumerate() {
            // Update calculator with OHLC data
            calculator.update(
                bar.close,
                bar.volume,
                Some(bar.high),
                Some(bar.low),
            );

            // Skip warmup period
            if !calculator.is_warmed_up() {
                continue;
            }

            // Get all current indicators (36 features)
            let indicators = calculator.current_indicators();

            // Calculate forward return (1-bar ahead)
            if i < bar_data.len() - 1 {
                let forward_return = (bar_data[i + 1].close / bar.close).ln();
                feature_matrix.push(indicators);
                returns.push(forward_return);
            }
        }

        info!("   Collected {} feature vectors", feature_matrix.len());

        if feature_matrix.is_empty() {
            warn!("   No features collected after warmup");
            continue;
        }

        // Calculate feature importance (correlation with returns)
        info!("\n📊 Feature Importance Analysis:");
        info!("   (Pearson correlation with 1-bar forward returns)\n");

        let mut correlations = Vec::new();

        // Get all unique feature names
        let feature_names: Vec<String> = feature_matrix[0].keys().cloned().collect();

        for feature_name in &feature_names {
            let mut feature_values = Vec::new();
            let mut valid_returns = Vec::new();

            // Collect feature values and corresponding returns
            for (features, ret) in feature_matrix.iter().zip(returns.iter()) {
                if let Some(&value) = features.get(feature_name) {
                    if value.is_finite() {
                        feature_values.push(value);
                        valid_returns.push(*ret);
                    }
                }
            }

            if feature_values.len() < 10 {
                continue;
            }

            // Calculate Pearson correlation
            let correlation = calculate_correlation(&feature_values, &valid_returns);
            correlations.push((feature_name.clone(), correlation, feature_values.len()));
        }

        // Sort by absolute correlation (strongest predictive power first)
        correlations.sort_by(|a, b| b.1.abs().partial_cmp(&a.1.abs()).unwrap());

        // Print top 20 features
        info!("   Top 20 Most Predictive Features:");
        info!("   {:<30} {:>12} {:>10}", "Feature", "Correlation", "N");
        info!("   {}", "-".repeat(55));

        for (i, (name, corr, n)) in correlations.iter().take(20).enumerate() {
            let emoji = if i < 10 { "🟢" } else { "🟡" };
            info!("   {:<30} {:>12.6} {:>10}  {}", name, corr, n, emoji);
        }

        // Categorize features
        info!("\n📋 Feature Categories:");

        let momentum_features: Vec<_> = correlations
            .iter()
            .filter(|(name, _, _)| {
                name.contains("rsi")
                    || name.contains("mfi")
                    || name.contains("cmf")
                    || name.contains("chaikin")
                    || name.contains("macd")
            })
            .collect();

        let volatility_features: Vec<_> = correlations
            .iter()
            .filter(|(name, _, _)| {
                name.contains("bollinger")
                    || name.contains("keltner")
                    || name.contains("donchian")
                    || name.contains("atr")
            })
            .collect();

        let volume_features: Vec<_> = correlations
            .iter()
            .filter(|(name, _, _)| {
                name.contains("obv")
                    || name.contains("vwap")
                    || name.contains("volume")
            })
            .collect();

        info!("   Momentum indicators: {} features", momentum_features.len());
        if !momentum_features.is_empty() {
            let avg_corr: f64 = momentum_features.iter().map(|(_, c, _)| c.abs()).sum::<f64>()
                / momentum_features.len() as f64;
            info!("      Average |correlation|: {:.6}", avg_corr);
        }

        info!("   Volatility indicators: {} features", volatility_features.len());
        if !volatility_features.is_empty() {
            let avg_corr: f64 = volatility_features.iter().map(|(_, c, _)| c.abs()).sum::<f64>()
                / volatility_features.len() as f64;
            info!("      Average |correlation|: {:.6}", avg_corr);
        }

        info!("   Volume indicators: {} features", volume_features.len());
        if !volume_features.is_empty() {
            let avg_corr: f64 = volume_features.iter().map(|(_, c, _)| c.abs()).sum::<f64>()
                / volume_features.len() as f64;
            info!("      Average |correlation|: {:.6}", avg_corr);
        }

        // Summary statistics
        info!("\n📈 Summary Statistics:");
        let all_corrs: Vec<f64> = correlations.iter().map(|(_, c, _)| c.abs()).collect();
        let mean_corr = all_corrs.iter().sum::<f64>() / all_corrs.len() as f64;
        let max_corr = all_corrs.iter().cloned().fold(0.0, f64::max);
        let min_corr = all_corrs.iter().cloned().fold(f64::INFINITY, f64::min);

        info!("   Total features: {}", correlations.len());
        info!("   Mean |correlation|: {:.6}", mean_corr);
        info!("   Max |correlation|: {:.6}", max_corr);
        info!("   Min |correlation|: {:.6}", min_corr);

        // Identify new features (enhanced set)
        let new_features: Vec<_> = correlations
            .iter()
            .filter(|(name, _, _)| {
                name.contains("mfi")
                    || name.contains("cmf")
                    || name.contains("chaikin")
                    || name.contains("keltner")
                    || name.contains("donchian")
                    || name.contains("obv")
                    || name.contains("vwap")
                    || name.contains("volume_oscillator")
            })
            .collect();

        info!("\n✨ NEW Features (20 added):");
        info!("   {} new features active", new_features.len());
        if !new_features.is_empty() {
            let new_avg_corr = new_features.iter().map(|(_, c, _)| c.abs()).sum::<f64>()
                / new_features.len() as f64;
            info!("   Average |correlation| of new features: {:.6}", new_avg_corr);

            info!("\n   Top 10 New Features:");
            let mut sorted_new = new_features.clone();
            sorted_new.sort_by(|a, b| b.1.abs().partial_cmp(&a.1.abs()).unwrap());

            for (name, corr, n) in sorted_new.iter().take(10) {
                info!("      {:<30} {:>12.6} {:>10}", name, corr, n);
            }
        }
    }

    info!("\n✅ Feature importance analysis complete!");
    info!("Next steps:");
    info!("  1. Review top predictive features");
    info!("  2. Retrain DQN with enhanced 36-feature set");
    info!("  3. Compare Sharpe ratios (baseline vs enhanced)");

    Ok(())
}

/// Calculate Pearson correlation coefficient between two vectors
fn calculate_correlation(x: &[f64], y: &[f64]) -> f64 {
    if x.len() != y.len() || x.is_empty() {
        return 0.0;
    }

    let n = x.len() as f64;

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

    // Calculate covariance and standard deviations
    let mut cov = 0.0;
    let mut var_x = 0.0;
    let var_y = 0.0;

    for (xi, yi) in x.iter().zip(y.iter()) {
        let dx = xi - mean_x;
        let dy = yi - mean_y;
        cov += dx * dy;
        var_x += dx * dx;
        let var_y = var_y + dy * dy;
    }

    // Avoid division by zero
    if var_x == 0.0 || var_y == 0.0 {
        return 0.0;
    }

    cov / (var_x * var_y).sqrt()
}