//! TFT Hyperparameter Optimization Adapter
//!
//! This module provides a production-ready adapter for optimizing TFT
//! hyperparameters using the generic optimization framework. It implements:
//!
//! - Parameter space with log-scale handling for learning rates
//! - Training wrapper that integrates with existing TFT pipeline
//! - Metrics extraction for quantile loss optimization
//!
//! ## Usage Example
//!
//! ```rust,no_run
//! use ml::hyperopt::ArgminOptimizer;
//! use ml::hyperopt::adapters::tft::{TFTTrainer, TFTParams};
//!
//! # async fn example() -> anyhow::Result<()> {
//! // Create trainer
//! let trainer = TFTTrainer::new(
//!     "test_data/ES_FUT_180d.parquet",
//!     50, // epochs per trial
//! )?;
//!
//! // Run optimization
//! let optimizer = ArgminOptimizer::with_trials(30, 5);
//! let result = optimizer.optimize(trainer)?;
//!
//! println!("Best learning rate: {}", result.best_params.learning_rate);
//! println!("Best batch size: {}", result.best_params.batch_size);
//! println!("Best validation loss: {:.6}", result.best_objective);
//! # Ok(())
//! # }
//! ```

use anyhow::Result;
use ml_core::device::MlDevice;
use serde::{Deserialize, Serialize};
use std::fs::OpenOptions;
use std::io::Write as IoWrite;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::features::extraction::OHLCVBar;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::trainers::tft::{TFTTrainer as RealTFTTrainer, TFTTrainerConfig};
use crate::MLError;

/// TFT model overhead in MB (weights + optimizer state + activations)
/// TFT with hidden_size=512, 4 heads, 3 layers: ~2M params × 4B × 6 ≈ 48MB base
const MODEL_OVERHEAD_MB: f64 = 500.0;
/// TFT per-sample memory in MB (sequence processing with attention + GRN blocks)
const MB_PER_SAMPLE: f64 = 0.02;

/// TFT hyperparameter space
///
/// Defines the hyperparameters to optimize for TFT training:
/// - Learning rate (log-scale: 1e-5 to 1e-3)
/// - Batch size (linear scale: 16 to 128)
/// - Hidden size (discrete: 128, 256, 512)
/// - Number of attention heads (discrete: 4, 8, 16)
/// - Dropout rate (linear scale: 0.0 to 0.3)
///
/// ## Parameter Scaling
///
/// - **Log-scale**: Learning rate (spans multiple orders)
/// - **Linear scale**: Batch size, dropout (span single order)
/// - **Discrete**: Hidden size, num_heads (power-of-2 values)
///
/// This scaling ensures efficient exploration by argmin's optimization.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TFTParams {
    /// Learning rate for AdamW optimizer (log-scale)
    pub learning_rate: f64,
    /// Batch size for training (linear scale, integer)
    pub batch_size: usize,
    /// Hidden dimension for attention layers (discrete: 128, 256, 512, 1024, 2048)
    pub hidden_size: usize,
    /// Number of attention heads (discrete: 4, 8, 16)
    pub num_heads: usize,
    /// Dropout rate for regularization (linear scale)
    pub dropout: f64,
    /// Signal high threshold in basis points (GPU backtest action mapping)
    pub signal_high_bps: f64,
    /// Signal low threshold in basis points (GPU backtest action mapping)
    pub signal_low_bps: f64,
}

impl Default for TFTParams {
    fn default() -> Self {
        Self {
            learning_rate: 1e-4,
            batch_size: 64,
            hidden_size: 256,
            num_heads: 8,
            dropout: 0.1,
            signal_high_bps: 10.0,
            signal_low_bps: 5.0,
        }
    }
}

impl ParameterSpace for TFTParams {
    fn continuous_bounds() -> Vec<(f64, f64)> {
        vec![
            (1e-5_f64.ln(), 1e-3_f64.ln()), // learning_rate (log scale)
            (16.0, 512.0),                  // batch_size (linear)
            (0.0, 4.0),                     // hidden_size_index (0=128, 1=256, 2=512, 3=1024, 4=2048)
            (0.0, 2.0),                     // num_heads_index (0=4, 1=8, 2=16)
            (0.0, 0.3),                     // dropout (linear)
            (1.0, 30.0),                    // signal_high_bps
            (0.5, 15.0),                    // signal_low_bps
        ]
    }

    fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
        if x.len() != 7 {
            return Err(MLError::ConfigError(format!("Expected 7 parameters, got {}", x.len())));
        }

        // Map continuous indices to discrete values
        let hidden_sizes = [128, 256, 512, 1024, 2048];
        let num_heads_options = [4, 8, 16];

        let hidden_size_idx = x[2].round().clamp(0.0, 4.0) as usize;
        let num_heads_idx = x[3].round().clamp(0.0, 2.0) as usize;

        Ok(Self {
            learning_rate: x[0].exp(),
            batch_size: x[1].round().max(16.0).min(512.0) as usize,
            hidden_size: hidden_sizes[hidden_size_idx],
            num_heads: num_heads_options[num_heads_idx],
            dropout: x[4].clamp(0.0, 0.3),
            signal_high_bps: x[5].clamp(1.0, 30.0),
            signal_low_bps: x[6].clamp(0.5, 15.0),
        })
    }

    fn to_continuous(&self) -> Vec<f64> {
        // Map discrete values back to continuous indices
        let hidden_size_idx = match self.hidden_size {
            128 => 0.0,
            256 => 1.0,
            512 => 2.0,
            1024 => 3.0,
            2048 => 4.0,
            _ => 2.0, // Default to 512
        };

        let num_heads_idx = match self.num_heads {
            4 => 0.0,
            8 => 1.0,
            16 => 2.0,
            _ => 1.0, // Default to 8
        };

        vec![
            self.learning_rate.ln(),
            self.batch_size as f64,
            hidden_size_idx,
            num_heads_idx,
            self.dropout,
            self.signal_high_bps,
            self.signal_low_bps,
        ]
    }

    fn param_names() -> Vec<&'static str> {
        vec![
            "learning_rate",
            "batch_size",
            "hidden_size",
            "num_heads",
            "dropout",
            "signal_high_bps",
            "signal_low_bps",
        ]
    }

    fn continuous_bounds_for(budget: &HardwareBudget) -> Vec<(f64, f64)> {
        let mut bounds = Self::continuous_bounds();
        if let Some(max_batch) = budget.max_batch_size(MODEL_OVERHEAD_MB, MB_PER_SAMPLE, 16.0, 4096.0) {
            if let Some(batch_bound) = bounds.get_mut(1) {
                batch_bound.1 = max_batch;
            }
        }

        // Cap hidden_size by VRAM (index 2 maps to {128, 256, 512, 1024, 2048})
        if budget.gpu_memory_mb < 8000 {
            // Small GPU (<8GB): force hidden_size=128
            bounds[2] = (0.0, 0.0);
        } else if budget.gpu_memory_mb < 16000 {
            // Medium GPU (8-16GB): cap at hidden_size=256
            bounds[2] = (0.0, 1.0);
        } else if budget.gpu_memory_mb < 25000 {
            // Large GPU (16-25GB): cap at hidden_size=512
            bounds[2] = (0.0, 2.0);
        } else if budget.gpu_memory_mb < 40000 {
            // XL GPU (25-40GB): cap at hidden_size=1024
            bounds[2] = (0.0, 3.0);
        } else {
            // ≥40GB (L40S 46GB, H100 80GB): allow full range [0, 4] → {128, 256, 512, 1024, 2048}
        }

        bounds
    }
}

/// TFT training metrics
///
/// Contains all relevant metrics from a TFT training run.
/// The primary optimization target is validation quantile loss.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TFTMetrics {
    /// Final validation loss (optimization target)
    pub val_loss: f64,
    /// Final training loss
    pub train_loss: f64,
    /// Validation RMSE
    pub val_rmse: f64,
    /// Number of epochs completed
    pub epochs_completed: usize,
    /// GPU walk-forward backtest Sharpe ratio
    pub backtest_sharpe: Option<f32>,
    /// GPU walk-forward backtest total trades
    pub backtest_trades: Option<u32>,
}

