kizzasi_core/

training.rs

//! Training infrastructure for SSM models
//!
//! Provides trainable versions of SSM components with automatic differentiation,
//! loss functions, and optimization utilities using candle-core.
//!
//! # Features
//!
//! - **TrainableSSM**: Differentiable SSM model with automatic gradient tracking
//! - **Trainer**: Full training loop with scheduler, metrics, and validation
//! - **Loss Functions**: MSE, MAE, Huber, Cross-Entropy
//! - **LR Scheduling**: Integrated support for all scheduler types
//! - **Metrics Tracking**: Automatic loss, LR, and gradient monitoring
//! - **Early Stopping**: Validation-based early stopping with patience

use crate::config::KizzasiConfig;
use crate::dataloader::TimeSeriesDataLoader;
use crate::device::DeviceConfig;
use crate::error::{CoreError, CoreResult};
use crate::metrics::{MetricsLogger, TrainingMetrics};
use crate::scheduler::LRScheduler;
use candle_core::{DType, Device, Tensor, Var};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};

/// Scheduler type enumeration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SchedulerType {
    Constant,
    Linear {
        warmup_steps: usize,
        final_lr: f64,
    },
    Cosine {
        warmup_steps: usize,
        min_lr: f64,
    },
    Step {
        milestones: Vec<usize>,
        decay_factor: f64,
    },
    Exponential {
        decay_rate: f64,
        decay_steps: usize,
    },
    OneCycle {
        warmup_pct: f64,
    },
    Polynomial {
        final_lr: f64,
        power: f64,
    },
}

/// Mixed precision training mode
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MixedPrecision {
    /// Full precision (FP32)
    None,
    /// Half precision (FP16) - faster but less stable
    FP16,
    /// Brain float 16 (BF16) - better stability than FP16
    BF16,
}

impl MixedPrecision {
    /// Convert to candle DType
    pub fn to_dtype(&self) -> DType {
        match self {
            MixedPrecision::None => DType::F32,
            MixedPrecision::FP16 => DType::F16,
            MixedPrecision::BF16 => DType::BF16,
        }
    }

    /// Check if mixed precision is enabled
    pub fn is_enabled(&self) -> bool {
        !matches!(self, MixedPrecision::None)
    }
}

/// Configuration for training
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Device configuration (CPU/CUDA/Metal)
    pub device_config: DeviceConfig,
    /// Learning rate (initial for schedulers)
    pub learning_rate: f64,
    /// Batch size
    pub batch_size: usize,
    /// Number of epochs
    pub epochs: usize,
    /// Weight decay (L2 regularization)
    pub weight_decay: f64,
    /// Gradient clipping threshold
    pub grad_clip: Option<f32>,
    /// Beta1 for Adam optimizer
    pub beta1: f64,
    /// Beta2 for Adam optimizer
    pub beta2: f64,
    /// Epsilon for Adam optimizer
    pub eps: f64,
    /// Learning rate scheduler type
    pub scheduler: Option<SchedulerType>,
    /// Enable metrics tracking
    pub track_metrics: bool,
    /// Log interval (batches)
    pub log_interval: usize,
    /// Validation split (0.0 to 1.0)
    pub validation_split: f32,
    /// Early stopping patience (epochs)
    pub early_stopping_patience: Option<usize>,
    /// Enable gradient checkpointing (saves memory by recomputing activations)
    pub use_gradient_checkpointing: bool,
    /// Checkpoint every N layers (None = checkpoint all layers)
    pub checkpoint_segment_size: Option<usize>,
    /// Mixed precision training mode
    pub mixed_precision: MixedPrecision,
    /// Loss scaling factor for mixed precision (to prevent underflow)
    pub loss_scale: f32,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            device_config: DeviceConfig::default(),
            learning_rate: 1e-4,
            batch_size: 32,
            epochs: 10,
            weight_decay: 1e-2,
            grad_clip: Some(1.0),
            beta1: 0.9,
            beta2: 0.999,
            eps: 1e-8,
            scheduler: None,
            track_metrics: true,
            log_interval: 10,
            validation_split: 0.2,
            early_stopping_patience: Some(5),
            use_gradient_checkpointing: false,
            checkpoint_segment_size: Some(2), // Checkpoint every 2 layers by default
            mixed_precision: MixedPrecision::None,
            loss_scale: 1.0, // No scaling by default
        }
    }
}

impl TrainingConfig {
    /// Set scheduler type
    pub fn with_scheduler(mut self, scheduler: SchedulerType) -> Self {
        self.scheduler = Some(scheduler);
        self
    }

    /// Disable metrics tracking
    pub fn without_metrics(mut self) -> Self {
        self.track_metrics = false;
        self
    }

    /// Set validation split
    pub fn with_validation_split(mut self, split: f32) -> Self {
        self.validation_split = split;
        self
    }

    /// Set early stopping patience
    pub fn with_early_stopping(mut self, patience: usize) -> Self {
        self.early_stopping_patience = Some(patience);
        self
    }

    /// Disable early stopping
    pub fn without_early_stopping(mut self) -> Self {
        self.early_stopping_patience = None;
        self
    }

    /// Enable gradient checkpointing for memory-efficient training
    pub fn with_gradient_checkpointing(mut self, segment_size: Option<usize>) -> Self {
        self.use_gradient_checkpointing = true;
        self.checkpoint_segment_size = segment_size;
        self
    }

    /// Disable gradient checkpointing
    pub fn without_gradient_checkpointing(mut self) -> Self {
        self.use_gradient_checkpointing = false;
        self
    }

    /// Enable mixed precision training (FP16)
    pub fn with_fp16(mut self) -> Self {
        self.mixed_precision = MixedPrecision::FP16;
        self.loss_scale = 128.0; // Default loss scale for FP16
        self
    }

    /// Enable mixed precision training (BF16)
    pub fn with_bf16(mut self) -> Self {
        self.mixed_precision = MixedPrecision::BF16;
        self.loss_scale = 1.0; // BF16 is more stable, doesn't need scaling
        self
    }

    /// Set mixed precision mode
    pub fn with_mixed_precision(mut self, mode: MixedPrecision, loss_scale: f32) -> Self {
        self.mixed_precision = mode;
        self.loss_scale = loss_scale;
        self
    }

    /// Disable mixed precision training
    pub fn without_mixed_precision(mut self) -> Self {
        self.mixed_precision = MixedPrecision::None;
        self.loss_scale = 1.0;
        self
    }
}

/// Trainable Selective SSM using candle Tensors
pub struct TrainableSSM {
    config: KizzasiConfig,
    training_config: TrainingConfig,
    device: Device,
    dtype: DType,
    // Learnable parameters
    embedding_weight: Var,
    a_matrices: Vec<Var>,
    b_matrices: Vec<Var>,
    c_matrices: Vec<Var>,
    d_vectors: Vec<Var>,
    output_proj: Var,
    // Layer normalization parameters
    ln_gamma: Vec<Var>,
    ln_beta: Vec<Var>,
    // Variable map for optimizer
    varmap: VarMap,
}

