foxhunt/ml/examples/train_tft_parquet.rs

//! TFT (Temporal Fusion Transformer) Training with Parquet Data
//!
//! Trains a TFT model using market data from Parquet files with lazy batch loading
//! to avoid OOM issues on large datasets. Uses the TFTParquetExt trait for efficient
//! memory management and 225-feature extraction (Wave C + Wave D).
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (20 epochs)
//! cargo run -p ml --example train_tft_parquet --release --features cuda
//!
//! # Custom configuration
//! cargo run -p ml --example train_tft_parquet --release --features cuda -- \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --epochs 50 \
//!   --batch-size 32 \
//!   --lookback-window 60 \
//!   --forecast-horizon 10
//! ```
//!
//! # Features
//!
//! - Lazy batch loading (10,000 rows at a time) to avoid OOM crashes
//! - 225-feature extraction (Wave C 201 + Wave D 24) from OHLCV bars
//! - Sliding window creation (configurable lookback/horizon)
//! - GPU-accelerated training (RTX 3050 Ti, 4GB VRAM)
//! - Automatic train/validation split (80/20)
//! - Model checkpointing and early stopping
//!
//! # Parquet Schema Requirements
//!
//! The Parquet file must follow Databento schema:
//! - Column 3: open (Float64)
//! - Column 4: high (Float64)
//! - Column 5: low (Float64)
//! - Column 6: close (Float64)
//! - Column 7: volume (UInt64)
//! - Column 9: ts_event (Timestamp[ns, UTC])

// Suppress warnings for unused dependencies in this example
// (examples have access to all crate dependencies but typically only use a subset)
#![allow(unused_crate_dependencies)]

// Use mimalloc allocator for 10-25% performance improvement
#[cfg(feature = "mimalloc-allocator")]
use mimalloc::MiMalloc;
#[cfg(feature = "mimalloc-allocator")]
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use tokio::sync::mpsc;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::checkpoint::FileSystemStorage;
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};

#[derive(Debug, Parser)]
#[command(
    name = "train_tft_parquet",
    about = "Train TFT model on Parquet market data with lazy loading"
)]
struct Opts {
    /// Parquet file path containing OHLCV bars (Databento schema)
    #[arg(long, default_value = "test_data/ES_FUT_small.parquet")]
    parquet_file: String,

    /// Number of training epochs
    #[arg(long, default_value = "3")]
    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,

    /// Validation batch size (defaults to match training batch_size)
    #[arg(long)]
    validation_batch_size: Option<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 (historical bars)
    #[arg(long, default_value = "60")]
    lookback_window: usize,

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

    /// Dropout rate for regularization
    #[arg(long, default_value = "0.1")]
    dropout_rate: f64,

    /// Number of LSTM layers
    #[arg(long, default_value = "2")]
    lstm_layers: usize,

    /// Quantiles for probabilistic forecasting (comma-separated)
    #[arg(long, default_value = "0.1,0.5,0.9")]
    quantiles: String,

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

    /// Use GPU for training (CUDA required)
    #[arg(long)]
    use_gpu: bool,

    /// Use INT8 quantization for memory efficiency (reduces VRAM usage by 3-8x)
    #[arg(long)]
    use_int8: bool,

    /// Use Quantization-Aware Training (1-2% better accuracy than PTQ)
    /// Trains with fake quantization, converts to INT8 at the end
    #[arg(long)]
    use_qat: bool,

    /// Number of batches for QAT calibration (default: 100)
    /// Higher values improve accuracy but increase training time
    #[arg(long, default_value = "100")]
    qat_calibration_batches: usize,

    /// Minimum batch size for QAT calibration OOM recovery (default: 2)
    /// If OOM occurs, batch size is halved automatically. Training aborts if below this threshold.
    #[arg(long, default_value = "2")]
    qat_min_batch_size: usize,

    /// ⚠️ WARNING: Gradient checkpointing NOT IMPLEMENTED for QAT (see QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md)
    ///
    /// This flag is IGNORED when --use-qat is enabled. For QAT memory reduction:
    /// 1. Use 2-phase workaround (calibrate → freeze observers → train with checkpointing)
    /// 2. See ml/docs/QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md for details
    ///
    /// For non-QAT training: Reduces GPU memory usage by 30-40% but increases training time by ~20%
    #[arg(long)]
    use_gradient_checkpointing: bool,

    /// Auto-detect optimal batch size based on available GPU memory
    /// Overrides --batch-size if enabled. Prevents OOM errors and maximizes GPU utilization.
    #[arg(long)]
    auto_batch_size: bool,

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

