//! TFT (Temporal Fusion Transformer) Training with Real DataBento Data
//!
//! Trains a TFT model using real market data from DataBento DBN files with proper
//! static covariates, time-varying features, and multi-horizon forecasting.
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (20 epochs)
//! cargo run -p ml --example train_tft_dbn --release --features cuda
//!
//! # Custom configuration
//! cargo run -p ml --example train_tft_dbn --release --features cuda -- \
//!   --epochs 50 \
//!   --batch-size 32 \
//!   --lookback 60 \
//!   --horizon 10
//! ```

use anyhow::{Context, Result};
use chrono::{DateTime, Datelike, TimeZone, Timelike, Utc};
use clap::Parser;
use dbn::decode::{DbnDecoder, DecodeRecordRef};
use dbn::OhlcvMsg;
use ndarray::{Array1, Array2};
use std::path::PathBuf;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::checkpoint::FileSystemStorage;
use ml::data_loaders::BarSamplingMethod;
use ml::features::config::FeatureConfig;
use ml::features::extraction::{extract_ml_features, OHLCVBar as ExtractorBar};
use ml::tft::training::TFTDataLoader;
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};

#[derive(Debug, Parser)]
#[command(
    name = "train_tft_dbn",
    about = "Train TFT model on real DataBento data"
)]
struct Opts {
    /// DBN file path (or directory containing multiple DBN files)
    #[arg(
        long,
        default_value = "test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn"
    )]
    data_path: String,

    /// Number of training epochs
    #[arg(long, default_value = "20")]
    epochs: usize,

    /// Learning rate
    #[arg(long, default_value = "0.001")]
    learning_rate: f64,

    /// Batch size (max 32 for 4GB VRAM)
    #[arg(long, default_value = "32")]
    batch_size: usize,

    /// Hidden dimension
    #[arg(long, default_value = "256")]
    hidden_dim: usize,

    /// Number of attention heads
    #[arg(long, default_value = "8")]
    num_attention_heads: usize,

    /// Lookback window
    #[arg(long, default_value = "60")]
    lookback_window: usize,

    /// Forecast horizon
    #[arg(long, default_value = "10")]
    forecast_horizon: usize,

    /// Training/validation split (0.0-1.0)
    #[arg(long, default_value = "0.8")]
    train_split: f64,

    /// Output directory for trained model
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Early stopping patience (epochs without improvement)
    #[arg(long, default_value = "20")]
    early_stopping_patience: usize,

    /// Early stopping threshold (minimum improvement)
    #[arg(long, default_value = "0.0001")]
    early_stopping_threshold: f64,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,

    /// Alternative bar sampling method (time, tick, volume, dollar, imbalance, run)
    #[arg(long, default_value = "time")]
    bar_method: String,

    /// Bar sampling threshold (tick count, volume, dollar value, imbalance, or run length)
    #[arg(long)]
    bar_threshold: Option<f64>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let opts = Opts::parse();

    // Setup logging
    let level = if opts.verbose {
        tracing::Level::DEBUG
    } else {
        tracing::Level::INFO
    };

    let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    info!("🚀 Starting TFT Training with Real DataBento Data");
    // Initialize Wave D feature configuration (225 features)
    let feature_config = FeatureConfig::wave_d();

    info!("Configuration:");
    // Validate feature count matches expected 225-dimensional input
    let total_features = feature_config.feature_count();
    assert_eq!(
        total_features, 225,
        "Feature config must provide exactly 225 features for Wave D (201 Wave C + 24 Wave D)"
    );
    info!("  • Data path: {}", opts.data_path);
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Hidden dimension: {}", opts.hidden_dim);
    info!("  • Attention heads: {}", opts.num_attention_heads);
    info!("  • Lookback window: {}", opts.lookback_window);
    info!("  • Forecast horizon: {}", opts.forecast_horizon);
    info!(
        "  • Feature count: {} (Wave D: Wave C 201 + Wave D 24)",
        total_features
    );
    info!(
        "  • Train/val split: {:.1}%/{:.1}%",
        opts.train_split * 100.0,
        (1.0 - opts.train_split) * 100.0
    );
    info!("  • GPU: CUDA MANDATORY (no CPU fallback)");
    info!(
        "  • Early stopping patience: {} epochs",
        opts.early_stopping_patience
    );
    info!(
        "  • Early stopping threshold: {:.2e}",
        opts.early_stopping_threshold
    );
    info!("  • Output directory: {}", opts.output_dir);
    info!("  • Bar sampling method: {}", opts.bar_method);
    if let Some(threshold) = opts.bar_threshold {
        info!("  • Bar threshold: {}", threshold);
    }

    // Create output directory
    let output_path = PathBuf::from(&opts.output_dir);
    if !output_path.exists() {
        std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
        info!("✅ Created output directory: {}", opts.output_dir);
    }

    // Configure alternative bar sampling (Wave B)
    let bar_sampling = match opts.bar_method.as_str() {
        "tick" => BarSamplingMethod::TickBars(opts.bar_threshold.unwrap_or(100.0) as usize),
        "volume" => BarSamplingMethod::VolumeBars(opts.bar_threshold.unwrap_or(10000.0)),
        "dollar" => BarSamplingMethod::DollarBars(opts.bar_threshold.unwrap_or(2_000_000.0)),
        "imbalance" => BarSamplingMethod::ImbalanceBars(opts.bar_threshold.unwrap_or(1000.0)),
        "run" => BarSamplingMethod::RunBars(opts.bar_threshold.unwrap_or(50.0) as usize),
        _ => BarSamplingMethod::TimeBars,
    };

    info!("✅ Bar sampling configured: {:?}", bar_sampling);

    // Load real market data from DBN files
    info!("\n📊 Loading real market data from DataBento...");

    // Note: load_dbn_ohlcv_bars will need to support alternative bar sampling
    // This requires extending the function to accept bar_sampling parameter
    // For now, it loads time-based bars

    // Check if path is a file or directory
    let path = std::path::Path::new(&opts.data_path);
    let bars = if path.is_dir() {
        // Load all .dbn files from directory
        info!("Loading DBN files from directory: {}", opts.data_path);
        let mut all_bars = Vec::new();

        for entry in std::fs::read_dir(path)? {
            let entry = entry?;
            let file_path = entry.path();

            if file_path.extension().and_then(|s| s.to_str()) == Some("dbn") {
                info!("  • Loading: {:?}", file_path.file_name().unwrap());
                let file_bars = load_dbn_ohlcv_bars(file_path.to_str().unwrap())
                    .await
                    .context(format!("Failed to load DBN file: {:?}", file_path))?;
                all_bars.extend(file_bars);
            }
        }

        // Sort by timestamp
        all_bars.sort_by_key(|b| b.timestamp);
        all_bars
    } else {
        // Load single file
        load_dbn_ohlcv_bars(&opts.data_path)
            .await
            .context("Failed to load DBN data")?
    };

    info!("✅ Loaded {} OHLCV bars from DataBento", bars.len());

    // Convert to TFT data structure with Wave D features (225)
    info!(
        "\n🔄 Converting to TFT data format with {} features...",
        total_features
    );
    let tft_data = convert_to_tft_data(
        &bars,
        opts.lookback_window,
        opts.forecast_horizon,
        &feature_config,
    )
    .context("Failed to convert to TFT format")?;

    info!("✅ Created {} TFT samples", tft_data.len());

    // Split into train/validation
    let split_idx = (tft_data.len() as f64 * opts.train_split) as usize;
    let train_data = tft_data[..split_idx].to_vec();
    let val_data = tft_data[split_idx..].to_vec();

    info!(
        "✅ Split: {} training, {} validation samples",
        train_data.len(),
        val_data.len()
    );

    // Create data loaders
    let train_loader = TFTDataLoader::new(train_data, opts.batch_size, true);
    let val_loader = TFTDataLoader::new(val_data, opts.batch_size, false);

    // Configure TFT trainer
    // Configure TFT trainer with 225 features (Wave D)
    // Static features: 10 (symbol metadata, volatility, liquidity)
    // Historical features: 225 (Wave C 201 + Wave D 24)
    // Future features: 10 (calendar features)
    let trainer_config = TFTTrainerConfig {
        epochs: opts.epochs,
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        hidden_dim: opts.hidden_dim,
        num_attention_heads: opts.num_attention_heads,
        dropout_rate: 0.1,
        lstm_layers: 2,
        quantiles: vec![0.1, 0.5, 0.9],
        lookback_window: opts.lookback_window,
        forecast_horizon: opts.forecast_horizon,
        use_gpu: true, // CUDA always required
        use_int8_quantization: false,
        use_qat: false,
        qat_calibration_batches: 100,
        checkpoint_dir: opts.output_dir.clone(),
    };

    // Create checkpoint storage
    let storage = std::sync::Arc::new(FileSystemStorage::new(output_path.clone()));

    // Create TFT trainer
    let mut trainer =
        TFTTrainer::new(trainer_config.clone(), storage).context("Failed to create TFT trainer")?;

    info!("✅ TFT trainer initialized");

    // Setup progress callback
    let (progress_tx, mut progress_rx) = mpsc::unbounded_channel();
    trainer.set_progress_callback(progress_tx);

    // Spawn progress monitor task
    let monitor_task = tokio::spawn(async move {
        while let Some(progress) = progress_rx.recv().await {
            info!("{}", progress.message);
            if let Some(loss) = progress.metrics.get("train_loss") {
                info!("  • Train loss: {:.6}", loss);
            }
            if let Some(val_loss) = progress.metrics.get("val_loss") {
                info!("  • Val loss: {:.6}", val_loss);
            }
            if let Some(quantile_loss) = progress.metrics.get("quantile_loss") {
                info!("  • Quantile loss: {:.6}", quantile_loss);
            }
            if let Some(rmse) = progress.metrics.get("rmse") {
                info!("  • RMSE: {:.6}", rmse);
            }
        }
    });

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let final_metrics = trainer
        .train(train_loader, val_loader)
        .await
        .context("Training failed")?;

    let training_duration = start_time.elapsed();

    // Wait for progress monitor to finish
    drop(trainer); // Drop trainer to close progress channel
    let _ = monitor_task.await;

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Training loss: {:.6}", final_metrics.train_loss);
    info!("  • Validation loss: {:.6}", final_metrics.val_loss);
    info!("  • Quantile loss: {:.6}", final_metrics.quantile_loss);
    info!("  • RMSE: {:.6}", final_metrics.rmse);
    info!(
        "  • Attention entropy: {:.4}",
        final_metrics.attention_entropy
    );
    info!(
        "  • Training duration: {:.1}s ({:.1} min)",
        training_duration.as_secs_f64(),
        training_duration.as_secs_f64() / 60.0
    );
    info!(
        "  • Training time: {:.1}s ({:.1} min)",
        final_metrics.training_time_seconds,
        final_metrics.training_time_seconds / 60.0
    );

    info!("\n💾 Model checkpoints saved to: {}", opts.output_dir);
    info!("\n🎉 TFT training with real DataBento data complete!");

    Ok(())
}