impl TrainableSSM {
    /// Create a new trainable SSM model
    pub fn new(config: KizzasiConfig, training_config: TrainingConfig) -> CoreResult<Self> {
        // Create device from configuration
        let device = training_config.device_config.create_device()?;

        // Use mixed precision dtype from training config
        let dtype = training_config.mixed_precision.to_dtype();

        let hidden_dim = config.get_hidden_dim();
        let state_dim = config.get_state_dim();
        let num_layers = config.get_num_layers();
        let input_dim = config.get_input_dim();
        let output_dim = config.get_output_dim();

        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, dtype, &device);

        // Initialize embedding layer
        let embedding_weight_tensor = vb
            .get_with_hints(
                (input_dim, hidden_dim),
                "embedding.weight",
                candle_nn::init::DEFAULT_KAIMING_NORMAL,
            )
            .map_err(|e| CoreError::Generic(format!("Failed to create embedding: {}", e)))?;
        let embedding_weight = Var::from_tensor(&embedding_weight_tensor)
            .map_err(|e| CoreError::Generic(format!("Failed to create embedding var: {}", e)))?;

        // Initialize SSM matrices for each layer
        let mut a_matrices = Vec::with_capacity(num_layers);
        let mut b_matrices = Vec::with_capacity(num_layers);
        let mut c_matrices = Vec::with_capacity(num_layers);
        let mut d_vectors = Vec::with_capacity(num_layers);
        let mut ln_gamma = Vec::with_capacity(num_layers);
        let mut ln_beta = Vec::with_capacity(num_layers);

