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

//! Hyperopt PSO adapter for Liquid CfC v2 networks
//!
//! Provides an 8-dimensional parameter space for Bayesian optimization of
//! Liquid Closed-form Continuous-time (CfC v2) neural network hyperparameters.
//!
//! ## Parameters
//!
//! | Dimension | Parameter       | Scale  | Range              |
//! |-----------|-----------------|--------|--------------------|
//! | 0         | learning_rate   | log    | [1e-5, 1e-2]      |
//! | 1         | hidden_size     | linear | [32, 512]          |
//! | 2         | backbone_depth  | linear | [1, 6]             |
//! | 3         | tau_min         | log    | [1e-3, 0.5]        |
//! | 4         | tau_max         | log    | [0.5, 10.0]        |
//! | 5         | dropout         | linear | [0.0, 0.5]         |
//! | 6         | gradient_clip   | log    | [0.1, 10.0]        |
//! | 7         | batch_size      | linear | [8, 512]           |

use std::sync::Arc;
use ml_core::device::MlDevice;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tracing::{info, warn};

use crate::features::extract_ml_features;
use crate::features::extraction::OHLCVBar;
use crate::gpu::DeviceConfig;
use crate::hyperopt::paths::TrainingPaths;
use crate::hyperopt::traits::{HardwareBudget, HyperparameterOptimizable, ParameterSpace};
use crate::liquid::candle_cfc::CfCTrainConfig;
use crate::liquid::adapter::LiquidTrainableAdapter;
use crate::training::unified_trainer::UnifiedTrainable;
use crate::MLError;

/// Liquid CfC model overhead in MB (ODE solver + backbone MLP)
const MODEL_OVERHEAD_MB: f64 = 40.0;
/// Liquid CfC per-sample memory in MB (ODE solver requires extra intermediate states)
const MB_PER_SAMPLE: f64 = 0.008;

/// Number of continuous dimensions in the Liquid CfC v2 parameter space.
const NUM_PARAMS: usize = 10;

/// Hyperparameters for Liquid CfC v2 network optimization.
///
/// All fields map to the CfC v2 architecture configuration. Log-scale
/// parameters (learning_rate, tau_min, tau_max, gradient_clip) are
/// transformed via `ln()`/`exp()` for uniform exploration across
/// orders of magnitude.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LiquidParams {
    /// Adam learning rate (log-scale, typically 1e-4 to 1e-3)
    pub learning_rate: f64,
    /// Hidden state dimension of the CfC cell
    pub hidden_size: usize,
    /// Number of layers in the backbone MLP
    pub backbone_depth: usize,
    /// Minimum time constant for ODE dynamics (log-scale)
    pub tau_min: f64,
    /// Maximum time constant for ODE dynamics (log-scale)
    pub tau_max: f64,
    /// Dropout probability applied after backbone layers
    pub dropout: f64,
    /// Maximum gradient norm for clipping (log-scale)
    pub gradient_clip: f64,
    /// Batch size for training (linear scale)
    pub batch_size: usize,
    /// 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 LiquidParams {
    fn default() -> Self {
        Self {
            learning_rate: 0.001,
            hidden_size: 128,
            backbone_depth: 2,
            tau_min: 0.01,
            tau_max: 1.0,
            dropout: 0.1,
            gradient_clip: 1.0,
            batch_size: 32,
            signal_high_bps: 10.0,
            signal_low_bps: 5.0,
        }
    }
}

