foxhunt/crates/ml/src/hyperopt/adapters/diffusion.rs

//! Hyperopt adapter for the Diffusion model.
//!
//! Defines `DiffusionParams` (ParameterSpace), `DiffusionMetrics`,
//! and `DiffusionTrainer` (HyperparameterOptimizable) for hyperparameter
//! optimization of the Diffusion model via the unified framework.

use candle_core::Device;
use std::path::PathBuf;
use std::sync::Arc;
use tracing::{info, warn};

use crate::diffusion::config::{DiffusionConfig, NoiseSchedule};
use crate::diffusion::trainable::DiffusionTrainableAdapter;
use crate::features::extract_ml_features;
use crate::features::extraction::OHLCVBar;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;

/// Diffusion model overhead in MB (weights + optimizer state, heavy due to denoiser)
const MODEL_OVERHEAD_MB: f64 = 100.0;
/// Diffusion per-sample memory in MB (noise + denoise activations)
const MB_PER_SAMPLE: f64 = 0.03;

/// Hyperparameters for Diffusion model hyperopt tuning.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DiffusionParams {
    /// Learning rate (log scale).
    pub learning_rate: f64,
    /// Number of diffusion timesteps.
    pub num_timesteps: usize,
    /// Number of DDIM sampling steps.
    pub sampling_steps: usize,
    /// Hidden dimension of the denoiser.
    pub hidden_dim: usize,
    /// Number of denoiser layers.
    pub num_layers: usize,
    /// Time embedding dimension.
    pub time_embed_dim: usize,
    /// Batch size for training.
    pub batch_size: usize,
    /// Weight decay (log scale).
    pub weight_decay: f64,
    /// Gradient clipping max norm (log scale).
    pub grad_clip: f64,
}

impl Default for DiffusionParams {
    fn default() -> Self {
        Self {
            learning_rate: 1e-4,
            num_timesteps: 1000,
            sampling_steps: 10,
            hidden_dim: 128,
            num_layers: 3,
            time_embed_dim: 32,
            batch_size: 32,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }
}

impl ParameterSpace for DiffusionParams {
    fn continuous_bounds() -> Vec<(f64, f64)> {
        vec![
            (1e-5_f64.ln(), 1e-3_f64.ln()),   // learning_rate (log)
            (100.0, 2000.0),                   // num_timesteps
            (5.0, 50.0),                       // sampling_steps
            (32.0, 2048.0),                    // hidden_dim
            (1.0, 6.0),                        // num_layers
            (8.0, 64.0),                       // time_embed_dim
            (4.0, 256.0),                      // batch_size
            (1e-6_f64.ln(), 1e-2_f64.ln()),    // weight_decay (log)
            (0.5_f64.ln(), 5.0_f64.ln()),      // grad_clip (log)
        ]
    }

    fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
        if x.len() != 9 {
            return Err(MLError::ConfigError(format!("Expected 9 params, got {}", x.len())));
        }
        Ok(Self {
            learning_rate: x.first().copied().unwrap_or(-9.21).exp(),
            num_timesteps: x.get(1).copied().unwrap_or(1000.0).round().max(100.0) as usize,
            sampling_steps: x.get(2).copied().unwrap_or(10.0).round().max(5.0) as usize,
            hidden_dim: x.get(3).copied().unwrap_or(128.0).round().max(32.0) as usize,
            num_layers: x.get(4).copied().unwrap_or(3.0).round().max(1.0) as usize,
            time_embed_dim: x.get(5).copied().unwrap_or(32.0).round().max(8.0) as usize,
            batch_size: x.get(6).copied().unwrap_or(32.0).round().max(4.0) as usize,
            weight_decay: x.get(7).copied().unwrap_or(-9.21).exp(),
            grad_clip: x.get(8).copied().unwrap_or(0.0).exp(),
        })
    }

    fn to_continuous(&self) -> Vec<f64> {
        vec![
            self.learning_rate.ln(),
            self.num_timesteps as f64,
            self.sampling_steps as f64,
            self.hidden_dim as f64,
            self.num_layers as f64,
            self.time_embed_dim as f64,
            self.batch_size as f64,
            self.weight_decay.ln(),
            self.grad_clip.ln(),
        ]
    }

    fn param_names() -> Vec<&'static str> {
        vec![
            "learning_rate", "num_timesteps", "sampling_steps",
            "hidden_dim", "num_layers", "time_embed_dim",
            "batch_size", "weight_decay", "grad_clip",
        ]
    }

    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, 4.0, 2048.0) {
            if let Some(batch_bound) = bounds.get_mut(6) {
                batch_bound.1 = max_batch;
            }
        }
        // Cap hidden_dim by VRAM (index 3)
        if budget.gpu_memory_mb < 8000 {
            bounds[3] = (32.0, 256.0);
        } else if budget.gpu_memory_mb < 16000 {
            bounds[3] = (32.0, 512.0);
        } else if budget.gpu_memory_mb < 40000 {
            bounds[3] = (32.0, 1024.0);
        } else {
            // ≥40GB (L40S 46GB, H100 80GB): allow full 2048
        }
        bounds
    }
}

