foxhunt/crates/ml-supervised/src/liquid/candle_cfc.rs

//! GPU-native `CfC` v2 (Closed-form Continuous-time) Neural Network
//!
//! Uses `GpuTensor` / `GpuLinear` (cudarc + cuBLAS) instead of Candle tensors.
//! This is the training path; the existing `FixedPoint` implementation in cells.rs/network.rs
//! remains the production inference path.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use serde::{Deserialize, Serialize};

use ml_core::MLError;

use crate::gpu_tensor::{
    gpu_add, gpu_cat_dim1, gpu_exp, gpu_mul, gpu_recip, gpu_scale, gpu_select_dim1,
    gpu_sigmoid, gpu_tanh, GpuLinear, GpuTensor,
};

/// Device configuration for `CfC` training -- re-exported from central gpu module.
pub use ml_core::gpu::DeviceConfig;

/// `CfC` v2 training configuration
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CfCTrainConfig {
    pub input_size: usize,
    pub hidden_size: usize,
    pub output_size: usize,
    pub backbone_hidden_sizes: Vec<usize>,
    pub tau_min: f64,
    pub tau_max: f64,
    pub learning_rate: f64,
    pub batch_size: usize,
    pub seq_len: usize,
    pub max_epochs: usize,
    pub early_stopping_patience: usize,
    pub dropout: f64,
    pub gradient_clip: f64,
    pub market_regime_adaptation: bool,
    pub device: DeviceConfig,
}

impl Default for CfCTrainConfig {
    fn default() -> Self {
        Self {
            input_size: 225,
            hidden_size: 128,
            output_size: 3,
            backbone_hidden_sizes: vec![128, 128],
            tau_min: 0.01,
            tau_max: 1.0,
            learning_rate: 0.001,
            batch_size: 64,
            seq_len: 128,
            max_epochs: 100,
            early_stopping_patience: 10,
            dropout: 0.1,
            gradient_clip: 1.0,
            market_regime_adaptation: true,
            device: DeviceConfig::Auto,
        }
    }
}

/// Backbone MLP that produces both f(x,h) and tau(x,h) outputs.
#[allow(missing_debug_implementations)]
pub struct BackboneMLP {
    layers: Vec<GpuLinear>,
    f_head: GpuLinear,
    tau_head: GpuLinear,
}

impl BackboneMLP {
    pub fn new(
        input_dim: usize,
        hidden_sizes: &[usize],
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        for &hidden_size in hidden_sizes {
            let layer = GpuLinear::new(current_dim, hidden_size, stream)?;
            layers.push(layer);
            current_dim = hidden_size;
        }

        let last_hidden = hidden_sizes.last().copied().ok_or_else(|| {
            MLError::ConfigError("backbone_hidden_sizes cannot be empty".to_owned())
        })?;

        let f_head = GpuLinear::new(last_hidden, last_hidden, stream)?;
        let tau_head = GpuLinear::new(last_hidden, last_hidden, stream)?;

        Ok(Self {
            layers,
            f_head,
            tau_head,
        })
    }

    pub fn forward(&self, input: &GpuTensor) -> Result<(GpuTensor, GpuTensor), MLError> {
        let mut x = input.clone();
        for layer in &self.layers {
            x = layer.forward(&x)?;
            x = gpu_tanh(&x)?;
        }

        let f_out = gpu_tanh(&self.f_head.forward(&x)?)?;
        let tau_out = self.tau_head.forward(&x)?;

        Ok((f_out, tau_out))
    }
}

/// `CfC` v2 Cell -- Closed-form Continuous-time cell
///
/// Update rule:
///   h(t+dt) = h(t) * exp(-dt/tau(t)) + f(x,h) * (1 - exp(-dt/tau(t)))
#[allow(missing_debug_implementations)]
pub struct CfCCell {
    backbone: BackboneMLP,
    pub hidden_size: usize,
    tau_min: f64,
    tau_max: f64,
}