/// OHLCV bar structure (intermediate format)
#[derive(Debug, Clone)]
struct OhlcvBar {
    timestamp: DateTime<Utc>,
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    volume: f64,
}

/// Load OHLCV bars from DBN file with price anomaly correction
async fn load_dbn_ohlcv_bars(file_path: &str) -> Result<Vec<OhlcvBar>> {
    debug!("Loading DBN file: {}", file_path);

    let mut decoder = DbnDecoder::from_file(file_path).context(format!(
        "Failed to create DBN decoder for file: {}",
        file_path
    ))?;

    let mut bars = Vec::new();
    let mut prev_close: Option<f64> = None;
    let mut corrections_applied = 0;

    while let Some(record_ref) = decoder
        .decode_record_ref()
        .context("Failed to decode DBN record")?
    {
        if let Some(ohlcv) = record_ref.get::<OhlcvMsg>() {
            // Convert timestamp
            let ts_nanos = ohlcv.hd.ts_event as i64;
            let secs = ts_nanos / 1_000_000_000;
            let nanos = (ts_nanos % 1_000_000_000) as u32;
            let timestamp = Utc
                .timestamp_opt(secs, nanos)
                .single()
                .ok_or_else(|| anyhow::anyhow!("Invalid timestamp: {}", ts_nanos))?;

            // Convert prices (DBN uses 9 decimal places)
            let mut open_f64 = ohlcv.open as f64 / 1_000_000_000.0;
            let mut high_f64 = ohlcv.high as f64 / 1_000_000_000.0;
            let mut low_f64 = ohlcv.low as f64 / 1_000_000_000.0;
            let mut close_f64 = ohlcv.close as f64 / 1_000_000_000.0;

            // Price anomaly detection (same logic as backtesting service)
            if let Some(prev) = prev_close {
                let pct_change = ((close_f64 - prev) / prev).abs();

                if pct_change > 0.5 && close_f64 < 1000.0 {
                    let corrected_close = close_f64 * 100.0;

                    if corrected_close >= 3000.0 && corrected_close <= 6000.0 {
                        open_f64 *= 100.0;
                        high_f64 *= 100.0;
                        low_f64 *= 100.0;
                        close_f64 = corrected_close;
                        corrections_applied += 1;

                        if corrections_applied <= 5 {
                            debug!(
                                "Applied 100x price correction at bar {} ({}% change)",
                                bars.len() + 1,
                                pct_change * 100.0
                            );
                        }
                    } else {
                        warn!(
                            "Skipping corrupted bar at index {} (timestamp: {})",
                            bars.len() + 1,
                            timestamp
                        );
                        prev_close = Some(prev);
                        continue;
                    }
                }
            }

            prev_close = Some(close_f64);

            let bar = OhlcvBar {
                timestamp,
                open: open_f64,
                high: high_f64,
                low: low_f64,
                close: close_f64,
                volume: ohlcv.volume as f64,
            };

            bars.push(bar);
        }
    }

    if corrections_applied > 0 {
        info!(
            "Applied {} automatic price corrections for encoding inconsistencies",
            corrections_applied
        );
    }

    Ok(bars)
}