/// Metrics returned from Diffusion hyperopt training.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DiffusionMetrics {
    /// Validation noise prediction loss.
    pub val_loss: f64,
    /// Training noise prediction loss.
    pub train_loss: f64,
    /// Number of epochs completed.
    pub epochs_completed: usize,
}

impl Default for DiffusionMetrics {
    fn default() -> Self {
        Self {
            val_loss: 0.0,
            train_loss: 0.0,
            epochs_completed: 0,
        }
    }
}

/// Diffusion model trainer for hyperparameter optimization.
///
/// Wraps the `DiffusionTrainableAdapter` (which implements `UnifiedTrainable`)
/// and drives it through the `HyperparameterOptimizable` interface.
///
/// The Diffusion model generates price paths via denoising diffusion.
/// Training uses sequence input: `(1, seq_len, feature_dim)`.
#[derive(Debug)]
pub struct DiffusionTrainer {
    data_dir: PathBuf,
    epochs: usize,
    device: Device,
    training_paths: TrainingPaths,
    early_stopping_patience: usize,
    trial_counter: usize,
    preloaded_bars: Option<Arc<[OHLCVBar]>>,
}

impl DiffusionTrainer {
    /// Create a new Diffusion trainer.
    ///
    /// # Arguments
    /// * `data_dir` - Directory containing .dbn/.dbn.zst files
    /// * `epochs` - Number of training epochs per trial
    pub fn new<P: Into<PathBuf>>(data_dir: P, epochs: usize) -> Result<Self, MLError> {
        let data_dir = data_dir.into();
        if !data_dir.exists() {
            return Err(MLError::ConfigError(format!("Data directory not found: {}", data_dir.display())));
        }

        let device = Device::new_cuda(0).unwrap_or_else(|e| {
            warn!("CUDA unavailable ({}), falling back to CPU", e);
            Device::Cpu
        });

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

        Ok(Self {
            data_dir,
            epochs,
            device,
            training_paths: TrainingPaths::new("/tmp/ml_training", "diffusion", "default"),
            early_stopping_patience: 10,
            trial_counter: 0,
            preloaded_bars: None,
        })
    }

    /// Set training paths configuration.
    pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self {
        self.training_paths = paths;
        self
    }

    /// Configure early stopping patience.
    pub fn with_early_stopping(mut self, patience: usize) -> Self {
        self.early_stopping_patience = patience;
        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()
    }
}