impl CfCCell {
    pub fn new(config: &CfCTrainConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let concat_dim = config.input_size + config.hidden_size;
        let backbone = BackboneMLP::new(concat_dim, &config.backbone_hidden_sizes, stream)?;

        let last_backbone = config.backbone_hidden_sizes.last().copied()
            .ok_or_else(|| MLError::ConfigError("backbone_hidden_sizes cannot be empty".to_owned()))?;
        if last_backbone != config.hidden_size {
            return Err(MLError::ConfigError(format!(
                "Last backbone hidden size ({}) must equal hidden_size ({})",
                last_backbone, config.hidden_size
            )));
        }

        Ok(Self {
            backbone,
            hidden_size: config.hidden_size,
            tau_min: config.tau_min,
            tau_max: config.tau_max,
        })
    }

    /// Single `CfC` step: `h_new` = h * decay + f * (1 - decay)
    pub fn step(&self, x: &GpuTensor, h: &GpuTensor, dt: f32) -> Result<GpuTensor, MLError> {
        if dt < 0.0 {
            return Err(MLError::InferenceError("dt must be non-negative".to_owned()));
        }

        // Concatenate input and hidden state: [batch, input_size + hidden_size]
        let xh = gpu_cat_dim1(x, h)?;
        let (f_out, tau_raw) = self.backbone.forward(&xh)?;

        // tau(t) = sigmoid(tau_raw) * (tau_max - tau_min) + tau_min
        let tau = gpu_sigmoid(&tau_raw)?;
        let tau_range = (self.tau_max - self.tau_min) as f32;
        let tau = gpu_scale(&tau, tau_range)?;
        let tau = crate::gpu_tensor::gpu_add_scalar(&tau, self.tau_min as f32)?;

        // decay = exp(-dt / tau)
        let tau_inv = gpu_recip(&tau)?;
        let neg_dt_over_tau = gpu_scale(&tau_inv, -dt)?;
        let decay = gpu_exp(&neg_dt_over_tau)?;

        // h_new = h * decay + f * (1 - decay)
        let one_minus_decay = crate::gpu_tensor::gpu_scalar_sub(1.0, &decay)?;
        let h_decay = gpu_mul(h, &decay)?;
        let f_contrib = gpu_mul(&f_out, &one_minus_decay)?;
        let h_new = gpu_add(&h_decay, &f_contrib)?;

        Ok(h_new)
    }
}

/// Complete `CfC` v2 Network for sequence processing
#[allow(missing_debug_implementations)]
pub struct CandleCfCNetwork {
    cell: CfCCell,
    output_layer: GpuLinear,
    config: CfCTrainConfig,
    stream: Arc<CudaStream>,
}

impl CandleCfCNetwork {
    pub fn new(config: &CfCTrainConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if config.input_size == 0 {
            return Err(MLError::ConfigError("Liquid CfC requires input_size > 0".to_owned()));
        }
        if config.hidden_size == 0 {
            return Err(MLError::ConfigError("Liquid CfC requires hidden_size > 0".to_owned()));
        }

        let cell = CfCCell::new(config, stream)?;
        let output_layer = GpuLinear::new(config.hidden_size, config.output_size, stream)?;

        Ok(Self {
            cell,
            output_layer,
            config: config.clone(),
            stream: Arc::clone(stream),
        })
    }

    /// Process full sequence `[batch, seq_len, features]` -> `[batch, output_size]`
    pub fn forward_sequence(&self, input: &GpuTensor, dt: f32) -> Result<GpuTensor, MLError> {
        if input.shape.len() != 3 {
            return Err(MLError::InvalidInput(format!(
                "Expected 3D input [batch, seq_len, features], got {:?}",
                input.shape
            )));
        }
        let batch_size = input.dim(0)?;
        let seq_len = input.dim(1)?;
        let _features = input.dim(2)?;

        let mut h = GpuTensor::zeros(&[batch_size, self.config.hidden_size], &self.stream)?;

        // Extract timesteps on GPU via gpu_select_dim1
        for t in 0..seq_len {
            let x_t = gpu_select_dim1(input, t)?;
            h = self.cell.step(&x_t, &h, dt)?;
        }

        self.output_layer.forward(&h)
    }