        for layer_idx in 0..num_layers {
            // A matrix: state transition (initialized for stability)
            let a_tensor = vb
                .get_with_hints(
                    (hidden_dim, state_dim),
                    &format!("ssm.layer_{}.a", layer_idx),
                    candle_nn::init::Init::Const(-0.5),
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create A matrix: {}", e)))?;
            let a = Var::from_tensor(&a_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create A var: {}", e)))?;
            a_matrices.push(a);

            // B matrix: input projection to state
            let b_tensor = vb
                .get_with_hints(
                    (hidden_dim, state_dim),
                    &format!("ssm.layer_{}.b", layer_idx),
                    candle_nn::init::DEFAULT_KAIMING_NORMAL,
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create B matrix: {}", e)))?;
            let b = Var::from_tensor(&b_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create B var: {}", e)))?;
            b_matrices.push(b);

            // C matrix: state to output projection
            let c_tensor = vb
                .get_with_hints(
                    (hidden_dim, state_dim),
                    &format!("ssm.layer_{}.c", layer_idx),
                    candle_nn::init::DEFAULT_KAIMING_NORMAL,
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create C matrix: {}", e)))?;
            let c = Var::from_tensor(&c_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create C var: {}", e)))?;
            c_matrices.push(c);

            // D vector: skip connection
            let d_tensor = vb
                .get_with_hints(
                    hidden_dim,
                    &format!("ssm.layer_{}.d", layer_idx),
                    candle_nn::init::Init::Const(1.0),
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create D vector: {}", e)))?;
            let d = Var::from_tensor(&d_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create D var: {}", e)))?;
            d_vectors.push(d);

            // Layer normalization parameters
            let gamma_tensor = vb
                .get_with_hints(
                    hidden_dim,
                    &format!("ln.layer_{}.gamma", layer_idx),
                    candle_nn::init::Init::Const(1.0),
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create LN gamma: {}", e)))?;
            let gamma = Var::from_tensor(&gamma_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create LN gamma var: {}", e)))?;
            ln_gamma.push(gamma);

            let beta_tensor = vb
                .get_with_hints(
                    hidden_dim,
                    &format!("ln.layer_{}.beta", layer_idx),
                    candle_nn::init::Init::Const(0.0),
                )
                .map_err(|e| CoreError::Generic(format!("Failed to create LN beta: {}", e)))?;
            let beta = Var::from_tensor(&beta_tensor)
                .map_err(|e| CoreError::Generic(format!("Failed to create LN beta var: {}", e)))?;
            ln_beta.push(beta);
        }

        // Output projection
        let output_proj_tensor = vb
            .get_with_hints(
                (hidden_dim, output_dim),
                "output.proj",
                candle_nn::init::DEFAULT_KAIMING_NORMAL,
            )
            .map_err(|e| {
                CoreError::Generic(format!("Failed to create output projection: {}", e))
            })?;
        let output_proj = Var::from_tensor(&output_proj_tensor)
            .map_err(|e| CoreError::Generic(format!("Failed to create output proj var: {}", e)))?;

        Ok(Self {
            config,
            training_config,
            device,
            dtype,
            embedding_weight,
            a_matrices,
            b_matrices,
            c_matrices,
            d_vectors,
            output_proj,
            ln_gamma,
            ln_beta,
            varmap,
        })
    }

    /// Forward pass for training (tracks gradients)
    ///
    /// # Arguments
    /// * `input` - Input tensor of shape [batch_size, seq_len, input_dim]
    /// * `targets` - Optional target tensor for loss computation
    ///
    /// # Returns
    /// Output tensor of shape [batch_size, seq_len, output_dim]
    pub fn forward(&self, input: &Tensor) -> CoreResult<Tensor> {
        // Embed input: [batch, seq, input_dim] -> [batch, seq, hidden_dim]
        // Reshape input to [batch * seq, input_dim] for matmul, then reshape back
        let batch_size = input
            .dim(0)
            .map_err(|e| CoreError::Generic(format!("Failed to get batch dimension: {}", e)))?;
        let seq_len = input
            .dim(1)
            .map_err(|e| CoreError::Generic(format!("Failed to get sequence dimension: {}", e)))?;
        let input_dim = input
            .dim(2)
            .map_err(|e| CoreError::Generic(format!("Failed to get input dimension: {}", e)))?;

        let x_flat = input
            .reshape((batch_size * seq_len, input_dim))
            .map_err(|e| CoreError::Generic(format!("Failed to reshape input: {}", e)))?;

        let hidden_dim = self.config.get_hidden_dim();
        let x_embedded = x_flat
            .matmul(self.embedding_weight.as_tensor())
            .map_err(|e| CoreError::Generic(format!("Embedding forward failed: {}", e)))?;

        let x = x_embedded
            .reshape((batch_size, seq_len, hidden_dim))
            .map_err(|e| CoreError::Generic(format!("Failed to reshape embedded: {}", e)))?;

        // Initialize hidden state
        let state_dim = self.config.get_state_dim();

        let mut h = Tensor::zeros(
            (batch_size, hidden_dim, state_dim),
            self.dtype,
            &self.device,
        )
        .map_err(|e| CoreError::Generic(format!("Failed to create hidden state: {}", e)))?;

        let mut x = x;

        // Process through each layer
        for layer_idx in 0..self.config.get_num_layers() {
            x = self.layer_norm(&x, layer_idx)?;
            x = self.ssm_layer(&x, &mut h, layer_idx)?;
        }

        // Project to output dimension: [batch, seq, hidden_dim] -> [batch, seq, output_dim]
        // Reshape to [batch * seq, hidden_dim], matmul, then reshape back
        let x_flat = x
            .reshape((batch_size * seq_len, hidden_dim))
            .map_err(|e| CoreError::Generic(format!("Failed to reshape for output: {}", e)))?;

        let output_dim = self.config.get_output_dim();
        let output_flat = x_flat
            .matmul(self.output_proj.as_tensor())
            .map_err(|e| CoreError::Generic(format!("Output projection failed: {}", e)))?;

        let output = output_flat
            .reshape((batch_size, seq_len, output_dim))
            .map_err(|e| CoreError::Generic(format!("Failed to reshape output: {}", e)))?;

        Ok(output)
    }

    /// Apply layer normalization
    fn layer_norm(&self, x: &Tensor, layer_idx: usize) -> CoreResult<Tensor> {
        const EPS: f64 = 1e-5;

        // Compute mean and variance along the last dimension
        let mean = x
            .mean_keepdim(candle_core::D::Minus1)
            .map_err(|e| CoreError::Generic(format!("Layer norm mean failed: {}", e)))?;
        let x_centered = x.broadcast_sub(&mean).map_err(|e| {
            CoreError::Generic(format!("Layer norm variance computation failed: {}", e))
        })?;
        let variance = x_centered
            .sqr()
            .map_err(|e| CoreError::Generic(format!("Layer norm variance sqr failed: {}", e)))?
            .mean_keepdim(candle_core::D::Minus1)
            .map_err(|e| CoreError::Generic(format!("Layer norm variance mean failed: {}", e)))?;

        // Normalize: (x - mean) / sqrt(variance + eps)
        let std = (variance.affine(1.0, EPS))
            .map_err(|e| CoreError::Generic(format!("Layer norm variance add eps failed: {}", e)))?
            .sqrt()
            .map_err(|e| CoreError::Generic(format!("Layer norm sqrt failed: {}", e)))?;

        let normalized = x_centered
            .broadcast_div(&std)
            .map_err(|e| CoreError::Generic(format!("Layer norm division failed: {}", e)))?;

        // Apply affine transformation
        let gamma = self.ln_gamma[layer_idx].as_tensor();
        let beta = self.ln_beta[layer_idx].as_tensor();

        normalized
            .broadcast_mul(gamma)
            .map_err(|e| CoreError::Generic(format!("Layer norm gamma mul failed: {}", e)))?
            .broadcast_add(beta)
            .map_err(|e| CoreError::Generic(format!("Layer norm beta add failed: {}", e)))
    }

    /// SSM layer computation
    fn ssm_layer(&self, x: &Tensor, _h: &mut Tensor, layer_idx: usize) -> CoreResult<Tensor> {
        let _a = self.a_matrices[layer_idx].as_tensor();
        let _b = self.b_matrices[layer_idx].as_tensor();
        let _c = self.c_matrices[layer_idx].as_tensor();
        let d = self.d_vectors[layer_idx].as_tensor();

        // Simplified SSM step (full implementation would include selective scan)
        // For now, implementing a basic skip connection
        // TODO: Implement proper selective scan mechanism with state evolution

        // For training, we process the entire sequence in parallel (teacher forcing)
        // Output: y = D * x (simplified - full version uses state)
        let y = x
            .broadcast_mul(d)
            .map_err(|e| CoreError::Generic(format!("Skip connection failed: {}", e)))?;

        Ok(y)
    }

    /// Create an optimizer for this model
    pub fn create_optimizer(&self) -> CoreResult<AdamW> {
        let params = ParamsAdamW {
            lr: self.training_config.learning_rate,
            beta1: self.training_config.beta1,
            beta2: self.training_config.beta2,
            eps: self.training_config.eps,
            weight_decay: self.training_config.weight_decay,
        };

        AdamW::new(self.varmap.all_vars(), params)
            .map_err(|e| CoreError::Generic(format!("Failed to create optimizer: {}", e)))
    }

    /// Get the variable map for loading/saving weights
    pub fn varmap(&self) -> &VarMap {
        &self.varmap
    }

    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }

    /// Get dtype
    pub fn dtype(&self) -> DType {
        self.dtype
    }

    /// Save model weights to a safetensors file
    ///
    /// # Arguments
    /// * `path` - Path to save the safetensors file
    ///
    /// # Example
    /// ```rust,ignore
    /// model.save_weights("model.safetensors")?;
    /// ```
    pub fn save_weights<P: AsRef<std::path::Path>>(&self, path: P) -> CoreResult<()> {
        self.varmap
            .save(path)
            .map_err(|e| CoreError::Generic(format!("Failed to save weights: {}", e)))
    }

    /// Load model weights from a safetensors file
    ///
    /// # Arguments
    /// * `path` - Path to the safetensors file
    ///
    /// # Example
    /// ```rust,ignore
    /// model.load_weights("model.safetensors")?;
    /// ```
    pub fn load_weights<P: AsRef<std::path::Path>>(&mut self, path: P) -> CoreResult<()> {
        self.varmap
            .load(path)
            .map_err(|e| CoreError::Generic(format!("Failed to load weights: {}", e)))
    }
}

/// Constraint-aware loss wrapper
///
/// Bridges kizzasi-logic constraints with candle tensor operations.
/// Allows combining task loss with constraint violations for constrained optimization.
///
/// # Examples
///
/// ```rust,ignore
/// use kizzasi_core::{ConstraintLoss, Loss};
///
/// let constraint_loss = ConstraintLoss::new(0.1);
///
/// // In training loop:
/// let task_loss = Loss::mse(&predictions, &targets)?;
/// let total_loss = constraint_loss.compute(&task_loss, &predictions, |pred| {
///     // Compute constraint violation from prediction
///     Ok(0.0)
/// })?;
/// ```
pub struct ConstraintLoss {
    /// Base weight for constraint violations
    constraint_weight: f32,
}

impl ConstraintLoss {
    /// Create a new constraint-aware loss
    pub fn new(constraint_weight: f32) -> Self {
        Self { constraint_weight }
    }

    /// Compute combined loss: task_loss + constraint_weight * constraint_penalty
    ///
    /// # Arguments
    /// * `task_loss` - Base task loss (MSE, MAE, etc.)
    /// * `prediction` - Model prediction tensor
    /// * `constraint_fn` - Function that computes constraint violation from prediction
    pub fn compute<F>(
        &self,
        task_loss: &Tensor,
        prediction: &Tensor,
        constraint_fn: F,
    ) -> CoreResult<Tensor>
    where
        F: Fn(&Tensor) -> CoreResult<f32>,
    {
        // Compute constraint violation
        let violation = constraint_fn(prediction)?;

        // Add constraint penalty to task loss
        // Create a scalar penalty value matching task_loss shape
        let penalty_value = self.constraint_weight * violation;

        // Use affine to add the penalty: task_loss + penalty = task_loss * 1.0 + penalty
        task_loss
            .affine(1.0, penalty_value as f64)
            .map_err(|e| CoreError::Generic(format!("Failed to add constraint penalty: {}", e)))
    }
}

/// Loss functions for training
pub struct Loss;

impl Loss {
    /// Mean Squared Error loss
    pub fn mse(predictions: &Tensor, targets: &Tensor) -> CoreResult<Tensor> {
        predictions
            .sub(targets)
            .map_err(|e| CoreError::Generic(format!("MSE subtraction failed: {}", e)))?
            .sqr()
            .map_err(|e| CoreError::Generic(format!("MSE square failed: {}", e)))?
            .mean_all()
            .map_err(|e| CoreError::Generic(format!("MSE mean failed: {}", e)))
    }

    /// Mean Absolute Error loss
    pub fn mae(predictions: &Tensor, targets: &Tensor) -> CoreResult<Tensor> {
        predictions
            .sub(targets)
            .map_err(|e| CoreError::Generic(format!("MAE subtraction failed: {}", e)))?
            .abs()
            .map_err(|e| CoreError::Generic(format!("MAE abs failed: {}", e)))?
            .mean_all()
            .map_err(|e| CoreError::Generic(format!("MAE mean failed: {}", e)))
    }

    /// Huber loss (smooth L1 loss)
    pub fn huber(predictions: &Tensor, targets: &Tensor, delta: f64) -> CoreResult<Tensor> {
        let diff = predictions
            .sub(targets)
            .map_err(|e| CoreError::Generic(format!("Huber subtraction failed: {}", e)))?;
        let abs_diff = diff
            .abs()
            .map_err(|e| CoreError::Generic(format!("Huber abs failed: {}", e)))?;

        // If |diff| <= delta: 0.5 * diff^2
        // If |diff| > delta: delta * (|diff| - 0.5 * delta)
        let squared = diff
            .sqr()
            .map_err(|e| CoreError::Generic(format!("Huber square failed: {}", e)))?
            .affine(0.5, 0.0)
            .map_err(|e| CoreError::Generic(format!("Huber mul 0.5 failed: {}", e)))?;

        let linear_offset = delta * delta * 0.5;
        let linear = abs_diff
            .affine(delta, -linear_offset)
            .map_err(|e| CoreError::Generic(format!("Huber linear computation failed: {}", e)))?;

        let mask = abs_diff
            .le(delta)
            .map_err(|e| CoreError::Generic(format!("Huber comparison failed: {}", e)))?
            .to_dtype(predictions.dtype())
            .map_err(|e| CoreError::Generic(format!("Huber mask conversion failed: {}", e)))?;

        // Invert mask: 1 - mask
        let inv_mask = mask
            .affine(-1.0, 1.0)
            .map_err(|e| CoreError::Generic(format!("Huber mask inversion failed: {}", e)))?;

        let loss = squared
            .mul(&mask)
            .map_err(|e| CoreError::Generic(format!("Huber squared mul failed: {}", e)))?
            .add(
                &linear
                    .mul(&inv_mask)
                    .map_err(|e| CoreError::Generic(format!("Huber linear mul failed: {}", e)))?,
            )
            .map_err(|e| CoreError::Generic(format!("Huber final add failed: {}", e)))?;

        loss.mean_all()
            .map_err(|e| CoreError::Generic(format!("Huber mean failed: {}", e)))
    }

    /// Cross-entropy loss for classification
    pub fn cross_entropy(logits: &Tensor, targets: &Tensor) -> CoreResult<Tensor> {
        // Log softmax
        let log_probs = candle_nn::ops::log_softmax(logits, candle_core::D::Minus1)
            .map_err(|e| CoreError::Generic(format!("Log softmax failed: {}", e)))?;

        // Negative log likelihood
        let nll = log_probs
            .mul(targets)
            .map_err(|e| CoreError::Generic(format!("NLL multiplication failed: {}", e)))?
            .sum_all()
            .map_err(|e| CoreError::Generic(format!("NLL sum failed: {}", e)))?
            .neg()
            .map_err(|e| CoreError::Generic(format!("NLL negation failed: {}", e)))?;

        // Average over batch
        let batch_size = logits
            .dim(0)
            .map_err(|e| CoreError::Generic(format!("Failed to get batch size: {}", e)))?;
        nll.affine(1.0 / batch_size as f64, 0.0)
            .map_err(|e| CoreError::Generic(format!("Cross entropy division failed: {}", e)))
    }
}

/// Training utilities with scheduler, metrics, and validation
pub struct Trainer {
    model: TrainableSSM,
    optimizer: AdamW,
    config: TrainingConfig,
    scheduler: Option<Box<dyn LRScheduler>>,
    metrics: TrainingMetrics,
    logger: MetricsLogger,
    current_step: usize,
}

impl Trainer {
    /// Create a new trainer
    pub fn new(model: TrainableSSM, config: TrainingConfig) -> CoreResult<Self> {
        let optimizer = model.create_optimizer()?;

        // Create scheduler based on config
        let scheduler = Self::create_scheduler(&config);

        let metrics = TrainingMetrics::new();

        let logger = MetricsLogger::new()
            .with_verbose(config.track_metrics)
            .with_log_interval(config.log_interval);

        Ok(Self {
            model,
            optimizer,
            config,
            scheduler,
            metrics,
            logger,
            current_step: 0,
        })
    }

    /// Create scheduler from config
    fn create_scheduler(config: &TrainingConfig) -> Option<Box<dyn LRScheduler>> {
        use crate::scheduler::*;

        config.scheduler.as_ref().map(|sched_type| {
            let total_steps = config.epochs * 100; // Rough estimate, can be updated later

            match sched_type {
                SchedulerType::Constant => {
                    Box::new(ConstantScheduler::new(config.learning_rate)) as Box<dyn LRScheduler>
                }
                SchedulerType::Linear {
                    warmup_steps,
                    final_lr,
                } => Box::new(LinearScheduler::new(
                    config.learning_rate,
                    *final_lr,
                    total_steps,
                    *warmup_steps,
                )) as Box<dyn LRScheduler>,
                SchedulerType::Cosine {
                    warmup_steps,
                    min_lr,
                } => Box::new(
                    CosineScheduler::new(config.learning_rate, total_steps, *warmup_steps)
                        .with_min_lr(*min_lr),
                ) as Box<dyn LRScheduler>,
                SchedulerType::Step {
                    milestones,
                    decay_factor,
                } => Box::new(StepScheduler::new(
                    config.learning_rate,
                    *decay_factor,
                    milestones.clone(),
                )) as Box<dyn LRScheduler>,
                SchedulerType::Exponential {
                    decay_rate,
                    decay_steps,
                } => Box::new(ExponentialScheduler::new(
                    config.learning_rate,
                    *decay_rate,
                    *decay_steps,
                )) as Box<dyn LRScheduler>,
                SchedulerType::OneCycle { warmup_pct } => Box::new(
                    OneCycleScheduler::new(config.learning_rate, total_steps)
                        .with_warmup_pct(*warmup_pct),
                ) as Box<dyn LRScheduler>,
                SchedulerType::Polynomial { final_lr, power } => Box::new(PolynomialScheduler::new(
                    config.learning_rate,
                    *final_lr,
                    total_steps,
                    *power,
                ))
                    as Box<dyn LRScheduler>,
            }
        })
    }

    /// Get current learning rate
    fn get_current_lr(&self) -> f64 {
        self.scheduler
            .as_ref()
            .map(|s| s.get_lr(self.current_step))
            .unwrap_or(self.config.learning_rate)
    }

    /// Train for one epoch
    pub fn train_epoch<F>(
        &mut self,
        data_loader: &[(Tensor, Tensor)],
        loss_fn: F,
    ) -> CoreResult<f32>
    where
        F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor>,
    {
        let mut total_loss = 0.0;
        let num_batches = data_loader.len();
        let epoch = self.current_step / num_batches.max(1);

        for (batch_idx, (inputs, targets)) in data_loader.iter().enumerate() {
            // Update learning rate from scheduler
            let lr = self.get_current_lr();
            if self.config.track_metrics {
                self.metrics.record_learning_rate(lr);
            }

            // Forward pass
            let predictions = self.model.forward(inputs)?;

            // Compute loss
            let loss = loss_fn(&predictions, targets)?;

            // Backward pass
            self.optimizer
                .backward_step(&loss)
                .map_err(|e| CoreError::Generic(format!("Backward step failed: {}", e)))?;

            // Accumulate loss
            let loss_val = loss
                .to_vec0::<f32>()
                .map_err(|e| CoreError::Generic(format!("Failed to extract loss value: {}", e)))?;
            total_loss += loss_val;

            // Track metrics
            if self.config.track_metrics {
                self.metrics.record_train_loss(epoch, loss_val);
                self.logger.log_batch(epoch, batch_idx, loss_val);

                // Compute and track gradient norm
                let grad_norm = self.compute_grad_norm()?;
                self.metrics.record_grad_norm(grad_norm);
            }

            // Gradient clipping if enabled
            if let Some(max_norm) = self.config.grad_clip {
                self.clip_gradients(max_norm)?;
            }

            self.current_step += 1;
        }

        Ok(total_loss / num_batches as f32)
    }

    /// Compute gradient norm
    fn compute_grad_norm(&self) -> CoreResult<f32> {
        // Placeholder: In candle, gradient norms would be computed from VarMap
        // For now, return a dummy value
        // TODO: Implement proper gradient norm computation when candle exposes gradient access
        Ok(1.0)
    }

    /// Clip gradients by global norm
    ///
    /// Note: Gradient clipping is handled internally by candle's optimizer.
    /// This is a placeholder for custom gradient clipping if needed.
    fn clip_gradients(&self, _max_norm: f32) -> CoreResult<()> {
        // Gradient clipping will be handled by the optimizer's built-in mechanism
        // or via custom backward hooks in future implementations
        Ok(())
    }

    /// Evaluate on validation data
    pub fn evaluate<F>(&self, data_loader: &[(Tensor, Tensor)], loss_fn: F) -> CoreResult<f32>
    where
        F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor>,
    {
        let mut total_loss = 0.0;
        let num_batches = data_loader.len();

        for (inputs, targets) in data_loader {
            // Forward pass (no gradient tracking needed)
            let predictions = self.model.forward(inputs)?;

            // Compute loss
            let loss = loss_fn(&predictions, targets)?;

            // Accumulate loss
            let loss_val = loss
                .to_vec0::<f32>()
                .map_err(|e| CoreError::Generic(format!("Failed to extract loss value: {}", e)))?;
            total_loss += loss_val;
        }

        Ok(total_loss / num_batches as f32)
    }

    /// Full training loop with validation and early stopping
    pub fn fit<F>(
        &mut self,
        mut train_loader: TimeSeriesDataLoader,
        mut val_loader: Option<TimeSeriesDataLoader>,
        loss_fn: F,
    ) -> CoreResult<()>
    where
        F: Fn(&Tensor, &Tensor) -> CoreResult<Tensor> + Copy,
    {
        use std::time::Instant;

        for epoch in 0..self.config.epochs {
            let epoch_start = Instant::now();

            // Shuffle training data
            train_loader.shuffle();

            // Prepare batches (simplified - actual implementation would iterate batches)
            // For now, this is a placeholder for the integration
            // TODO: Implement proper batch iteration with TimeSeriesDataLoader
            let train_batches: Vec<(Tensor, Tensor)> = Vec::new();

            // Train for one epoch
            let train_loss = self.train_epoch(&train_batches, loss_fn)?;

            // Validation
            let val_loss = if let Some(ref mut _val_data) = val_loader {
                let val_batches: Vec<(Tensor, Tensor)> = Vec::new();
                let val_loss = self.evaluate(&val_batches, loss_fn)?;

                if self.config.track_metrics {
                    self.metrics.record_val_loss(epoch, val_loss);
                }

                Some(val_loss)
            } else {
                None
            };

            // Track epoch duration
            let epoch_duration = epoch_start.elapsed().as_secs_f64();
            if self.config.track_metrics {
                self.metrics.record_epoch_duration(epoch, epoch_duration);
            }

            // Log epoch metrics
            let current_lr = self.get_current_lr();
            self.logger
                .log_epoch(epoch, train_loss, val_loss, current_lr);

            // Early stopping check
            if let Some(patience) = self.config.early_stopping_patience {
                if !self.metrics.is_improving(patience) {
                    tracing::info!("Early stopping triggered at epoch {}", epoch);
                    break;
                }
            }
        }

        // Log training summary
        if self.config.track_metrics {
            let summary = self.metrics.summary();
            self.logger.log_summary(&summary);
        }

        Ok(())
    }

    /// Get reference to the model
    pub fn model(&self) -> &TrainableSSM {
        &self.model
    }

    /// Get mutable reference to the model
    pub fn model_mut(&mut self) -> &mut TrainableSSM {
        &mut self.model
    }

    /// Get reference to training metrics
    pub fn metrics(&self) -> &TrainingMetrics {
        &self.metrics
    }

    /// Get mutable reference to training metrics
    pub fn metrics_mut(&mut self) -> &mut TrainingMetrics {
        &mut self.metrics
    }

    /// Get current training step
    pub fn current_step(&self) -> usize {
        self.current_step
    }

    /// Save checkpoint to disk
    ///
    /// Saves model weights, optimizer state, training configuration, metrics, and metadata.
    ///
    /// # Arguments
    /// * `path` - Directory to save checkpoint files
    /// * `name` - Checkpoint name (without extension)
    ///
    /// # Example
    /// ```rust,ignore
    /// trainer.save_checkpoint("checkpoints", "epoch_10")?;
    /// // Creates: checkpoints/epoch_10.safetensors and checkpoints/epoch_10.json
    /// ```
    pub fn save_checkpoint<P: AsRef<std::path::Path>>(
        &self,
        path: P,
        name: &str,
    ) -> CoreResult<()> {
        use std::fs;
        use std::path::PathBuf;

        let checkpoint_dir = path.as_ref();
        fs::create_dir_all(checkpoint_dir).map_err(|e| {
            CoreError::Generic(format!("Failed to create checkpoint directory: {}", e))
        })?;

        // Save model weights to safetensors
        let weights_path: PathBuf = checkpoint_dir.join(format!("{}.safetensors", name));
        self.model
            .save_weights(&weights_path)
            .map_err(|e| CoreError::Generic(format!("Failed to save model weights: {}", e)))?;

        // Create checkpoint metadata
        let metadata = CheckpointMetadata {
            version: env!("CARGO_PKG_VERSION").to_string(),
            timestamp: chrono::Utc::now().to_rfc3339(),
            current_step: self.current_step,
            current_epoch: self.metrics.summary().total_epochs,
            config: self.config.clone(),
            metrics: self.metrics.clone(),
        };

        // Save metadata to JSON
        let metadata_path: PathBuf = checkpoint_dir.join(format!("{}.json", name));
        let metadata_json = serde_json::to_string_pretty(&metadata).map_err(|e| {
            CoreError::Generic(format!("Failed to serialize checkpoint metadata: {}", e))
        })?;

        fs::write(&metadata_path, metadata_json).map_err(|e| {
            CoreError::Generic(format!("Failed to write checkpoint metadata: {}", e))
        })?;

        tracing::info!(
            "Checkpoint saved: weights={}, metadata={}",
            weights_path.display(),
            metadata_path.display()
        );

        Ok(())
    }

    /// Load checkpoint and resume training
    ///
    /// Creates a new Trainer from a saved checkpoint, restoring model weights,
    /// configuration, and training state.
    ///
    /// # Arguments
    /// * `path` - Directory containing checkpoint files
    /// * `name` - Checkpoint name (without extension)
    /// * `model_config` - Model configuration (must match saved model)
    ///
    /// # Example
    /// ```rust,ignore
    /// let trainer = Trainer::load_checkpoint("checkpoints", "epoch_10", model_config)?;
    /// // Continue training from epoch 10
    /// ```
    pub fn load_checkpoint<P: AsRef<std::path::Path>>(
        path: P,
        name: &str,
        model_config: KizzasiConfig,
    ) -> CoreResult<Self> {
        use std::fs;
        use std::path::PathBuf;

        let checkpoint_dir = path.as_ref();

        // Load metadata from JSON
        let metadata_path: PathBuf = checkpoint_dir.join(format!("{}.json", name));
        let metadata_json = fs::read_to_string(&metadata_path).map_err(|e| {
            CoreError::Generic(format!("Failed to read checkpoint metadata: {}", e))
        })?;

        let metadata: CheckpointMetadata = serde_json::from_str(&metadata_json).map_err(|e| {
            CoreError::Generic(format!("Failed to parse checkpoint metadata: {}", e))
        })?;

        // Load model weights
        let weights_path: PathBuf = checkpoint_dir.join(format!("{}.safetensors", name));
        let mut model = TrainableSSM::new(model_config, metadata.config.clone())?;
        model
            .load_weights(&weights_path)
            .map_err(|e| CoreError::Generic(format!("Failed to load model weights: {}", e)))?;

        // Create trainer with loaded state
        let optimizer = model.create_optimizer()?;
        let scheduler = Self::create_scheduler(&metadata.config);

        let logger = MetricsLogger::new()
            .with_verbose(metadata.config.track_metrics)
            .with_log_interval(metadata.config.log_interval);

        tracing::info!(
            "Checkpoint loaded: version={}, step={}, epoch={}",
            metadata.version,
            metadata.current_step,
            metadata.current_epoch
        );

        Ok(Self {
            model,
            optimizer,
            config: metadata.config,
            scheduler,
            metrics: metadata.metrics,
            logger,
            current_step: metadata.current_step,
        })
    }

    /// Save checkpoint with automatic naming (epoch-based)
    ///
    /// Convenience method that automatically names checkpoints based on current epoch.
    ///
    /// # Example
    /// ```rust,ignore
    /// trainer.save_checkpoint_auto("checkpoints")?;
    /// // Creates: checkpoints/checkpoint_epoch_5.safetensors, etc.
    /// ```
    pub fn save_checkpoint_auto<P: AsRef<std::path::Path>>(&self, path: P) -> CoreResult<()> {
        let current_epoch = self.metrics.summary().total_epochs;
        let name = format!("checkpoint_epoch_{}", current_epoch);
        self.save_checkpoint(path, &name)
    }

    /// Save checkpoint if this is the best epoch (lowest validation loss)
    ///
    /// Automatically saves a "best" checkpoint when validation loss improves.
    ///
    /// # Example
    /// ```rust,ignore
    /// // After each validation epoch
    /// trainer.save_best_checkpoint("checkpoints")?;
    /// ```
    pub fn save_best_checkpoint<P: AsRef<std::path::Path>>(&self, path: P) -> CoreResult<()> {
        let summary = self.metrics.summary();

        // Only save if this is the best epoch
        // Note: total_epochs is 1-indexed (count), best_epoch is 0-indexed (epoch number)
        if let (Some(best_epoch), Some(_best_loss)) = (summary.best_epoch, summary.best_val_loss) {
            // Current epoch is total_epochs - 1 (convert from count to 0-indexed)
            let current_epoch = summary.total_epochs.saturating_sub(1);
            if current_epoch == best_epoch {
                tracing::info!("New best validation loss! Saving best checkpoint");
                return self.save_checkpoint(path, "best");
            }
        }

        Ok(())
    }
}

/// Checkpoint metadata for training state persistence
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMetadata {
    /// Package version when checkpoint was created
    pub version: String,
    /// ISO 8601 timestamp
    pub timestamp: String,
    /// Current training step
    pub current_step: usize,
    /// Current epoch number
    pub current_epoch: usize,
    /// Training configuration
    pub config: TrainingConfig,
    /// Training metrics history
    pub metrics: TrainingMetrics,
}

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

    #[test]
    fn test_trainable_ssm_creation() {
        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();

        let model = TrainableSSM::new(config, training_config);
        assert!(model.is_ok());
    }

    #[test]
    fn test_forward_pass() {
        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();

        let model = TrainableSSM::new(config, training_config).unwrap();
        let device = model.device().clone();

        // Create dummy input: [batch=2, seq=10, input_dim=3]
        let input = Tensor::randn(0f32, 1.0, (2, 10, 3), &device).unwrap();

        let output = model.forward(&input);
        if let Err(e) = &output {
            panic!("Forward pass failed: {:?}", e);
        }

        let output = output.unwrap();
        assert_eq!(output.dims(), &[2, 10, 3]);
    }

    #[test]
    fn test_mse_loss() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 3.0], &device).unwrap();
        let targets = Tensor::new(&[1.5f32, 2.5, 3.5], &device).unwrap();

        let loss = Loss::mse(&predictions, &targets).unwrap();
        let loss_val = loss.to_vec0::<f32>().unwrap();

        // Expected: mean((0.5^2 + 0.5^2 + 0.5^2)) = 0.25
        assert!((loss_val - 0.25).abs() < 1e-5);
    }

    #[test]
    fn test_training_config_default() {
        let config = TrainingConfig::default();
        assert_eq!(config.learning_rate, 1e-4);
        assert_eq!(config.batch_size, 32);
        assert_eq!(config.epochs, 10);
        assert!(config.track_metrics);
        assert_eq!(config.validation_split, 0.2);
        assert_eq!(config.early_stopping_patience, Some(5));
    }

    #[test]
    fn test_training_config_with_scheduler() {
        let config = TrainingConfig::default().with_scheduler(SchedulerType::Cosine {
            warmup_steps: 100,
            min_lr: 1e-6,
        });

        assert!(config.scheduler.is_some());
        if let Some(SchedulerType::Cosine {
            warmup_steps,
            min_lr,
        }) = config.scheduler
        {
            assert_eq!(warmup_steps, 100);
            assert_eq!(min_lr, 1e-6);
        } else {
            panic!("Expected Cosine scheduler");
        }
    }

    #[test]
    fn test_training_config_builder() {
        let config = TrainingConfig::default()
            .with_validation_split(0.15)
            .with_early_stopping(10)
            .without_metrics();

        assert_eq!(config.validation_split, 0.15);
        assert_eq!(config.early_stopping_patience, Some(10));
        assert!(!config.track_metrics);
    }

    #[test]
    fn test_trainer_with_scheduler() {
        let model_config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default().with_scheduler(SchedulerType::Linear {
            warmup_steps: 50,
            final_lr: 1e-6,
        });

        let model = TrainableSSM::new(model_config, training_config.clone()).unwrap();
        let trainer = Trainer::new(model, training_config);

        assert!(trainer.is_ok());
        let trainer = trainer.unwrap();
        assert!(trainer.scheduler.is_some());
    }

    #[test]
    fn test_trainer_metrics_tracking() {
        let model_config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();
        let model = TrainableSSM::new(model_config, training_config.clone()).unwrap();
        let trainer = Trainer::new(model, training_config).unwrap();

        // Check that metrics are initialized
        assert_eq!(trainer.metrics().current_step(), 0);
        assert_eq!(trainer.current_step(), 0);
    }

    #[test]
    fn test_scheduler_type_constant() {
        let config = TrainingConfig::default().with_scheduler(SchedulerType::Constant);

        assert!(config.scheduler.is_some());
    }

    #[test]
    fn test_scheduler_type_step() {
        let config = TrainingConfig::default().with_scheduler(SchedulerType::Step {
            milestones: vec![100, 200, 300],
            decay_factor: 0.1,
        });

        if let Some(SchedulerType::Step {
            milestones,
            decay_factor,
        }) = config.scheduler
        {
            assert_eq!(milestones, vec![100, 200, 300]);
            assert_eq!(decay_factor, 0.1);
        } else {
            panic!("Expected Step scheduler");
        }
    }

    #[test]
    fn test_scheduler_type_onecycle() {
        let config =
            TrainingConfig::default().with_scheduler(SchedulerType::OneCycle { warmup_pct: 0.3 });

        if let Some(SchedulerType::OneCycle { warmup_pct }) = config.scheduler {
            assert_eq!(warmup_pct, 0.3);
        } else {
            panic!("Expected OneCycle scheduler");
        }
    }

    #[test]
    fn test_mae_loss() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 3.0], &device).unwrap();
        let targets = Tensor::new(&[1.5f32, 2.5, 3.5], &device).unwrap();

        let loss = Loss::mae(&predictions, &targets).unwrap();
        let loss_val = loss.to_vec0::<f32>().unwrap();

        // Expected: mean(|0.5| + |0.5| + |0.5|) = 0.5
        assert!((loss_val - 0.5).abs() < 1e-5);
    }

    #[test]
    fn test_huber_loss() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 5.0], &device).unwrap();
        let targets = Tensor::new(&[1.1f32, 2.1, 3.0], &device).unwrap();

        let loss = Loss::huber(&predictions, &targets, 1.0).unwrap();
        let loss_val = loss.to_vec0::<f32>().unwrap();

        // Huber loss is smooth L1
        assert!(loss_val > 0.0);
        assert!(loss_val < 2.0); // Should be less than L1 loss for large errors
    }

    #[test]
    fn test_constraint_loss_creation() {
        let constraint_loss = ConstraintLoss::new(0.5);
        assert_eq!(constraint_loss.constraint_weight, 0.5);
    }

    #[test]
    fn test_constraint_loss_no_violation() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 3.0], &device).unwrap();
        let targets = Tensor::new(&[1.5f32, 2.5, 3.5], &device).unwrap();

        let task_loss = Loss::mse(&predictions, &targets).unwrap();
        let task_loss_val = task_loss.to_vec0::<f32>().unwrap();

        let constraint_loss = ConstraintLoss::new(0.5);

        // No constraint violation
        let total_loss = constraint_loss
            .compute(&task_loss, &predictions, |_pred| Ok(0.0))
            .unwrap();
        let total_loss_val = total_loss.to_vec0::<f32>().unwrap();

        // Should equal task loss when no violation
        assert!((total_loss_val - task_loss_val).abs() < 1e-5);
    }

    #[test]
    fn test_constraint_loss_with_violation() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 3.0], &device).unwrap();
        let targets = Tensor::new(&[1.5f32, 2.5, 3.5], &device).unwrap();

        let task_loss = Loss::mse(&predictions, &targets).unwrap();
        let task_loss_val = task_loss.to_vec0::<f32>().unwrap();

        let constraint_loss = ConstraintLoss::new(0.5);

        // Constraint violation of 1.0
        let total_loss = constraint_loss
            .compute(&task_loss, &predictions, |_pred| Ok(1.0))
            .unwrap();
        let total_loss_val = total_loss.to_vec0::<f32>().unwrap();

        // Should be task_loss + 0.5 * 1.0 = task_loss + 0.5
        let expected = task_loss_val + 0.5;
        assert!((total_loss_val - expected).abs() < 1e-5);
    }

    #[test]
    fn test_constraint_loss_scaling() {
        let device = Device::Cpu;
        let predictions = Tensor::new(&[1.0f32, 2.0, 3.0], &device).unwrap();
        let targets = Tensor::new(&[1.5f32, 2.5, 3.5], &device).unwrap();

        let task_loss = Loss::mse(&predictions, &targets).unwrap();
        let task_loss_val = task_loss.to_vec0::<f32>().unwrap();

        // Test different constraint weights
        let weights = [0.1, 0.5, 1.0, 2.0];
        let violation = 1.5;

        for &weight in &weights {
            let constraint_loss = ConstraintLoss::new(weight);
            let total_loss = constraint_loss
                .compute(&task_loss, &predictions, |_pred| Ok(violation))
                .unwrap();
            let total_loss_val = total_loss.to_vec0::<f32>().unwrap();

            let expected = task_loss_val + weight * violation;
            assert!(
                (total_loss_val - expected).abs() < 1e-4,
                "Weight {} failed: got {}, expected {}",
                weight,
                total_loss_val,
                expected
            );
        }
    }

    #[test]
    fn test_checkpoint_save_load() {
        use std::env;
        use std::fs;

        let temp_dir = env::temp_dir().join("kizzasi_checkpoint_test");
        fs::create_dir_all(&temp_dir).unwrap();

        // Create a model
        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig {
            epochs: 5,
            learning_rate: 1e-3,
            ..Default::default()
        };

        let model = TrainableSSM::new(config.clone(), training_config.clone()).unwrap();
        let trainer = Trainer::new(model, training_config).unwrap();

        // Save checkpoint
        trainer
            .save_checkpoint(&temp_dir, "test_checkpoint")
            .unwrap();

        // Verify files exist
        assert!(temp_dir.join("test_checkpoint.safetensors").exists());
        assert!(temp_dir.join("test_checkpoint.json").exists());

        // Load checkpoint
        let loaded_trainer =
            Trainer::load_checkpoint(&temp_dir, "test_checkpoint", config).unwrap();

        // Verify loaded config matches
        assert_eq!(loaded_trainer.config.epochs, 5);
        assert_eq!(loaded_trainer.config.learning_rate, 1e-3);
        assert_eq!(loaded_trainer.current_step, 0);

        // Clean up
        fs::remove_dir_all(&temp_dir).unwrap();
    }

    #[test]
    fn test_checkpoint_auto_save() {
        use std::env;
        use std::fs;

        let temp_dir = env::temp_dir().join("kizzasi_checkpoint_auto_test");
        fs::create_dir_all(&temp_dir).unwrap();

        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();
        let model = TrainableSSM::new(config, training_config.clone()).unwrap();
        let mut trainer = Trainer::new(model, training_config).unwrap();

        // Record some metrics to simulate training
        trainer.metrics.record_train_loss(0, 0.5);

        // Save checkpoint with auto naming
        trainer.save_checkpoint_auto(&temp_dir).unwrap();

        // Verify file exists with auto-generated name
        assert!(temp_dir.join("checkpoint_epoch_1.safetensors").exists());
        assert!(temp_dir.join("checkpoint_epoch_1.json").exists());

        // Clean up
        fs::remove_dir_all(&temp_dir).unwrap();
    }

    #[test]
    fn test_checkpoint_best_save() {
        use std::env;
        use std::fs;

        let temp_dir = env::temp_dir().join("kizzasi_checkpoint_best_test");
        fs::create_dir_all(&temp_dir).unwrap();

        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();
        let model = TrainableSSM::new(config, training_config.clone()).unwrap();
        let mut trainer = Trainer::new(model, training_config).unwrap();

        // Simulate training epoch 0 (not best yet)
        trainer.metrics.record_train_loss(0, 1.2);
        trainer.metrics.record_val_loss(0, 1.0);
        trainer.save_best_checkpoint(&temp_dir).unwrap();

        // Epoch 0 is the best so far, so checkpoint should be saved
        assert!(temp_dir.join("best.safetensors").exists());
        assert!(temp_dir.join("best.json").exists());

        // Simulate training epoch 1 with worse loss (should not overwrite)
        trainer.metrics.record_train_loss(1, 0.9);
        trainer.metrics.record_val_loss(1, 1.2);

        // Remove old best to test that it doesn't get overwritten
        fs::remove_file(temp_dir.join("best.safetensors")).unwrap();
        fs::remove_file(temp_dir.join("best.json")).unwrap();

        trainer.save_best_checkpoint(&temp_dir).unwrap();
        // Should not save because epoch 1 is not the best
        assert!(!temp_dir.join("best.safetensors").exists());

        // Clean up
        fs::remove_dir_all(&temp_dir).unwrap();
    }

    #[test]
    fn test_checkpoint_metadata() {
        use std::env;
        use std::fs;

        let temp_dir = env::temp_dir().join("kizzasi_checkpoint_metadata_test");
        fs::create_dir_all(&temp_dir).unwrap();

        let config = KizzasiConfig::new()
            .input_dim(3)
            .output_dim(3)
            .hidden_dim(64)
            .state_dim(8)
            .num_layers(2);

        let training_config = TrainingConfig::default();
        let model = TrainableSSM::new(config, training_config.clone()).unwrap();
        let mut trainer = Trainer::new(model, training_config).unwrap();

        // Add some metrics
        trainer.metrics.record_train_loss(0, 0.5);
        trainer.metrics.record_val_loss(0, 0.45);

        // Save checkpoint
        trainer.save_checkpoint(&temp_dir, "metadata_test").unwrap();

        // Load and verify metadata
        let metadata_path = temp_dir.join("metadata_test.json");
        let metadata_json = fs::read_to_string(&metadata_path).unwrap();
        let metadata: CheckpointMetadata = serde_json::from_str(&metadata_json).unwrap();

        assert_eq!(metadata.version, env!("CARGO_PKG_VERSION"));
        assert!(!metadata.timestamp.is_empty());
        assert_eq!(metadata.current_step, 0);
        assert!(metadata.metrics.val_loss(0).is_some());
        assert_eq!(metadata.metrics.val_loss(0).unwrap(), 0.45);

        // Clean up
        fs::remove_dir_all(&temp_dir).unwrap();
    }
}