/// TFT trainer for hyperparameter optimization
///
/// This struct wraps the TFT training pipeline and implements
/// `HyperparameterOptimizable` for use with optimization backends.
///
/// ## Configuration
///
/// - **Data dir**: Directory with .dbn.zst files (Databento OHLCV data)
/// - **Epochs**: Number of training epochs per trial
/// - **Device**: CUDA GPU (falls back to CPU if unavailable)
/// - **Features**: 42 market features from extraction pipeline
///
/// ## Optimized Hyperparameters
///
/// The following are optimized by `TFTParams`:
/// - Learning rate
/// - Batch size
/// - Hidden size
/// - Number of attention heads
/// - Dropout
#[derive(Debug)]
pub struct TFTTrainer {
    data_dir: PathBuf,
    epochs: usize,
    device: MlDevice,
    training_paths: TrainingPaths,
    early_stopping_patience: usize,
    trial_counter: usize,
    preloaded_bars: Option<Arc<[OHLCVBar]>>,
}

impl TFTTrainer {
    /// Create a new TFT trainer
    ///
    /// # Arguments
    ///
    /// * `data_dir` - Directory containing .dbn.zst files with OHLCV data
    /// * `epochs` - Number of training epochs per trial
    pub fn new<P: Into<PathBuf>>(data_dir: P, epochs: usize) -> Result<Self> {
        let data_dir = data_dir.into();

        if !data_dir.exists() {
            return Err(MLError::ConfigError(format!("Data directory not found: {}", data_dir.display()))
            .into());
        }

        let device = MlDevice::cuda(0)
            .map_err(|e| MLError::ConfigError(format!("CUDA GPU required for TFT hyperopt: {}", e)))?;

        info!(
            "TFT Trainer initialized: Device={:?}, Data={}, Epochs per trial={}",
            device,
            data_dir.display(),
            epochs
        );

        let training_paths = TrainingPaths::new("/tmp/ml_training", "tft", "default");

        Ok(Self {
            data_dir,
            epochs,
            device,
            training_paths,
            early_stopping_patience: 10,
            trial_counter: 0,
            preloaded_bars: None,
        })
    }

    /// Configure early stopping parameters (overrides defaults)
    ///
    /// NOTE: Currently this value is stored but not used by the TFT trainer,
    /// which has hardcoded patience of 20 epochs. This method exists for
    /// CLI consistency and future extensibility.
    ///
    /// # Arguments
    ///
    /// * `patience` - Epochs without improvement before stopping
    ///
    /// # Returns
    ///
    /// Self for method chaining
    pub fn with_early_stopping(mut self, patience: usize) -> Self {
        self.early_stopping_patience = patience;
        self
    }

    /// Set training paths configuration
    ///
    /// # Arguments
    ///
    /// * `paths` - Training paths configuration
    ///
    /// # Returns
    ///
    /// Self for method chaining
    pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self {
        info!(
            "TFT training paths configured: Run={:?}, Checkpoints={:?}",
            paths.run_dir(),
            paths.checkpoints_dir()
        );
        self.training_paths = paths;
        self
    }

    /// Preload DBN data once for reuse across all hyperopt trials.
    pub fn preload_data(&mut self) -> Result<(), MLError> {
        if self.preloaded_bars.is_some() {
            return Ok(());
        }
        let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir)
            .map_err(|e| MLError::ModelError(format!("Failed to preload DBN data: {}", e)))?;
        info!("Preloaded {} OHLCV bars for reuse across trials", bars.len());
        self.preloaded_bars = Some(Arc::from(bars));
        Ok(())
    }

    /// Check if data has been preloaded.
    pub fn has_preloaded_data(&self) -> bool {
        self.preloaded_bars.is_some()
    }
}

/// Write a log entry to the training log file
fn write_training_log_tft(logs_dir: &std::path::Path, message: &str) -> Result<(), std::io::Error> {
    let log_file = logs_dir.join("training.log");
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(log_file)?;

    let timestamp = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S");
    writeln!(file, "[{}] {}", timestamp, message)?;
    Ok(())
}

/// Write trial results to JSON file
fn write_trial_result_tft(
    hyperopt_dir: &std::path::Path,
    trial_result: &crate::hyperopt::traits::TrialResult<TFTParams>,
) -> Result<(), std::io::Error> {
    let trials_file = hyperopt_dir.join("trials.json");

    // Read existing trials (if any)
    let mut all_trials = if trials_file.exists() {
        let content = std::fs::read_to_string(&trials_file)?;
        serde_json::from_str::<Vec<serde_json::Value>>(&content).unwrap_or_default()
    } else {
        Vec::new()
    };

    // Append new trial
    let trial_json = serde_json::to_value(trial_result)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
    all_trials.push(trial_json);

    // Write back to file (pretty printed)
    let content = serde_json::to_string_pretty(&all_trials)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
    std::fs::write(&trials_file, content)?;

    Ok(())
}