    /// Forward compatible with `UnifiedTrainable` (3D input, default dt=0.01)
    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
        self.forward_sequence(input, 0.01)
    }

    /// Approximate parameter count for reporting purposes
    pub fn param_count(&self) -> usize {
        let backbone_params =
            self.config
                .backbone_hidden_sizes
                .iter()
                .enumerate()
                .fold(0, |acc, (i, &size)| {
                    let in_d = if i == 0 {
                        self.config.input_size + self.config.hidden_size
                    } else {
                        self.config
                            .backbone_hidden_sizes
                            .get(i.saturating_sub(1))
                            .copied()
                            .unwrap_or(size)
                    };
                    acc + in_d * size + size
                });
        let last_h = self
            .config
            .backbone_hidden_sizes
            .last()
            .copied()
            .unwrap_or(0);
        let heads = 2 * (last_h * last_h + last_h);
        let output = self.config.hidden_size * self.config.output_size + self.config.output_size;
        backbone_params + heads + output
    }
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_cfc_train_config_default() {
        let config = CfCTrainConfig::default();
        assert_eq!(config.input_size, 225);
        assert_eq!(config.hidden_size, 128);
        assert_eq!(config.output_size, 3);
        assert_eq!(config.backbone_hidden_sizes, vec![128, 128]);
        assert!(config.tau_min > 0.0);
        assert!(config.tau_max > config.tau_min);
    }

    #[test]
    fn test_cfc_network_creation() {
        let stream = test_stream();
        let config = CfCTrainConfig::default();
        let network = CandleCfCNetwork::new(&config, &stream).unwrap();
        assert!(network.param_count() > 0);
    }

    #[test]
    fn test_cfc_network_forward_sequence() {
        let stream = test_stream();
        let config = CfCTrainConfig {
            input_size: 16,
            hidden_size: 32,
            output_size: 3,
            backbone_hidden_sizes: vec![32],
            seq_len: 10,
            batch_size: 4,
            ..CfCTrainConfig::default()
        };
        let network = CandleCfCNetwork::new(&config, &stream).unwrap();
        let input = GpuTensor::randn(&[4, 10, 16], 1.0, &stream).unwrap();
        let output = network.forward_sequence(&input, 0.01).unwrap();
        assert_eq!(output.shape, vec![4, 3]);
    }

    #[test]
    fn test_cfc_network_forward_trait() {
        let stream = test_stream();
        let config = CfCTrainConfig {
            input_size: 16,
            hidden_size: 32,
            output_size: 3,
            backbone_hidden_sizes: vec![32],
            seq_len: 10,
            ..CfCTrainConfig::default()
        };
        let network = CandleCfCNetwork::new(&config, &stream).unwrap();
        let input = GpuTensor::randn(&[4, 10, 16], 1.0, &stream).unwrap();
        let output = network.forward(&input).unwrap();
        assert_eq!(output.shape, vec![4, 3]);
    }

    #[test]
    fn test_cfc_cell_step() {
        let stream = test_stream();
        let config = CfCTrainConfig {
            input_size: 8,
            hidden_size: 16,
            output_size: 3,
            backbone_hidden_sizes: vec![16],
            ..CfCTrainConfig::default()
        };
        let cell = CfCCell::new(&config, &stream).unwrap();

        let x = GpuTensor::randn(&[4, 8], 1.0, &stream).unwrap();
        let h = GpuTensor::zeros(&[4, 16], &stream).unwrap();
        let h_new = cell.step(&x, &h, 0.01).unwrap();
        assert_eq!(h_new.shape, vec![4, 16]);
    }

    #[test]
    fn test_backbone_empty_hidden_sizes_errors() {
        let stream = test_stream();
        let result = BackboneMLP::new(100, &[], &stream);
        assert!(result.is_err());
    }

    #[test]
    fn test_device_config_resolve_cpu() {
        let cpu = DeviceConfig::Cpu;
        let device = cpu.resolve().unwrap();
        assert!(device.is_cpu());
    }

    #[test]
    fn test_device_config_resolve_auto() {
        let auto = DeviceConfig::Auto;
        let device = auto.resolve().unwrap();
        drop(device);
    }

    #[test]
    fn test_cfc_train_config_serde_roundtrip() {
        let config = CfCTrainConfig::default();
        let json = serde_json::to_string(&config).unwrap();
        let deserialized: CfCTrainConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.input_size, config.input_size);
        assert_eq!(deserialized.hidden_size, config.hidden_size);
    }
}