impl HyperparameterOptimizable for DiffusionTrainer {
    type Params = DiffusionParams;
    type Metrics = DiffusionMetrics;

    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 Diffusion: lr={:.6}, timesteps={}, sampling={}, hidden={}, layers={}, batch={}",
            params.learning_rate, params.num_timesteps, params.sampling_steps,
            params.hidden_dim, params.num_layers, params.batch_size
        );

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

        let bars = if let Some(ref preloaded) = self.preloaded_bars {
            preloaded.to_vec()
        } 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 features = extract_ml_features(&bars)
            .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?;

        if features.len() < 2 {
            return Err(MLError::ModelError("Insufficient features extracted".to_owned()));
        }

        let feature_dim = 51;

        // Diffusion uses flat input: data_dim = seq_len * feature_dim
        // For hyperopt, we use feature_dim=1 and seq_len=feature_dim (51)
        // so data_dim = 51, matching our feature vectors.
        let pairs = crate::hyperopt::shared_data::build_flat_pairs(
            &features, &bars, feature_dim, &self.device,
        )?;

        let split = (pairs.len() as f64 * 0.8) as usize;
        let train_data = &pairs[..split];
        let val_data = &pairs[split..];

        if val_data.len() < 2 {
            return Err(MLError::ModelError("Validation set too small".to_owned()));
        }

        let config = DiffusionConfig {
            num_timesteps: params.num_timesteps,
            sampling_steps: params.sampling_steps,
            seq_len: feature_dim, // Treat feature vector as a sequence
            feature_dim: 1,       // Univariate per position
            hidden_dim: params.hidden_dim,
            num_layers: params.num_layers,
            time_embed_dim: params.time_embed_dim,
            schedule: NoiseSchedule::Cosine,
            learning_rate: params.learning_rate,
            weight_decay: params.weight_decay,
            grad_clip: params.grad_clip,
        };

        let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result<DiffusionMetrics, MLError> {
            let mut model = DiffusionTrainableAdapter::new(config, self.device.clone())?;

            let mut best_val_loss = f64::MAX;
            let mut patience_counter = 0_usize;
            let mut last_train_loss = 0.0_f64;

            for epoch in 0..self.epochs {
                let mut epoch_loss = 0.0_f64;
                let mut batch_count = 0_usize;

                for (input, target) in train_data {
                    let pred = model.forward(input)?;
                    let loss = model.compute_loss(&pred, target)?;
                    let loss_val = loss
                        .to_scalar::<f32>()
                        .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))?
                        as f64;

                    if loss_val.is_nan() || loss_val.is_infinite() {
                        return Ok(DiffusionMetrics {
                            val_loss: 1000.0,
                            train_loss: 1000.0,
                            epochs_completed: epoch,
                        });
                    }

                    model.backward(&loss)?;
                    model.optimizer_step()?;
                    epoch_loss += loss_val;
                    batch_count += 1;
                }

                last_train_loss = if batch_count > 0 {
                    epoch_loss / batch_count as f64
                } else {
                    0.0
                };

                let val_loss = model.validate(val_data)?;

                if val_loss < best_val_loss {
                    best_val_loss = val_loss;
                    patience_counter = 0;
                } else {
                    patience_counter += 1;
                }

                if epoch % 10 == 0 {
                    info!(
                        "  Epoch {}/{}: train_loss={:.6}, val_loss={:.6}",
                        epoch + 1, self.epochs, last_train_loss, val_loss
                    );
                }

                if patience_counter >= self.early_stopping_patience {
                    info!("Early stopping at epoch {}", epoch + 1);
                    return Ok(DiffusionMetrics {
                        val_loss: best_val_loss,
                        train_loss: last_train_loss,
                        epochs_completed: epoch + 1,
                    });
                }
            }

            Ok(DiffusionMetrics {
                val_loss: best_val_loss,
                train_loss: last_train_loss,
                epochs_completed: self.epochs,
            })
        }));

        let metrics = match training_result {
            Ok(Ok(m)) => m,
            Ok(Err(e)) => {
                warn!("Diffusion training failed: {}", e);
                DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 }
            }
            Err(_) => {
                warn!("Diffusion training panicked (likely OOM)");
                DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 }
            }
        };

        info!(
            "Diffusion training completed: val_loss={:.6}, train_loss={:.6}, epochs={}",
            metrics.val_loss, metrics.train_loss, metrics.epochs_completed
        );

        let duration_secs = trial_start.elapsed().as_secs_f64();
        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();
        crate::hyperopt::shared_data::write_trial_result_json(
            &self.training_paths.hyperopt_dir(),
            &trial_result,
        ).ok();

        if self.device.is_cuda() {
            std::thread::sleep(std::time::Duration::from_millis(100));
        }

        Ok(metrics)
    }

    fn extract_objective(metrics: &Self::Metrics) -> f64 {
        metrics.val_loss
    }
}

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

    #[test]
    fn test_bounds_count_matches_param_names() {
        let bounds = DiffusionParams::continuous_bounds();
        let names = DiffusionParams::param_names();
        assert_eq!(bounds.len(), names.len());
    }

    #[test]
    fn test_roundtrip_continuous() {
        let params = DiffusionParams::default();
        let continuous = params.to_continuous();
        let restored = DiffusionParams::from_continuous(&continuous).unwrap();
        assert!((params.learning_rate - restored.learning_rate).abs() < 1e-6);
        assert_eq!(params.num_timesteps, restored.num_timesteps);
        assert_eq!(params.hidden_dim, restored.hidden_dim);
        assert_eq!(params.num_layers, restored.num_layers);
    }

    #[test]
    fn test_from_continuous_wrong_length_errors() {
        let result = DiffusionParams::from_continuous(&[0.1, 0.2]);
        assert!(result.is_err());
    }

    #[test]
    fn test_bounds_are_valid() {
        for (min, max) in DiffusionParams::continuous_bounds() {
            assert!(min < max, "Invalid bounds: {min} >= {max}");
        }
    }

    #[test]
    fn test_default_within_bounds() {
        let params = DiffusionParams::default();
        let continuous = params.to_continuous();
        let bounds = DiffusionParams::continuous_bounds();
        for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() {
            assert!(
                *val >= *min && *val <= *max,
                "Param {} ({}) = {} outside [{}, {}]",
                i,
                DiffusionParams::param_names().get(i).unwrap_or(&"?"),
                val, min, max,
            );
        }
    }

    #[test]
    fn test_metrics_default() {
        let metrics = DiffusionMetrics::default();
        assert_eq!(metrics.val_loss, 0.0);
        assert_eq!(metrics.train_loss, 0.0);
        assert_eq!(metrics.epochs_completed, 0);
    }

    #[test]
    fn test_batch_size_static_upper_bound() {
        let bounds = DiffusionParams::continuous_bounds();
        // batch_size is param index 6
        let (_, max_batch) = bounds.get(6).copied().unwrap_or((4.0, 256.0));
        assert!(max_batch <= 256.0, "Max static batch_size should be ≤256");
    }
}