impl HyperparameterOptimizable for TFTTrainer {
    type Params = TFTParams;
    type Metrics = TFTMetrics;

    #[allow(clippy::unwrap_in_result)]
    fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
        let trial_start = std::time::Instant::now();

        let current_trial = self.trial_counter;
        self.trial_counter += 1;

        info!(
            "Training TFT: lr={:.6}, batch={}, hidden={}, heads={}, dropout={:.3}",
            params.learning_rate,
            params.batch_size,
            params.hidden_size,
            params.num_heads,
            params.dropout
        );

        std::fs::create_dir_all(self.training_paths.logs_dir()).ok();
        write_training_log_tft(
            &self.training_paths.logs_dir(),
            &format!("=== Starting TFT Trial ===\nParams: {:#?}", params),
        )
        .ok();

        if params.hidden_size % params.num_heads != 0 {
            warn!(
                "Hidden size {} not divisible by num_heads {}, adjusting",
                params.hidden_size, params.num_heads
            );
            return Ok(TFTMetrics {
                val_loss: 1000.0,
                train_loss: 1000.0,
                val_rmse: 1000.0,
                epochs_completed: 0,
                backtest_sharpe: None,
                backtest_trades: None,
            });
        }

        self.training_paths.create_all().map_err(|e| {
            MLError::ModelError(format!("Failed to create training directories: {}", e))
        })?;

        let trainer_config = TFTTrainerConfig {
            epochs: self.epochs,
            learning_rate: params.learning_rate,
            batch_size: params.batch_size,
            auto_batch_size: false,
            hidden_dim: params.hidden_size,
            num_attention_heads: params.num_heads,
            dropout_rate: params.dropout,
            lstm_layers: 2,
            quantiles: vec![0.1, 0.5, 0.9],
            lookback_window: 60,
            forecast_horizon: 10,
            use_gpu: true,
            use_gradient_checkpointing: false,
            validation_frequency: 1,
            validation_batch_size: params.batch_size,
            max_validation_batches: None,
            checkpoint_dir: self
                .training_paths
                .checkpoints_dir()
                .to_string_lossy()
                .to_string(),
        };

        let checkpoint_storage = std::sync::Arc::new(crate::checkpoint::MemoryStorage::new());
        let mut trainer = RealTFTTrainer::new(trainer_config, checkpoint_storage)
            .map_err(|e| MLError::ModelError(format!("Failed to create TFT trainer: {}", e)))?;