impl ParameterSpace for LiquidParams {
    fn continuous_bounds() -> Vec<(f64, f64)> {
        vec![
            (1e-5_f64.ln(), 1e-2_f64.ln()), // 0: learning_rate (log scale)
            (32.0, 2048.0),                  // 1: hidden_size (linear, rounded to int)
            (1.0, 6.0),                      // 2: backbone_depth (linear, rounded to int)
            (1e-3_f64.ln(), 0.5_f64.ln()),   // 3: tau_min (log scale)
            (0.5_f64.ln(), 10.0_f64.ln()),   // 4: tau_max (log scale)
            (0.0, 0.5),                      // 5: dropout (linear)
            (0.1_f64.ln(), 10.0_f64.ln()),   // 6: gradient_clip (log scale)
            (8.0, 512.0),                    // 7: batch_size (linear)
            (1.0, 30.0),                     // 8: signal_high_bps
            (0.5, 15.0),                     // 9: signal_low_bps
        ]
    }

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

        let learning_rate = x
            .get(0)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing learning_rate at index 0".to_owned()))?
            .exp();
        let hidden_size = x
            .get(1)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing hidden_size at index 1".to_owned()))?
            .round() as usize;
        let backbone_depth = x
            .get(2)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing backbone_depth at index 2".to_owned()))?
            .round() as usize;
        let tau_min = x
            .get(3)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing tau_min at index 3".to_owned()))?
            .exp();
        let tau_max = x
            .get(4)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing tau_max at index 4".to_owned()))?
            .exp();
        let dropout = *x.get(5).ok_or_else(|| MLError::ConfigError("LiquidParams: missing dropout at index 5".to_owned()))?;
        let gradient_clip = x
            .get(6)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing gradient_clip at index 6".to_owned()))?
            .exp();
        let batch_size = x
            .get(7)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing batch_size at index 7".to_owned()))?
            .round() as usize;

        let signal_high_bps = x
            .get(8)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing signal_high_bps at index 8".to_owned()))?
            .clamp(1.0, 30.0);
        let signal_low_bps = x
            .get(9)
            .ok_or_else(|| MLError::ConfigError("LiquidParams: missing signal_low_bps at index 9".to_owned()))?
            .clamp(0.5, 15.0);

        Ok(Self {
            learning_rate,
            hidden_size,
            backbone_depth,
            tau_min,
            tau_max,
            dropout,
            gradient_clip,
            batch_size,
            signal_high_bps,
            signal_low_bps,
        })
    }

    fn to_continuous(&self) -> Vec<f64> {
        vec![
            self.learning_rate.ln(),     // 0: log scale
            self.hidden_size as f64,     // 1: linear
            self.backbone_depth as f64,  // 2: linear
            self.tau_min.ln(),           // 3: log scale
            self.tau_max.ln(),           // 4: log scale
            self.dropout,                // 5: linear
            self.gradient_clip.ln(),     // 6: log scale
            self.batch_size as f64,      // 7: linear
            self.signal_high_bps,        // 8: linear
            self.signal_low_bps,         // 9: linear
        ]
    }

    fn param_names() -> Vec<&'static str> {
        vec![
            "learning_rate",
            "hidden_size",
            "backbone_depth",
            "tau_min",
            "tau_max",
            "dropout",
            "gradient_clip",
            "batch_size",
            "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, 8.0, 2048.0) {
            if let Some(batch_bound) = bounds.get_mut(7) {
                batch_bound.1 = max_batch;
            }
        }
        // Cap hidden_size by VRAM (index 1)
        if budget.gpu_memory_mb < 8000 {
            bounds[1] = (32.0, 256.0);    // Small GPU: cap hidden at 256
        } else if budget.gpu_memory_mb < 16000 {
            bounds[1] = (32.0, 512.0);    // Medium: cap at 512
        } else if budget.gpu_memory_mb < 40000 {
            bounds[1] = (32.0, 1024.0);   // Large: cap at 1024
        } else {
            // ≥40GB (L40S 46GB, H100 80GB): allow full 2048
        }
        bounds
    }
}

/// Metrics returned from Liquid CfC hyperopt training.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct LiquidMetrics {
    /// Validation loss (primary objective for optimization)
    pub val_loss: f64,
    /// Training loss at end of run
    pub train_loss: 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>,
}