    /// Maximum validation batches to run (default: unlimited, use 50 for 4GB GPUs)
    /// Limits validation to N batches to reduce memory usage. Each batch uses ~10MB,
    /// so 50 batches = ~500MB vs 1760MB for full validation (176 batches).
    #[arg(long)]
    max_validation_batches: Option<usize>,
}

#[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")?;

    #[cfg(feature = "mimalloc-allocator")]
    info!("🚀 Using mimalloc allocator for improved performance");
    #[cfg(not(feature = "mimalloc-allocator"))]
    info!("ℹ️  Using system allocator (consider --features mimalloc-allocator for 10-25% speedup)");
    info!("🚀 Starting TFT Training with Parquet Data (Lazy Loading)");
    info!("");
    info!("Configuration:");
    info!("  • Parquet file: {}", opts.parquet_file);
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!(
        "  • Validation batch size: {}",
        opts.validation_batch_size.unwrap_or(opts.batch_size)
    );
    if let Some(max_val_batches) = opts.max_validation_batches {
        info!(
            "  • Max validation batches: {} (memory optimization)",
            max_val_batches
        );
    } else {
        info!("  • Max validation batches: unlimited");
    }
    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!("  • Dropout rate: {}", opts.dropout_rate);
    info!("  • LSTM layers: {}", opts.lstm_layers);
    info!("  • Quantiles: {}", opts.quantiles);
    info!("  • Feature count: 225 (Wave C 201 + Wave D 24)");
    info!("  • GPU enabled: {}", opts.use_gpu);
    info!("  • INT8 quantization: {}", opts.use_int8);
    info!("  • Quantization-Aware Training: {}", opts.use_qat);
    if opts.use_qat {
        info!(
            "  • QAT calibration batches: {}",
            opts.qat_calibration_batches
        );
    }
    info!(
        "  • Gradient checkpointing: {}",
        opts.use_gradient_checkpointing
    );
    if opts.use_gradient_checkpointing {
        if opts.use_qat {
            warn!("⚠️  WARNING: --use-gradient-checkpointing is IGNORED with --use-qat (not implemented)");
            warn!("    → For QAT memory reduction, use 2-phase workaround:");
            warn!("    → See ml/docs/QAT_GRADIENT_CHECKPOINTING_WORKAROUND.md");
        } else {
            info!("    → Expected: 30-40% memory reduction, ~20% slower training");
        }
    }
    info!("  • Output directory: {}", opts.output_dir);
    info!("");

    // Verify Parquet file exists
    let parquet_path = PathBuf::from(&opts.parquet_file);
    if !parquet_path.exists() {
        return Err(anyhow::anyhow!(
            "Parquet file not found: {}",
            opts.parquet_file
        ));
    }

    // 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);
    }

    // Parse quantiles
    let quantiles: Vec<f64> = opts
        .quantiles
        .split(',')
        .map(|s| s.trim().parse::<f64>())
        .collect::<Result<Vec<_>, _>>()
        .context("Failed to parse quantiles")?;

    if quantiles.is_empty() || quantiles.len() > 10 {
        return Err(anyhow::anyhow!(
            "Invalid number of quantiles: {} (must be 1-10)",
            quantiles.len()
        ));
    }

    info!(
        "📊 Quantiles for probabilistic forecasting: {:?}",
        quantiles
    );

    // Configure TFT trainer
    // Static features: 5 (symbol metadata)
    // Historical features: 210 (Wave C 201 features + Wave D 24 features - 5 static - 10 known)
    // Future features: 10 (calendar features, time-based)
    // Total input features: 225 (5 + 10 + 210)
    let trainer_config = TFTTrainerConfig {
        epochs: opts.epochs,
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        auto_batch_size: opts.auto_batch_size,
        validation_batch_size: opts.validation_batch_size.unwrap_or(opts.batch_size),
        hidden_dim: opts.hidden_dim,
        num_attention_heads: opts.num_attention_heads,
        dropout_rate: opts.dropout_rate,
        lstm_layers: opts.lstm_layers,
        quantiles,
        lookback_window: opts.lookback_window,
        forecast_horizon: opts.forecast_horizon,
        use_gpu: opts.use_gpu,
        use_int8_quantization: opts.use_int8,
        use_qat: opts.use_qat,
        qat_calibration_batches: opts.qat_calibration_batches,
        qat_min_batch_size: opts.qat_min_batch_size,
        qat_warmup_epochs: 10, // Default: 10 epochs LR warmup after calibration
        qat_cooldown_factor: 0.1, // Default: 10x LR reduction in final 10% of training
        use_gradient_checkpointing: opts.use_gradient_checkpointing,
        max_validation_batches: opts.max_validation_batches,
        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 with {} quantiles",
        trainer_config.quantiles.len()
    );

    if opts.use_qat {
        info!("🧠 Quantization-Aware Training (QAT) enabled");
        info!(
            "   Phase 1: Calibration ({} batches) - collecting activation statistics",
            opts.qat_calibration_batches
        );
        info!("   Phase 2: Training with fake quantization - simulating INT8 ops");
        info!("   Phase 3: Conversion to true INT8 model");
        info!("   Expected: 1-2% better accuracy than post-training quantization");
    } else if opts.use_int8 {
        info!("⚡ INT8 quantization enabled (PTQ mode) - expect 3-8x memory reduction");
        info!("   Memory usage: ~125MB (vs ~1GB FP32)");
    }

    // 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);
            }
            if let Some(attention_entropy) = progress.metrics.get("attention_entropy") {
                info!("  • Attention entropy: {:.4}", attention_entropy);
            }
        }
    });

    // Train the model using lazy-loading Parquet pipeline
    info!("");
    info!("🏋️  Starting training with lazy-loading Parquet pipeline...");
    info!("   (Loading 10,000 rows at a time to avoid OOM)");
    info!("");
    let start_time = std::time::Instant::now();

    let final_metrics = trainer
        .train_from_parquet(&opts.parquet_file)
        .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!("");
    info!("✅ Training completed successfully!");
    info!("");
    info!("📊 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!(
        "  • Reported training time: {:.1}s ({:.1} min)",
        final_metrics.training_time_seconds,
        final_metrics.training_time_seconds / 60.0
    );
    info!("");
    info!("💾 Model checkpoints saved to: {}", opts.output_dir);
    info!("");
    info!("🎉 TFT training with Parquet data complete!");

    Ok(())
}

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

    #[test]
    fn test_cli_parsing() {
        // Test default parameters
        let args = vec!["train_tft_parquet"];
        let opts = Opts::try_parse_from(args).expect("Failed to parse default args");

        assert_eq!(opts.parquet_file, "test_data/ES_FUT_180d.parquet");
        assert_eq!(opts.epochs, 20);
        assert_eq!(opts.batch_size, 32);
        assert_eq!(opts.lookback_window, 60);
        assert_eq!(opts.forecast_horizon, 10);
        assert_eq!(opts.learning_rate, 0.001);
        assert!(!opts.use_gpu); // Default is false (flag not set)
    }

    #[test]
    fn test_cli_custom_parameters() {
        // Test custom parameters with GPU enabled
        let args = vec![
            "train_tft_parquet",
            "--parquet-file",
            "test_data/NQ_FUT_180d.parquet",
            "--epochs",
            "50",
            "--batch-size",
            "16",
            "--lookback-window",
            "120",
            "--forecast-horizon",
            "20",
            "--learning-rate",
            "0.0005",
            "--use-gpu", // Flag to enable GPU
            "--verbose",
        ];

        let opts = Opts::try_parse_from(args).expect("Failed to parse custom args");

        assert_eq!(opts.parquet_file, "test_data/NQ_FUT_180d.parquet");
        assert_eq!(opts.epochs, 50);
        assert_eq!(opts.batch_size, 16);
        assert_eq!(opts.lookback_window, 120);
        assert_eq!(opts.forecast_horizon, 20);
        assert_eq!(opts.learning_rate, 0.0005);
        assert!(opts.use_gpu); // GPU should be enabled
        assert!(opts.verbose);
    }

    #[test]
    fn test_quantile_parsing() {
        let quantiles_str = "0.1,0.5,0.9";
        let quantiles: Vec<f64> = quantiles_str
            .split(',')
            .map(|s| s.trim().parse::<f64>())
            .collect::<Result<Vec<_>, _>>()
            .expect("Failed to parse quantiles");

        assert_eq!(quantiles.len(), 3);
        assert_eq!(quantiles[0], 0.1);
        assert_eq!(quantiles[1], 0.5);
        assert_eq!(quantiles[2], 0.9);
    }

    #[test]
    fn test_invalid_quantiles() {
        let quantiles_str = "invalid,0.5,0.9";
        let result: Result<Vec<f64>, _> = quantiles_str
            .split(',')
            .map(|s| s.trim().parse::<f64>())
            .collect();

        assert!(result.is_err(), "Should fail on invalid quantile");
    }
}