        // Load bars from DBN files (use preloaded cache if available)
        let bars = if let Some(ref preloaded) = self.preloaded_bars {
            preloaded.to_vec() // cpu-side Arc<[OHLCVBar]> clone
        } else {
            super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir)
                .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?
        };

        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| MLError::ModelError(format!("Failed to create async runtime: {}", e)))?;

        let training_metrics = runtime
            .block_on(trainer.train_from_bars(&bars))
            .map_err(|e| MLError::TrainingError(format!("TFT training failed: {}", e)))?;

        // Map from crate::trainers::tft::TrainingMetrics to TFTMetrics
        let metrics = TFTMetrics {
            val_loss: training_metrics.val_loss,
            train_loss: training_metrics.train_loss,
            val_rmse: training_metrics.rmse,
            epochs_completed: self.epochs,
            backtest_sharpe: None,
            backtest_trades: None,
        };

        info!(
            "Training completed: train_loss={:.6}, val_loss={:.6}, rmse={:.4}",
            metrics.train_loss, metrics.val_loss, metrics.val_rmse
        );

        // END: Add trial completion logging
        let duration_secs = trial_start.elapsed().as_secs_f64();
        write_training_log_tft(
            &self.training_paths.logs_dir(),
            &format!(
                "Training completed in {:.2}s: val_loss={:.6}, train_loss={:.6}, rmse={:.4}",
                duration_secs, metrics.val_loss, metrics.train_loss, metrics.val_rmse
            ),
        )
        .ok();

        // CRITICAL FIX: Explicit memory cleanup to prevent OOM between trials
        // Drop trainer and sync CUDA to free GPU/RAM
        info!("Cleaning up resources...");
        drop(trainer);

        // Sync CUDA to ensure GPU memory is freed
        if self.device.is_cuda() {
            // Force CUDA synchronization to release GPU memory
            std::thread::sleep(std::time::Duration::from_millis(100));
        }
        info!("Resource cleanup complete");

        // Write trial result to JSON (ensure directory exists first)
        let trial_result = crate::hyperopt::traits::TrialResult {
            trial_num: current_trial,
            params,
            objective: Self::extract_objective(&metrics),
            duration_secs,
            metrics: None,
        };

        std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
        write_trial_result_tft(&self.training_paths.hyperopt_dir(), &trial_result).ok();

        Ok(metrics)
    }

    fn extract_objective(metrics: &Self::Metrics) -> f64 {
        // Prefer GPU backtest fitness when available
        if let (Some(sharpe), Some(trades)) = (metrics.backtest_sharpe, metrics.backtest_trades) {
            return crate::cuda_pipeline::signal_adapter::backtest_fitness(sharpe, trades, 30, 0.0);
        }
        // Fallback: minimize validation loss
        metrics.val_loss
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tft::{TFTConfig, TemporalFusionTransformer};

    #[test]
    fn test_tft_params_roundtrip() {
        let params = TFTParams {
            learning_rate: 0.0001,
            batch_size: 64,
            hidden_size: 256,
            num_heads: 8,
            dropout: 0.1,
            signal_high_bps: 10.0,
            signal_low_bps: 5.0,
        };

        let continuous = params.to_continuous();
        let recovered = TFTParams::from_continuous(&continuous).unwrap();

        assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
        assert_eq!(recovered.batch_size, params.batch_size);
        assert_eq!(recovered.hidden_size, params.hidden_size);
        assert_eq!(recovered.num_heads, params.num_heads);
        assert!((recovered.dropout - params.dropout).abs() < 1e-10);
    }

    #[test]
    fn test_tft_params_bounds() {
        let bounds = TFTParams::continuous_bounds();
        assert_eq!(bounds.len(), 7);

        // Check log-scale bounds are reasonable
        assert!(bounds[0].0 < bounds[0].1); // learning_rate

        // Check discrete bounds
        assert_eq!(bounds[2], (0.0, 4.0)); // hidden_size_index
        assert_eq!(bounds[3], (0.0, 2.0)); // num_heads_index

        // Check linear bounds
        assert_eq!(bounds[1], (16.0, 512.0)); // batch_size
        assert_eq!(bounds[4], (0.0, 0.3)); // dropout

        // Check signal threshold bounds
        assert_eq!(bounds[5], (1.0, 30.0)); // signal_high_bps
        assert_eq!(bounds[6], (0.5, 15.0)); // signal_low_bps
    }

    #[test]
    fn test_tft_discrete_params() {
        // Test all discrete hidden_size values
        for (idx, expected_size) in [(0.0, 128), (1.0, 256), (2.0, 512)] {
            let x = vec![(-4.6_f64).ln(), 64.0, idx, 1.0, 0.1, 10.0, 5.0];
            let params = TFTParams::from_continuous(&x).unwrap();
            assert_eq!(params.hidden_size, expected_size);
        }

        // Test all discrete num_heads values
        for (idx, expected_heads) in [(0.0, 4), (1.0, 8), (2.0, 16)] {
            let x = vec![(-4.6_f64).ln(), 64.0, 1.0, idx, 0.1, 10.0, 5.0];
            let params = TFTParams::from_continuous(&x).unwrap();
            assert_eq!(params.num_heads, expected_heads);
        }
    }

    #[test]
    fn test_param_names() {
        let names = TFTParams::param_names();
        assert_eq!(names.len(), 7);
        assert_eq!(names[0], "learning_rate");
        assert_eq!(names[1], "batch_size");
        assert_eq!(names[2], "hidden_size");
        assert_eq!(names[3], "num_heads");
        assert_eq!(names[4], "dropout");
        assert_eq!(names[5], "signal_high_bps");
        assert_eq!(names[6], "signal_low_bps");
    }

    #[test]
    fn test_tft_config_api_match() {
        // Verify TFTConfig field names match actual TFT implementation
        let params = TFTParams::default();

        let config = TFTConfig {
            input_dim: 51,                  // TODO: should be 42 market features (was 51, needs CUDA alignment)
            hidden_dim: params.hidden_size, // ✅ Was: hidden_size
            num_heads: params.num_heads,    // ✅ Correct
            num_layers: 2,                  // ✅ Correct
            prediction_horizon: 10,
            sequence_length: 60,
            num_quantiles: 3,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 36,       // WAVE 10: 51 - 5 - 10 = 36
            learning_rate: params.learning_rate,
            batch_size: params.batch_size,
            dropout_rate: params.dropout, // ✅ Was: dropout
            l2_regularization: 1e-4,
            use_flash_attention: true,
            memory_efficient: true,
            max_inference_latency_us: 50,
            target_throughput_pps: 100_000,
        };

        // Verify config matches WAVE 10 feature split
        assert_eq!(
            config.num_static_features + config.num_known_features + config.num_unknown_features,
            51,
            "TFT config must use 51 total features (WAVE 10: Proxy OFI removed)"
        );

        // Verify num_heads divides hidden_dim evenly
        assert_eq!(
            config.hidden_dim % config.num_heads,
            0,
            "hidden_dim must be divisible by num_heads"
        );
    }

    #[test]
    fn test_tft_trainer_creation() {
        // Test that TFTTrainer can be created with correct device handling
        // This test doesn't require actual Parquet file, just validates construction
        let result = TFTTrainer::new("nonexistent.parquet", 10);

        // Should fail with config error (file not found)
        assert!(result.is_err());

        let err = result.unwrap_err();
        let err_msg = format!("{:?}", err);
        assert!(
            err_msg.contains("not found"),
            "Should report file not found"
        );
    }

    #[test]
    fn test_tft_model_creation_with_params() {
        // Test that TFT model can be created with different hyperparameter configs
        let test_cases = vec![
            (128, 4),  // Small model
            (256, 8),  // Default model
            (512, 16), // Large model
        ];

        for (hidden_size, num_heads) in test_cases {
            let config = TFTConfig {
                input_dim: 225,
                hidden_dim: hidden_size,
                num_heads,
                num_layers: 2,
                prediction_horizon: 10,
                sequence_length: 60,
                num_quantiles: 3,
                num_static_features: 5,
                num_known_features: 10,
                num_unknown_features: 210,
                learning_rate: 1e-4,
                batch_size: 32,
                dropout_rate: 0.1,
                l2_regularization: 1e-4,
                use_flash_attention: true,
                memory_efficient: true,
                max_inference_latency_us: 50,
                target_throughput_pps: 100_000,
            };

            let device = MlDevice::cuda(0).expect("CUDA required");
            let stream = device.cuda_stream().expect("cuda stream");
            let result = TemporalFusionTransformer::new_with_stream(config, Some(std::sync::Arc::clone(stream)));

            assert!(
                result.is_ok(),
                "Failed to create TFT with hidden_dim={}, num_heads={}",
                hidden_size,
                num_heads
            );
        }
    }

    #[test]
    fn test_parameter_space_coverage() {
        // Verify parameter space covers production requirements
        let bounds = TFTParams::continuous_bounds();

        // Learning rate: 1e-5 to 1e-3 (log scale)
        let lr_min = bounds[0].0.exp();
        let lr_max = bounds[0].1.exp();
        assert!((lr_min - 1e-5).abs() < 1e-10, "LR min should be 1e-5");
        assert!((lr_max - 1e-3).abs() < 1e-10, "LR max should be 1e-3");

        // Batch size: 16 to 512 (linear scale, VRAM-capped at runtime)
        assert_eq!(bounds[1], (16.0, 512.0));

        // Hidden size: [128, 256, 512, 1024, 2048] via index [0, 4]
        assert_eq!(bounds[2], (0.0, 4.0));

        // Num heads: [4, 8, 16] via index [0, 2]
        assert_eq!(bounds[3], (0.0, 2.0));

        // Dropout: 0.0 to 0.3 (linear scale)
        assert_eq!(bounds[4], (0.0, 0.3));
    }
}