/// Liquid CfC v2 trainer for hyperparameter optimization.
///
/// Wraps the `LiquidTrainableAdapter` (which implements `UnifiedTrainable`)
/// and drives it through the `HyperparameterOptimizable` interface.
#[derive(Debug)]
pub struct LiquidTrainer {
    data_dir: PathBuf,
    epochs: usize,
    device: MlDevice,
    training_paths: TrainingPaths,
    early_stopping_patience: usize,
    trial_counter: usize,
    preloaded_bars: Option<Arc<[OHLCVBar]>>,
}

impl LiquidTrainer {
    /// Create a new Liquid CfC 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 = MlDevice::cuda(0)
            .map_err(|e| MLError::ConfigError(format!("CUDA GPU required for Liquid hyperopt: {}", e)))?;

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

        Ok(Self {
            data_dir,
            epochs,
            device,
            training_paths: TrainingPaths::new("/tmp/ml_training", "liquid", "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 LiquidTrainer {
    type Params = LiquidParams;
    type Metrics = LiquidMetrics;

    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 Liquid CfC: lr={:.6}, hidden={}, depth={}, tau=[{:.3},{:.3}], dropout={:.3}",
            params.learning_rate, params.hidden_size, params.backbone_depth,
            params.tau_min, params.tau_max, params.dropout
        );

        self.training_paths.create_all().map_err(|e| {
            MLError::ModelError(format!("Failed to create training directories: {}", 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)))?
        };

        // Extract 51-dim features
        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()));
        }

        // Build (input, target) pairs: input = feature[i], target = close price change
        let feature_dim = 51;
        let pairs = crate::hyperopt::shared_data::build_flat_pairs(&features, &bars, feature_dim, &self.device)?;

        // Split train/val 80/20
        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()));
        }

        // Build backbone hidden sizes from depth param
        let backbone = vec![params.hidden_size; params.backbone_depth];

        let config = CfCTrainConfig {
            input_size: feature_dim,
            hidden_size: params.hidden_size,
            output_size: 1,
            backbone_hidden_sizes: backbone,
            tau_min: params.tau_min,
            tau_max: params.tau_max,
            learning_rate: params.learning_rate,
            batch_size: params.batch_size,
            seq_len: 1,
            max_epochs: self.epochs,
            early_stopping_patience: self.early_stopping_patience,
            dropout: params.dropout,
            gradient_clip: params.gradient_clip,
            market_regime_adaptation: false,
            device: DeviceConfig::Cuda(0),
        };

        let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result<LiquidMetrics, MLError> {
            let mut model = LiquidTrainableAdapter::new(config)?;

            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 loss_accum = 0.0_f64;
                let mut batch_count = 0_usize;

                let stream = self.device.cuda_stream()
                    .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
                for (input, target) in train_data {
                    let input_host = input.to_host(stream)?;
                    let target_host = target.to_host(stream)?;
                    let loss_val = model.forward_loss(&input_host, &target_host)?;

                    if batch_count == 0 && (loss_val.is_nan() || loss_val.is_infinite()) {
                        return Ok(LiquidMetrics {
                            val_loss: 1000.0,
                            train_loss: 1000.0,
                            epochs_completed: epoch,
                            backtest_sharpe: None,
                            backtest_trades: None,
                        });
                    }

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

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

                // Validate
                let mut vl_accum = 0.0_f64;
                let mut vl_count = 0_usize;
                for (vi, vt) in val_data {
                    let vi_h = vi.to_host(stream)?;
                    let vt_h = vt.to_host(stream)?;
                    if let Ok(vl) = model.forward_loss(&vi_h, &vt_h) {
                        vl_accum += vl;
                        vl_count += 1;
                    }
                }
                let val_loss = if vl_count > 0 { vl_accum / vl_count as f64 } else { 1000.0 };

                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(LiquidMetrics {
                        val_loss: best_val_loss,
                        train_loss: last_train_loss,
                        epochs_completed: epoch + 1,
                        backtest_sharpe: None,
                        backtest_trades: None,
                    });
                }
            }

            Ok(LiquidMetrics {
                val_loss: best_val_loss,
                train_loss: last_train_loss,
                epochs_completed: self.epochs,
                backtest_sharpe: None,
                backtest_trades: None,
            })
        }));

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

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

        // Write trial result
        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();

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

        Ok(metrics)
    }

    fn extract_objective(metrics: &Self::Metrics) -> f64 {
        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);
        }
        metrics.val_loss
    }
}

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

    #[test]
    fn test_liquid_params_default() {
        let params = LiquidParams::default();
        assert!((params.learning_rate - 0.001).abs() < 1e-10);
        assert_eq!(params.hidden_size, 128);
        assert_eq!(params.backbone_depth, 2);
        assert!((params.tau_min - 0.01).abs() < 1e-10);
        assert!((params.tau_max - 1.0).abs() < 1e-10);
        assert!((params.dropout - 0.1).abs() < 1e-10);
        assert!((params.gradient_clip - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_liquid_params_roundtrip() {
        let params = LiquidParams::default();
        let continuous = params.to_continuous();
        let restored = LiquidParams::from_continuous(&continuous).unwrap();
        assert!((params.learning_rate - restored.learning_rate).abs() < 1e-6);
        assert_eq!(params.hidden_size, restored.hidden_size);
        assert_eq!(params.backbone_depth, restored.backbone_depth);
        assert!((params.tau_min - restored.tau_min).abs() < 1e-6);
        assert!((params.tau_max - restored.tau_max).abs() < 1e-6);
        assert!((params.dropout - restored.dropout).abs() < 1e-10);
        assert!((params.gradient_clip - restored.gradient_clip).abs() < 1e-6);
    }

    #[test]
    fn test_liquid_params_bounds() {
        let bounds = LiquidParams::continuous_bounds();
        assert_eq!(bounds.len(), NUM_PARAMS);
        for (lo, hi) in &bounds {
            assert!(lo < hi, "Lower bound {} must be less than upper {}", lo, hi);
        }
    }

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

    #[test]
    fn test_liquid_params_within_bounds() {
        let params = LiquidParams::default();
        let continuous = params.to_continuous();
        let bounds = LiquidParams::continuous_bounds();

        for (i, (val, (lo, hi))) in continuous.iter().zip(bounds.iter()).enumerate() {
            assert!(
                val >= lo && val <= hi,
                "Param {} ({}): value {} outside bounds [{}, {}]",
                i,
                LiquidParams::param_names().get(i).unwrap_or(&"unknown"),
                val,
                lo,
                hi
            );
        }
    }

    #[test]
    fn test_liquid_params_wrong_dimension() {
        let too_short = vec![0.0; 3];
        let result = LiquidParams::from_continuous(&too_short);
        assert!(result.is_err());

        let too_long = vec![0.0; 12];
        let result = LiquidParams::from_continuous(&too_long);
        assert!(result.is_err());
    }

    #[test]
    fn test_liquid_params_extreme_values() {
        let bounds = LiquidParams::continuous_bounds();
        // Build a vector from lower bounds
        let lower: Vec<f64> = bounds.iter().map(|(lo, _)| *lo).collect();
        let params_lo = LiquidParams::from_continuous(&lower).unwrap();
        assert!(params_lo.learning_rate > 0.0);
        assert!(params_lo.hidden_size >= 32);
        assert!(params_lo.backbone_depth >= 1);

        // Build a vector from upper bounds
        let upper: Vec<f64> = bounds.iter().map(|(_, hi)| *hi).collect();
        let params_hi = LiquidParams::from_continuous(&upper).unwrap();
        assert!(params_hi.learning_rate <= 0.011); // 1e-2 with float tolerance
        assert!(params_hi.hidden_size <= 2048);
        assert!(params_hi.backbone_depth <= 6);
    }
}