foxhunt/ml/examples/train_tft_dbn.rs

//! 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, Timelike, TimeZone, Utc};
use dbn::decode::{DecodeRecordRef, DbnDecoder};
use dbn::OhlcvMsg;
use ndarray::{Array1, Array2};
use std::path::PathBuf;
use structopt::StructOpt;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::checkpoint::FileSystemStorage;
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};
use ml::tft::training::TFTDataLoader;
use ml::data_loaders::BarSamplingMethod;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    // 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");
    info!("Configuration:");
    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!("  • 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
    info!("\n🔄 Converting to TFT data format...");
    let tft_data = convert_to_tft_data(
        &bars,
        opts.lookback_window,
        opts.forecast_horizon,
    ).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
    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
        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 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 format
///
/// TFT expects:
/// - Static features: Symbol metadata, exchange, trading hours
/// - Historical features: Past OHLCV, volume, spreads, returns
/// - Future features: Known future events (calendar features)
/// - Targets: Multi-horizon price forecast
fn convert_to_tft_data(
    bars: &[OhlcvBar],
    lookback_window: usize,
    forecast_horizon: usize,
) -> 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()
        ));
    }

    let mut tft_samples = Vec::new();

    // Calculate statistics for static features
    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();

    // Create sliding windows
    for i in 0..bars.len() - lookback_window - forecast_horizon + 1 {
        // Static features: Symbol metadata (10 features)
        // [mean_price, price_std, mean_volume, volume_std, hour_of_day, day_of_week, is_morning, is_afternoon, volatility, liquidity]
        let first_bar = &bars[i];
        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[i..i + 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;  // Simple liquidity measure

        let static_features = Array1::from_vec(vec![
            mean_price / 5000.0,      // Normalize around ES futures price
            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,       // Scale volatility
            liquidity / 100.0,
        ]);

        // Historical features: Past OHLCV + derived features (50 features per timestep)
        // [open, high, low, close, volume, returns, high-low spread, close-open,
        //  SMA_5, SMA_20, EMA_12, RSI_14, MACD, volatility_5, volatility_20,
        //  volume_sma, price_change_pct, intraday_range, typical_price, ...]
        let mut hist_features = Vec::new();

        for t in 0..lookback_window {
            let bar = &bars[i + t];
            let prev_bar = if t > 0 { &bars[i + t - 1] } else { bar };

            // Basic OHLCV (normalized)
            let open = bar.open / mean_price;
            let high = bar.high / mean_price;
            let low = bar.low / mean_price;
            let close = bar.close / mean_price;
            let volume = bar.volume / mean_volume;

            // Derived features
            let returns = ((bar.close / prev_bar.close).ln() * 100.0).min(10.0).max(-10.0);
            let spread = (bar.high - bar.low) / bar.close;
            let body = (bar.close - bar.open) / bar.close;

            // Calculate SMA_5 (using available history)
            let sma_5 = if t >= 4 {
                let sum: f64 = (0..5).map(|j| bars[i + t - j].close).sum();
                sum / 5.0 / mean_price
            } else {
                close
            };

            // Calculate SMA_20 (using available history)
            let sma_20 = if t >= 19 {
                let sum: f64 = (0..20).map(|j| bars[i + t - j].close).sum();
                sum / 20.0 / mean_price
            } else {
                close
            };

            // Calculate EMA_12 (simplified)
            let ema_12 = close;  // Use close as proxy for EMA (proper EMA would need state)

            // Calculate RSI_14 (simplified)
            let rsi_14 = if t >= 14 {
                let recent_returns: Vec<f64> = (1..=14).map(|j| {
                    (bars[i + t - j + 1].close / bars[i + t - j].close).ln()
                }).collect();

                let gains: f64 = recent_returns.iter().filter(|r| **r > 0.0).sum();
                let losses: f64 = recent_returns.iter().filter(|r| **r < 0.0).map(|r| -r).sum();

                if losses < 1e-10 {
                    100.0
                } else {
                    let rs = gains / losses;
                    100.0 - (100.0 / (1.0 + rs))
                }
            } else {
                50.0  // Neutral RSI
            } / 100.0;  // Normalize to [0, 1]

            // MACD (simplified: close - SMA_20)
            let macd = (close - sma_20) / sma_20;

            // Volatility (5-period and 20-period)
            let vol_5 = if t >= 5 {
                let returns: Vec<f64> = (1..=5).map(|j| {
                    (bars[i + t - j + 1].close / bars[i + t - j].close).ln()
                }).collect();
                let mean = returns.iter().sum::<f64>() / returns.len() as f64;
                (returns.iter().map(|r| (r - mean).powi(2)).sum::<f64>() / returns.len() as f64).sqrt() * 100.0
            } else {
                0.01
            };

            let vol_20 = volatility * 100.0;  // Use pre-calculated volatility

            // Volume features
            let volume_sma = if t >= 4 {
                let sum: f64 = (0..5).map(|j| bars[i + t - j].volume).sum();
                sum / 5.0 / mean_volume
            } else {
                volume
            };

            let volume_change_pct = if t > 0 {
                ((bar.volume - prev_bar.volume) / prev_bar.volume).min(2.0).max(-2.0)
            } else {
                0.0
            };

            // Price metrics
            let intraday_range = spread;
            let typical_price = ((bar.high + bar.low + bar.close) / 3.0) / mean_price;
            let weighted_price = ((bar.high + bar.low + bar.close * 2.0) / 4.0) / mean_price;

            // Time features
            let hour_sin = (hour * 2.0 * std::f64::consts::PI / 24.0).sin();
            let hour_cos = (hour * 2.0 * std::f64::consts::PI / 24.0).cos();
            let day_sin = (day_of_week * 2.0 * std::f64::consts::PI / 7.0).sin();
            let day_cos = (day_of_week * 2.0 * std::f64::consts::PI / 7.0).cos();

            // Momentum features
            let momentum_5 = if t >= 5 {
                (bar.close / bars[i + t - 5].close - 1.0).min(0.1).max(-0.1)
            } else {
                0.0
            };

            let momentum_20 = if t >= 20 {
                (bar.close / bars[i + t - 20].close - 1.0).min(0.2).max(-0.2)
            } else {
                0.0
            };

            // Order flow proxy (volume * sign of price change)
            let order_flow = volume * returns.signum();

            // Combine all features (50 total)
            let mut features = vec![
                open, high, low, close, volume,           // 5: Basic OHLCV
                returns, spread, body,                     // 3: Price dynamics
                sma_5, sma_20, ema_12,                    // 3: Moving averages
                rsi_14, macd,                             // 2: Momentum indicators
                vol_5, vol_20,                            // 2: Volatility
                volume_sma, volume_change_pct,            // 2: Volume indicators
                intraday_range, typical_price, weighted_price, // 3: Price metrics
                hour_sin, hour_cos, day_sin, day_cos,     // 4: Time features
                momentum_5, momentum_20,                  // 2: Momentum
                order_flow,                               // 1: Order flow
            ];

            // Pad remaining features to reach 50 (add technical ratios and cross-features)
            while features.len() < 50 {
                let idx = features.len();
                match idx {
                    28 => features.push(close / sma_5 - 1.0),           // Price vs SMA_5
                    29 => features.push(close / sma_20 - 1.0),          // Price vs SMA_20
                    30 => features.push(volume / volume_sma - 1.0),     // Volume ratio
                    31 => features.push(spread * volume),                // Spread-volume
                    32 => features.push(returns * volume),               // Return-volume
                    33 => features.push(high / sma_20 - 1.0),           // High vs SMA
                    34 => features.push(low / sma_20 - 1.0),            // Low vs SMA
                    35 => features.push(vol_5 / (vol_20 + 1e-6)),       // Vol ratio
                    36 => features.push(rsi_14 - 0.5),                  // RSI deviation
                    37 => features.push((sma_5 / sma_20 - 1.0).min(0.1).max(-0.1)), // SMA cross
                    38 => features.push(body * volume),                  // Body-volume
                    39 => features.push(returns.abs()),                  // Absolute returns
                    40 => features.push((high - close) / (high - low + 1e-6)), // Upper shadow
                    41 => features.push((close - low) / (high - low + 1e-6)),   // Lower shadow
                    42 => features.push((typical_price - close).abs()), // Price deviation
                    43 => features.push(momentum_5 * momentum_20),       // Momentum product
                    44 => features.push(is_morning * volume),            // Morning volume
                    45 => features.push(is_afternoon * volume),          // Afternoon volume
                    46 => features.push(volatility * returns.abs()),     // Vol-return
                    47 => features.push((close - typical_price).signum()), // Price bias
                    48 => features.push(order_flow.abs()),               // Order flow magnitude
                    49 => features.push((volume - volume_sma).abs()),    // Volume surprise
                    _ => features.push(0.0),
                }
            }

            hist_features.extend(features);
        }

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

        // Future features: Known future events (10 features per timestep)
        // [hour, day_of_week, is_weekend, is_morning, is_afternoon, week_of_month,
        //  month, quarter, is_month_start, is_month_end]
        let mut fut_features = Vec::new();

        for t in 0..forecast_horizon {
            let future_bar = &bars[i + 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[i + 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 tft_data = convert_to_tft_data(&bars, 60, 10).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, 50], "Historical features should be [60, 50]");
        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());
    }
}