/// Convert OHLCV bars to TFT data structure using production feature extraction
///
/// ## Wave D Feature Set (225 total):
/// Uses `ml::features::extraction::extract_ml_features()` for all 225 features:
/// ### Wave C Base (indices 0-200): 201 features
/// - OHLCV: 5 features
/// - Technical Indicators: 21 features (SMA, EMA, RSI, MACD, Bollinger, ATR)
/// - Microstructure: 3 features (Roll, Amihud, Corwin-Schultz)
/// - Statistical Features: 172 features (quantiles, correlations, rolling stats)
///
/// ### Wave D Regime Detection (indices 201-224): 24 features
/// - CUSUM Statistics (201-210): 10 features
/// - ADX & Directional Indicators (211-215): 5 features
/// - Regime Transition Probabilities (216-220): 5 features
/// - Adaptive Strategy Metrics (221-224): 4 features
fn convert_to_tft_data(
    bars: &[OhlcvBar],
    lookback_window: usize,
    forecast_horizon: usize,
    _feature_config: &FeatureConfig,
) -> Result<Vec<(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)>> {
    if bars.len() < lookback_window + forecast_horizon {
        return Err(anyhow::anyhow!(
            "Not enough data: need {} bars, got {}",
            lookback_window + forecast_horizon,
            bars.len()
        ));
    }

    // Convert OhlcvBar to ExtractorBar for feature extraction
    let extractor_bars: Vec<ExtractorBar> = bars
        .iter()
        .map(|b| ExtractorBar {
            timestamp: b.timestamp,
            open: b.open,
            high: b.high,
            low: b.low,
            close: b.close,
            volume: b.volume,
        })
        .collect();

    // Extract 225-dimensional features using production pipeline
    info!("🔍 Extracting 225-dim features via production pipeline...");
    let feature_vectors =
        extract_ml_features(&extractor_bars).context("Failed to extract ML features")?;

    info!(
        "✅ Extracted {} feature vectors (225-dim each)",
        feature_vectors.len()
    );

    // Calculate statistics for static features and normalization
    let prices: Vec<f64> = bars.iter().map(|b| b.close).collect();
    let mean_price = prices.iter().sum::<f64>() / prices.len() as f64;
    let price_std =
        (prices.iter().map(|p| (p - mean_price).powi(2)).sum::<f64>() / prices.len() as f64).sqrt();

    let volumes: Vec<f64> = bars.iter().map(|b| b.volume).collect();
    let mean_volume = volumes.iter().sum::<f64>() / volumes.len() as f64;
    let volume_std = (volumes
        .iter()
        .map(|v| (v - mean_volume).powi(2))
        .sum::<f64>()
        / volumes.len() as f64)
        .sqrt();

    let mut tft_samples = Vec::new();

    // Note: feature_vectors starts AFTER warmup period (50 bars)
    // So feature_vectors[0] corresponds to bars[50]
    const WARMUP_PERIOD: usize = 50;

    // Create sliding windows from feature vectors
    for i in 0..feature_vectors.len() {
        // Check if we have enough data for this window
        if i + lookback_window + forecast_horizon > feature_vectors.len() {
            break;
        }

        // Calculate the bar index in the original bars array
        let bar_idx = WARMUP_PERIOD + i;
        let first_bar = &bars[bar_idx];

        // Static features: Symbol metadata (10 features)
        let hour = first_bar.timestamp.hour() as f64;
        let day_of_week = first_bar.timestamp.weekday().num_days_from_monday() as f64;
        let is_morning = if hour < 12.0 { 1.0 } else { 0.0 };
        let is_afternoon = if hour >= 12.0 && hour < 17.0 {
            1.0
        } else {
            0.0
        };

        // Calculate volatility over lookback window
        let lookback_slice = &bars[bar_idx..bar_idx + lookback_window];
        let returns: Vec<f64> = lookback_slice
            .windows(2)
            .map(|w| (w[1].close / w[0].close).ln())
            .collect();
        let volatility = if returns.len() > 1 {
            let mean_return = returns.iter().sum::<f64>() / returns.len() as f64;
            (returns
                .iter()
                .map(|r| (r - mean_return).powi(2))
                .sum::<f64>()
                / returns.len() as f64)
                .sqrt()
        } else {
            0.01
        };

        let liquidity = mean_volume / mean_price;

        let static_features = Array1::from_vec(vec![
            mean_price / 5000.0,
            price_std / 100.0,
            mean_volume / 1000.0,
            volume_std / 1000.0,
            hour / 24.0,
            day_of_week / 7.0,
            is_morning,
            is_afternoon,
            volatility * 100.0,
            liquidity / 100.0,
        ]);

        // Historical features: Use production-extracted 225-dim feature vectors
        let mut hist_features = Vec::with_capacity(lookback_window * 225);
        for t in 0..lookback_window {
            let feature_vec = &feature_vectors[i + t];
            hist_features.extend_from_slice(feature_vec);
        }

        let historical_features = Array2::from_shape_vec((lookback_window, 225), hist_features)?;

        // Future features: Known future events (10 features per timestep)
        let mut fut_features = Vec::with_capacity(forecast_horizon * 10);

        for t in 0..forecast_horizon {
            let future_bar = &bars[bar_idx + lookback_window + t];
            let fut_hour = future_bar.timestamp.hour() as f64;
            let fut_day = future_bar.timestamp.weekday().num_days_from_monday() as f64;
            let is_weekend = if fut_day >= 5.0 { 1.0 } else { 0.0 };
            let fut_is_morning = if fut_hour < 12.0 { 1.0 } else { 0.0 };
            let fut_is_afternoon = if fut_hour >= 12.0 && fut_hour < 17.0 {
                1.0
            } else {
                0.0
            };
            let week_of_month = ((future_bar.timestamp.day() - 1) / 7) as f64;
            let month = future_bar.timestamp.month() as f64;
            let quarter = ((month - 1.0) / 3.0).floor();
            let is_month_start = if future_bar.timestamp.day() <= 5 {
                1.0
            } else {
                0.0
            };
            let is_month_end = if future_bar.timestamp.day() >= 25 {
                1.0
            } else {
                0.0
            };

            fut_features.extend(vec![
                fut_hour / 24.0,
                fut_day / 7.0,
                is_weekend,
                fut_is_morning,
                fut_is_afternoon,
                week_of_month / 4.0,
                month / 12.0,
                quarter / 4.0,
                is_month_start,
                is_month_end,
            ]);
        }

        let future_features = Array2::from_shape_vec((forecast_horizon, 10), fut_features)?;

        // Targets: Multi-horizon price forecast (normalized)
        let targets: Vec<f64> = (0..forecast_horizon)
            .map(|t| bars[bar_idx + lookback_window + t].close / mean_price)
            .collect();

        let target_array = Array1::from_vec(targets);

        tft_samples.push((
            static_features,
            historical_features,
            future_features,
            target_array,
        ));
    }

    Ok(tft_samples)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_load_dbn_bars() {
        let dbn_file = "test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn";

        if !std::path::Path::new(dbn_file).exists() {
            eprintln!("Skipping test: DBN file not found");
            return;
        }

        let result = load_dbn_ohlcv_bars(dbn_file).await;
        assert!(
            result.is_ok(),
            "Failed to load DBN bars: {:?}",
            result.err()
        );

        let bars = result.unwrap();
        assert!(!bars.is_empty(), "Should load bars");
        println!("✅ Loaded {} bars", bars.len());
    }

    #[tokio::test]
    async fn test_convert_to_tft_format() {
        let dbn_file = "test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn";

        if !std::path::Path::new(dbn_file).exists() {
            eprintln!("Skipping test: DBN file not found");
            return;
        }

        let bars = load_dbn_ohlcv_bars(dbn_file).await.unwrap();
        let feature_config = FeatureConfig::wave_d();
        let tft_data = convert_to_tft_data(&bars, 60, 10, &feature_config).unwrap();

        assert!(!tft_data.is_empty(), "Should create TFT samples");

        let (static_feat, hist_feat, fut_feat, targets) = &tft_data[0];

        // Verify shapes
        assert_eq!(
            static_feat.len(),
            10,
            "Static features should have 10 dimensions"
        );
        assert_eq!(
            hist_feat.shape(),
            &[60, 225],
            "Historical features should be [60, 225] (Wave D)"
        );
        assert_eq!(
            fut_feat.shape(),
            &[10, 10],
            "Future features should be [10, 10]"
        );
        assert_eq!(targets.len(), 10, "Targets should have 10 timesteps");

        println!("✅ TFT data structure validated:");
        println!("  • Static features: {:?}", static_feat.shape());
        println!("  • Historical features: {:?}", hist_feat.shape());
        println!("  • Future features: {:?}", fut_feat.shape());
        println!("  • Targets: {:?}", targets.shape());
    }
}