sklears-neural 0.2.0

//! Optimization solvers for neural networks.

use crate::NeuralResult;
use scirs2_core::ndarray::{Array1, Array2};
use std::collections::VecDeque;

/// Gradient clipping methods
#[derive(Debug, Clone, Copy, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum GradientClipping {
    /// No gradient clipping
    None,
    /// Clip gradients by norm (global norm clipping)
    ByNorm(f64),
    /// Clip gradients by value (element-wise clipping)
    ByValue(f64),
}

/// Gradient clipping utilities
pub struct GradientClipper;

impl GradientClipper {
    /// Apply gradient clipping to weights and biases gradients
    pub fn clip_gradients(
        weight_grads: &mut [Array2<f64>],
        bias_grads: &mut [Array1<f64>],
        clipping: GradientClipping,
    ) {
        match clipping {
            GradientClipping::None => {}
            GradientClipping::ByNorm(max_norm) => {
                Self::clip_by_global_norm(weight_grads, bias_grads, max_norm);
            }
            GradientClipping::ByValue(max_value) => {
                Self::clip_by_value(weight_grads, bias_grads, max_value);
            }
        }
    }

    /// Clip gradients by global norm
    fn clip_by_global_norm(
        weight_grads: &mut [Array2<f64>],
        bias_grads: &mut [Array1<f64>],
        max_norm: f64,
    ) {
        // Compute global norm
        let mut global_norm_sq = 0.0;

        for grad in weight_grads.iter() {
            global_norm_sq += grad.iter().map(|x| x * x).sum::<f64>();
        }

        for grad in bias_grads.iter() {
            global_norm_sq += grad.iter().map(|x| x * x).sum::<f64>();
        }

        let global_norm = global_norm_sq.sqrt();

        // Apply clipping if necessary
        if global_norm > max_norm {
            let clip_coeff = max_norm / global_norm;

            for grad in weight_grads.iter_mut() {
                grad.mapv_inplace(|x| x * clip_coeff);
            }

            for grad in bias_grads.iter_mut() {
                grad.mapv_inplace(|x| x * clip_coeff);
            }
        }
    }

    /// Clip gradients by value (element-wise)
    fn clip_by_value(
        weight_grads: &mut [Array2<f64>],
        bias_grads: &mut [Array1<f64>],
        max_value: f64,
    ) {
        for grad in weight_grads.iter_mut() {
            grad.mapv_inplace(|x| x.max(-max_value).min(max_value));
        }

        for grad in bias_grads.iter_mut() {
            grad.mapv_inplace(|x| x.max(-max_value).min(max_value));
        }
    }
}
/// Solver algorithms for weight optimization
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Solver {
    /// Limited-memory BFGS algorithm (two-loop recursion, [`LbfgsSolver`])
    Lbfgs,
    /// Stochastic Gradient Descent
    Sgd,
    /// Adam optimizer
    #[default]
    Adam,
    /// AdamW optimizer (Adam with decoupled weight decay, [`AdamWSolver`])
    AdamW,
    /// RMSprop optimizer ([`RMSpropSolver`])
    RMSprop,
    /// Nadam optimizer (Nesterov-accelerated Adam, [`NadamSolver`])
    Nadam,
    /// LARS optimizer (Layer-wise Adaptive Rate Scaling)
    Lars,
    /// LAMB optimizer (Layer-wise Adaptive Moments optimizer for Batch training)
    Lamb,
}

/// Learning rate schedule for SGD
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum LearningRateSchedule {
    /// Constant learning rate
    #[default]
    Constant,
    /// Inverted scaling: lr = lr_init / (t^power)
    InvScaling,
    /// Adaptive learning rate based on loss improvement
    Adaptive,
    /// Exponential decay: lr = lr_init * decay_rate^(step / decay_steps)
    ExponentialDecay,
    /// Cosine annealing: lr = lr_min + (lr_max - lr_min) * (1 + cos(pi * step / max_steps)) / 2
    CosineAnnealing,
    /// Step decay: lr = lr_init * decay_factor^floor(step / step_size)
    StepDecay,
    /// Warmup: lr increases linearly from 0 to lr_init over warmup_steps
    Warmup,
    /// Cyclical learning rate: lr cycles between lr_min and lr_max
    CyclicalLR,
}

/// SGD solver for neural network optimization
#[derive(Debug, Clone)]
pub struct SgdSolver {
    /// Current learning rate, adjusted by the schedule each step
    pub learning_rate: f64,
    /// Momentum coefficient; 0.0 disables momentum
    pub momentum: f64,
    /// Whether to use Nesterov accelerated momentum
    pub nesterovs_momentum: bool,
    /// Learning rate schedule that determines how the rate changes over time
    pub schedule: LearningRateSchedule,
    /// Exponent for inverse-scaling schedule: `lr = eta0 / (t + 1)^power_t`
    pub power_t: f64,
    /// Initial learning rate used as the reference for schedule computations
    pub eta0: f64,
    /// Multiplicative decay factor per `decay_steps` for exponential schedule
    pub decay_rate: f64,
    /// Number of steps between learning rate decay events for exponential schedule
    pub decay_steps: usize,
    /// Total steps budget used by cosine annealing and warmup schedules
    pub max_steps: usize,
    /// Minimum learning rate floor for cosine annealing and cyclical schedules
    pub lr_min: f64,
    /// Maximum learning rate ceiling for cyclical LR schedule
    pub lr_max: f64,
    /// Half-cycle length in steps for step-decay and cyclical LR
    pub step_size: usize,
    /// Multiplicative factor applied at each step boundary for step-decay
    pub decay_factor: f64,
    /// Number of linear-warmup steps before reaching the full learning rate
    pub warmup_steps: usize,
    /// Full cycle length in steps for cyclical LR schedule
    pub cycle_length: usize,
    velocity_weights: Vec<Array2<f64>>,
    velocity_biases: Vec<Array1<f64>>,
    t: usize,
}

impl SgdSolver {
    /// Create a new SGD solver with basic parameters; advanced schedule params use sensible defaults
    pub fn new(
        learning_rate: f64,
        momentum: f64,
        nesterovs_momentum: bool,
        schedule: LearningRateSchedule,
        power_t: f64,
        eta0: f64,
    ) -> Self {
        Self {
            learning_rate,
            momentum,
            nesterovs_momentum,
            schedule,
            power_t,
            eta0,
            // Default values for additional parameters
            decay_rate: 0.9,
            decay_steps: 1000,
            max_steps: 10000,
            lr_min: 0.0,
            lr_max: learning_rate,
            step_size: 1000,
            decay_factor: 0.1,
            warmup_steps: 1000,
            cycle_length: 2000,
            velocity_weights: Vec::new(),
            velocity_biases: Vec::new(),
            t: 0,
        }
    }

    /// Create a new SGD solver with full configuration
    pub fn new_with_schedule_params(
        learning_rate: f64,
        momentum: f64,
        nesterovs_momentum: bool,
        schedule: LearningRateSchedule,
        power_t: f64,
        eta0: f64,
        decay_rate: f64,
        decay_steps: usize,
        max_steps: usize,
        lr_min: f64,
        lr_max: f64,
        step_size: usize,
        decay_factor: f64,
        warmup_steps: usize,
        cycle_length: usize,
    ) -> Self {
        Self {
            learning_rate,
            momentum,
            nesterovs_momentum,
            schedule,
            power_t,
            eta0,
            decay_rate,
            decay_steps,
            max_steps,
            lr_min,
            lr_max,
            step_size,
            decay_factor,
            warmup_steps,
            cycle_length,
            velocity_weights: Vec::new(),
            velocity_biases: Vec::new(),
            t: 0,
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.velocity_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.velocity_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.t = 0;
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        self.t += 1;
        let lr = self.get_learning_rate();

        for i in 0..weights.len() {
            // Update velocity for weights
            self.velocity_weights[i] =
                &self.velocity_weights[i] * self.momentum - &weight_grads[i] * lr;

            if self.nesterovs_momentum {
                // Nesterov momentum
                weights[i] =
                    &weights[i] + &self.velocity_weights[i] * self.momentum - &weight_grads[i] * lr;
            } else {
                // Standard momentum
                weights[i] = &weights[i] + &self.velocity_weights[i];
            }

            // Update velocity for biases
            self.velocity_biases[i] =
                &self.velocity_biases[i] * self.momentum - &bias_grads[i] * lr;

            if self.nesterovs_momentum {
                biases[i] =
                    &biases[i] + &self.velocity_biases[i] * self.momentum - &bias_grads[i] * lr;
            } else {
                biases[i] = &biases[i] + &self.velocity_biases[i];
            }
        }

        Ok(())
    }

    fn get_learning_rate(&self) -> f64 {
        match self.schedule {
            LearningRateSchedule::Constant => self.learning_rate,
            LearningRateSchedule::InvScaling => self.eta0 / (self.t as f64).powf(self.power_t),
            LearningRateSchedule::Adaptive => {
                // This would need loss history to implement properly
                // For now, return constant rate
                self.learning_rate
            }
            LearningRateSchedule::ExponentialDecay => {
                // lr = lr_init * decay_rate^(step / decay_steps)
                self.learning_rate
                    * self
                        .decay_rate
                        .powf(self.t as f64 / self.decay_steps as f64)
            }
            LearningRateSchedule::CosineAnnealing => {
                // lr = lr_min + (lr_max - lr_min) * (1 + cos(pi * step / max_steps)) / 2
                let progress = (self.t as f64 / self.max_steps as f64).min(1.0);
                self.lr_min
                    + (self.learning_rate - self.lr_min)
                        * (1.0 + (std::f64::consts::PI * progress).cos())
                        / 2.0
            }
            LearningRateSchedule::StepDecay => {
                // lr = lr_init * decay_factor^floor(step / step_size)
                let step_count = self.t / self.step_size;
                self.learning_rate * self.decay_factor.powi(step_count as i32)
            }
            LearningRateSchedule::Warmup => {
                // Linear warmup: lr increases linearly from 0 to lr_init over warmup_steps
                if self.t <= self.warmup_steps {
                    self.learning_rate * (self.t as f64 / self.warmup_steps as f64)
                } else {
                    self.learning_rate
                }
            }
            LearningRateSchedule::CyclicalLR => {
                // Triangular cyclical learning rate
                let cycle = 1.0 + (self.t as f64 / (2.0 * self.cycle_length as f64));
                let x = (2.0 * cycle.fract() - 1.0).abs();
                self.lr_min + (self.lr_max - self.lr_min) * (1.0 - x).max(0.0)
            }
        }
    }
}

/// Adam optimizer for neural network optimization
#[derive(Debug, Clone)]
pub struct AdamSolver {
    /// Step size used to scale parameter updates
    pub learning_rate: f64,
    /// Exponential decay rate for the first moment estimate
    pub beta1: f64,
    /// Exponential decay rate for the second moment estimate
    pub beta2: f64,
    /// Small constant added to denominator to prevent division by zero
    pub epsilon: f64,
    m_weights: Vec<Array2<f64>>,
    v_weights: Vec<Array2<f64>>,
    m_biases: Vec<Array1<f64>>,
    v_biases: Vec<Array1<f64>>,
    t: usize,
}

impl AdamSolver {
    /// Create a new Adam optimizer with the given hyperparameters
    pub fn new(learning_rate: f64, beta1: f64, beta2: f64, epsilon: f64) -> Self {
        Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            m_weights: Vec::new(),
            v_weights: Vec::new(),
            m_biases: Vec::new(),
            v_biases: Vec::new(),
            t: 0,
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.m_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.v_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.m_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.v_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.t = 0;
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        self.t += 1;
        let lr_t = self.learning_rate * ((1.0 - self.beta2.powi(self.t as i32)).sqrt())
            / (1.0 - self.beta1.powi(self.t as i32));

        for i in 0..weights.len() {
            // Update biased first moment estimate for weights
            self.m_weights[i] =
                &self.m_weights[i] * self.beta1 + &weight_grads[i] * (1.0 - self.beta1);

            // Update biased second raw moment estimate for weights
            self.v_weights[i] = &self.v_weights[i] * self.beta2
                + &weight_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            // Update weights
            let update =
                &self.m_weights[i] / (&self.v_weights[i].mapv(|x| x.sqrt()) + self.epsilon);
            weights[i] = &weights[i] - &update * lr_t;

            // Update biased first moment estimate for biases
            self.m_biases[i] = &self.m_biases[i] * self.beta1 + &bias_grads[i] * (1.0 - self.beta1);

            // Update biased second raw moment estimate for biases
            self.v_biases[i] = &self.v_biases[i] * self.beta2
                + &bias_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            // Update biases
            let bias_update =
                &self.m_biases[i] / (&self.v_biases[i].mapv(|x| x.sqrt()) + self.epsilon);
            biases[i] = &biases[i] - &bias_update * lr_t;
        }

        Ok(())
    }
}

/// AdamW optimizer with decoupled weight decay
#[derive(Debug, Clone)]
pub struct AdamWSolver {
    /// Step size used to scale parameter updates
    pub learning_rate: f64,
    /// Exponential decay rate for the first moment estimate
    pub beta1: f64,
    /// Exponential decay rate for the second moment estimate
    pub beta2: f64,
    /// Small constant for numerical stability in denominator
    pub epsilon: f64,
    /// L2 weight decay coefficient applied independently of the gradient
    pub weight_decay: f64,
    m_weights: Vec<Array2<f64>>,
    v_weights: Vec<Array2<f64>>,
    m_biases: Vec<Array1<f64>>,
    v_biases: Vec<Array1<f64>>,
    t: usize,
}

impl AdamWSolver {
    /// Create a new AdamW optimizer with the given hyperparameters
    pub fn new(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
        weight_decay: f64,
    ) -> Self {
        Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
            m_weights: Vec::new(),
            v_weights: Vec::new(),
            m_biases: Vec::new(),
            v_biases: Vec::new(),
            t: 0,
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.m_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.v_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.m_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.v_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.t = 0;
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        self.t += 1;
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);
        let lr_t = self.learning_rate * (bias_correction2.sqrt()) / bias_correction1;

        for i in 0..weights.len() {
            // Update biased first moment estimate for weights
            self.m_weights[i] =
                &self.m_weights[i] * self.beta1 + &weight_grads[i] * (1.0 - self.beta1);

            // Update biased second raw moment estimate for weights
            self.v_weights[i] = &self.v_weights[i] * self.beta2
                + &weight_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            // AdamW update with decoupled weight decay
            let update =
                &self.m_weights[i] / (&self.v_weights[i].mapv(|x| x.sqrt()) + self.epsilon);
            weights[i] =
                &weights[i] * (1.0 - self.learning_rate * self.weight_decay) - &update * lr_t;

            // Update biased first moment estimate for biases (no weight decay)
            self.m_biases[i] = &self.m_biases[i] * self.beta1 + &bias_grads[i] * (1.0 - self.beta1);

            // Update biased second raw moment estimate for biases
            self.v_biases[i] = &self.v_biases[i] * self.beta2
                + &bias_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            // Update biases (no weight decay)
            let bias_update =
                &self.m_biases[i] / (&self.v_biases[i].mapv(|x| x.sqrt()) + self.epsilon);
            biases[i] = &biases[i] - &bias_update * lr_t;
        }

        Ok(())
    }
}

/// RMSprop optimizer
#[derive(Debug, Clone)]
pub struct RMSpropSolver {
    /// Step size controlling the magnitude of parameter updates
    pub learning_rate: f64,
    /// Smoothing constant for the running mean-square of gradients
    pub alpha: f64,
    /// Small constant for numerical stability in denominator
    pub epsilon: f64,
    /// Momentum coefficient; 0.0 disables momentum
    pub momentum: f64,
    /// When true, uses the centered variant that normalizes by gradient mean
    pub centered: bool,
    v_weights: Vec<Array2<f64>>,
    v_biases: Vec<Array1<f64>>,
    momentum_weights: Vec<Array2<f64>>,
    momentum_biases: Vec<Array1<f64>>,
    g_weights: Vec<Array2<f64>>,
    g_biases: Vec<Array1<f64>>,
}

impl RMSpropSolver {
    /// Create a new RMSprop optimizer with the given hyperparameters
    pub fn new(
        learning_rate: f64,
        alpha: f64,
        epsilon: f64,
        momentum: f64,
        centered: bool,
    ) -> Self {
        Self {
            learning_rate,
            alpha,
            epsilon,
            momentum,
            centered,
            v_weights: Vec::new(),
            v_biases: Vec::new(),
            momentum_weights: Vec::new(),
            momentum_biases: Vec::new(),
            g_weights: Vec::new(),
            g_biases: Vec::new(),
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.v_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.v_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.momentum_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.momentum_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();

        if self.centered {
            self.g_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
            self.g_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        }
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        for i in 0..weights.len() {
            // Update accumulated squared gradients for weights
            self.v_weights[i] = &self.v_weights[i] * self.alpha
                + &weight_grads[i].mapv(|x| x * x) * (1.0 - self.alpha);

            let denominator = if self.centered {
                // Update accumulated gradients for centered RMSprop
                self.g_weights[i] =
                    &self.g_weights[i] * self.alpha + &weight_grads[i] * (1.0 - self.alpha);
                (&self.v_weights[i] - &self.g_weights[i].mapv(|x| x * x))
                    .mapv(|x| (x + self.epsilon).sqrt())
            } else {
                self.v_weights[i].mapv(|x| (x + self.epsilon).sqrt())
            };

            // Compute update
            let update = &weight_grads[i] / &denominator;

            if self.momentum > 0.0 {
                // Apply momentum
                self.momentum_weights[i] =
                    &self.momentum_weights[i] * self.momentum + &update * self.learning_rate;
                weights[i] = &weights[i] - &self.momentum_weights[i];
            } else {
                // Direct update
                weights[i] = &weights[i] - &update * self.learning_rate;
            }

            // Same process for biases
            self.v_biases[i] = &self.v_biases[i] * self.alpha
                + &bias_grads[i].mapv(|x| x * x) * (1.0 - self.alpha);

            let bias_denominator = if self.centered {
                self.g_biases[i] =
                    &self.g_biases[i] * self.alpha + &bias_grads[i] * (1.0 - self.alpha);
                (&self.v_biases[i] - &self.g_biases[i].mapv(|x| x * x))
                    .mapv(|x| (x + self.epsilon).sqrt())
            } else {
                self.v_biases[i].mapv(|x| (x + self.epsilon).sqrt())
            };

            let bias_update = &bias_grads[i] / &bias_denominator;

            if self.momentum > 0.0 {
                self.momentum_biases[i] =
                    &self.momentum_biases[i] * self.momentum + &bias_update * self.learning_rate;
                biases[i] = &biases[i] - &self.momentum_biases[i];
            } else {
                biases[i] = &biases[i] - &bias_update * self.learning_rate;
            }
        }

        Ok(())
    }
}

/// Nadam optimizer (Nesterov-accelerated Adam)
#[derive(Debug, Clone)]
pub struct NadamSolver {
    /// Step size for parameter updates
    pub learning_rate: f64,
    /// Exponential decay rate for first moment estimates
    pub beta1: f64,
    /// Exponential decay rate for second raw moment estimates
    pub beta2: f64,
    /// Small constant added to denominator for numerical stability
    pub epsilon: f64,
    m_weights: Vec<Array2<f64>>,
    v_weights: Vec<Array2<f64>>,
    m_biases: Vec<Array1<f64>>,
    v_biases: Vec<Array1<f64>>,
    t: usize,
}

impl NadamSolver {
    /// Create a new Nadam optimizer with the given hyperparameters
    pub fn new(learning_rate: f64, beta1: f64, beta2: f64, epsilon: f64) -> Self {
        Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            m_weights: Vec::new(),
            v_weights: Vec::new(),
            m_biases: Vec::new(),
            v_biases: Vec::new(),
            t: 0,
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.m_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.v_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.m_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.v_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.t = 0;
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        self.t += 1;
        let beta1_t = self.beta1.powi(self.t as i32);
        let beta2_t = self.beta2.powi(self.t as i32);
        let _lr_t = self.learning_rate / (1.0 - beta1_t);

        for i in 0..weights.len() {
            // Update biased first moment estimate
            self.m_weights[i] =
                &self.m_weights[i] * self.beta1 + &weight_grads[i] * (1.0 - self.beta1);

            // Update biased second raw moment estimate
            self.v_weights[i] = &self.v_weights[i] * self.beta2
                + &weight_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            // Bias-corrected second moment estimate
            let v_corrected = &self.v_weights[i] / (1.0 - beta2_t);

            // Nadam update: combines current gradient with momentum term
            let m_hat = (&self.m_weights[i] * self.beta1 + &weight_grads[i] * (1.0 - self.beta1))
                / (1.0 - beta1_t);
            let update = &m_hat / (&v_corrected.mapv(|x| x.sqrt()) + self.epsilon);
            weights[i] = &weights[i] - &update * self.learning_rate;

            // Same for biases
            self.m_biases[i] = &self.m_biases[i] * self.beta1 + &bias_grads[i] * (1.0 - self.beta1);
            self.v_biases[i] = &self.v_biases[i] * self.beta2
                + &bias_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            let v_bias_corrected = &self.v_biases[i] / (1.0 - beta2_t);
            let m_bias_hat = (&self.m_biases[i] * self.beta1 + &bias_grads[i] * (1.0 - self.beta1))
                / (1.0 - beta1_t);
            let bias_update = &m_bias_hat / (&v_bias_corrected.mapv(|x| x.sqrt()) + self.epsilon);
            biases[i] = &biases[i] - &bias_update * self.learning_rate;
        }

        Ok(())
    }
}

/// LARS optimizer (Layer-wise Adaptive Rate Scaling)
/// Designed for large batch training with layer-wise adaptive learning rates
#[derive(Debug, Clone)]
pub struct LarsSolver {
    /// Global base learning rate before layer-wise scaling
    pub learning_rate: f64,
    /// Momentum coefficient for gradient accumulation
    pub momentum: f64,
    /// L2 regularization coefficient added to gradients before momentum update
    pub weight_decay: f64,
    /// LARS-specific scaling coefficient; typically 0.001
    pub lars_coefficient: f64,
    /// Small constant added to denominator to avoid division by zero
    pub epsilon: f64,
    /// Trust ratio coefficient scaling the layer-wise learning rate; typically 1.0
    pub trust_coefficient: f64,
    momentum_weights: Vec<Array2<f64>>,
    momentum_biases: Vec<Array1<f64>>,
}

impl LarsSolver {
    /// Create a new LARS optimizer with the given hyperparameters
    pub fn new(
        learning_rate: f64,
        momentum: f64,
        weight_decay: f64,
        lars_coefficient: f64,
        epsilon: f64,
        trust_coefficient: f64,
    ) -> Self {
        Self {
            learning_rate,
            momentum,
            weight_decay,
            lars_coefficient,
            epsilon,
            trust_coefficient,
            momentum_weights: Vec::new(),
            momentum_biases: Vec::new(),
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.momentum_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.momentum_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        for i in 0..weights.len() {
            // Compute norms for layer-wise adaptation
            let weight_norm = weights[i].iter().map(|x| x * x).sum::<f64>().sqrt();
            let grad_norm = weight_grads[i].iter().map(|x| x * x).sum::<f64>().sqrt();

            // Compute layer-wise learning rate using LARS formula
            let lars_lr = if grad_norm > self.epsilon && weight_norm > self.epsilon {
                self.learning_rate * self.lars_coefficient * weight_norm
                    / (grad_norm + self.weight_decay * weight_norm)
            } else {
                self.learning_rate
            };

            // Apply trust coefficient (layer-wise trust ratio)
            let effective_lr = lars_lr * self.trust_coefficient;

            // Add weight decay to gradients
            let weight_grad_with_decay = &weight_grads[i] + &weights[i] * self.weight_decay;

            // Update momentum for weights
            self.momentum_weights[i] =
                &self.momentum_weights[i] * self.momentum + &weight_grad_with_decay * effective_lr;

            // Update weights
            weights[i] = &weights[i] - &self.momentum_weights[i];

            // For biases, use standard SGD with momentum (no LARS scaling)
            self.momentum_biases[i] =
                &self.momentum_biases[i] * self.momentum + &bias_grads[i] * self.learning_rate;
            biases[i] = &biases[i] - &self.momentum_biases[i];
        }

        Ok(())
    }
}

/// LAMB optimizer (Layer-wise Adaptive Moments optimizer for Batch training)
/// Combines layer-wise adaptation from LARS with Adam's moments
#[derive(Debug, Clone)]
pub struct LambSolver {
    /// Global base learning rate before layer-wise trust ratio scaling
    pub learning_rate: f64,
    /// Exponential decay rate for first moment estimates
    pub beta1: f64,
    /// Exponential decay rate for second raw moment estimates
    pub beta2: f64,
    /// Small constant added to denominator for numerical stability
    pub epsilon: f64,
    /// L2 regularization coefficient added to gradients before moment update
    pub weight_decay: f64,
    /// Maximum trust ratio coefficient capping the layer-wise learning rate; typically 1.0
    pub trust_coefficient: f64,
    m_weights: Vec<Array2<f64>>,
    v_weights: Vec<Array2<f64>>,
    m_biases: Vec<Array1<f64>>,
    v_biases: Vec<Array1<f64>>,
    t: usize,
}

impl LambSolver {
    /// Create a new LAMB optimizer with the given hyperparameters
    pub fn new(
        learning_rate: f64,
        beta1: f64,
        beta2: f64,
        epsilon: f64,
        weight_decay: f64,
        trust_coefficient: f64,
    ) -> Self {
        Self {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
            trust_coefficient,
            m_weights: Vec::new(),
            v_weights: Vec::new(),
            m_biases: Vec::new(),
            v_biases: Vec::new(),
            t: 0,
        }
    }

    /// Initialize optimizer state buffers to match shapes of `weights` and `biases`
    pub fn initialize(&mut self, weights: &[Array2<f64>], biases: &[Array1<f64>]) {
        self.m_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.v_weights = weights.iter().map(|w| Array2::zeros(w.dim())).collect();
        self.m_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.v_biases = biases.iter().map(|b| Array1::zeros(b.len())).collect();
        self.t = 0;
    }

    /// Perform one optimization step using the provided gradients; updates `weights` and `biases` in place
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        self.t += 1;

        for i in 0..weights.len() {
            // Add weight decay to gradients
            let weight_grad_with_decay = &weight_grads[i] + &weights[i] * self.weight_decay;

            // Update biased first moment estimate
            self.m_weights[i] =
                &self.m_weights[i] * self.beta1 + &weight_grad_with_decay * (1.0 - self.beta1);

            // Update biased second raw moment estimate
            self.v_weights[i] = &self.v_weights[i] * self.beta2
                + &weight_grad_with_decay.mapv(|x| x * x) * (1.0 - self.beta2);

            // Bias correction
            let m_hat = &self.m_weights[i] / (1.0 - self.beta1.powi(self.t as i32));
            let v_hat = &self.v_weights[i] / (1.0 - self.beta2.powi(self.t as i32));

            // Compute Adam update
            let adam_update = &m_hat / (&v_hat.mapv(|x| x.sqrt()) + self.epsilon);

            // Compute norms for layer-wise adaptation
            let weight_norm = weights[i].iter().map(|x| x * x).sum::<f64>().sqrt();
            let adam_update_norm = adam_update.iter().map(|x| x * x).sum::<f64>().sqrt();

            // Compute layer-wise learning rate using LAMB formula
            let trust_ratio = if weight_norm > 0.0 && adam_update_norm > 0.0 {
                (weight_norm / adam_update_norm).min(self.trust_coefficient)
            } else {
                1.0
            };

            // Apply layer-wise adaptive learning rate
            let effective_lr = self.learning_rate * trust_ratio;

            // Update weights
            weights[i] = &weights[i] - &adam_update * effective_lr;

            // For biases, use standard Adam (no layer-wise adaptation)
            self.m_biases[i] = &self.m_biases[i] * self.beta1 + &bias_grads[i] * (1.0 - self.beta1);
            self.v_biases[i] = &self.v_biases[i] * self.beta2
                + &bias_grads[i].mapv(|x| x * x) * (1.0 - self.beta2);

            let m_bias_hat = &self.m_biases[i] / (1.0 - self.beta1.powi(self.t as i32));
            let v_bias_hat = &self.v_biases[i] / (1.0 - self.beta2.powi(self.t as i32));

            let bias_update = &m_bias_hat / (&v_bias_hat.mapv(|x| x.sqrt()) + self.epsilon);
            biases[i] = &biases[i] - &bias_update * self.learning_rate;
        }

        Ok(())
    }
}

/// Limited-memory BFGS (L-BFGS) optimizer.
///
/// Maintains a short history of parameter and gradient differences and uses the
/// classic two-loop recursion (Nocedal & Wright, *Numerical Optimization*,
/// Algorithm 7.4) to compute the quasi-Newton search direction `d = -H_k g_k`
/// without ever forming the inverse Hessian. The configured `learning_rate` is
/// applied as a fixed step length in place of a Wolfe line search, because the
/// streaming per-batch `update_params` interface does not expose the loss
/// function a line search would require. With an empty history the direction
/// reduces to scaled steepest descent, and curvature pairs that fail the
/// `sᵀy > epsilon` condition are skipped to keep the implicit Hessian positive
/// definite.
#[derive(Debug, Clone)]
pub struct LbfgsSolver {
    /// Step length applied to the L-BFGS search direction
    pub learning_rate: f64,
    /// Maximum number of `(s, y)` correction pairs retained
    pub history_size: usize,
    /// Curvature threshold; pairs with `sᵀy` at or below this are skipped
    pub epsilon: f64,
    s_history: VecDeque<Vec<f64>>,
    y_history: VecDeque<Vec<f64>>,
    rho_history: VecDeque<f64>,
    prev_params: Option<Vec<f64>>,
    prev_grad: Option<Vec<f64>>,
}

impl LbfgsSolver {
    /// Create a new L-BFGS optimizer with the given step length, memory depth,
    /// and curvature threshold.
    pub fn new(learning_rate: f64, history_size: usize, epsilon: f64) -> Self {
        Self {
            learning_rate,
            history_size: history_size.max(1),
            epsilon,
            s_history: VecDeque::new(),
            y_history: VecDeque::new(),
            rho_history: VecDeque::new(),
            prev_params: None,
            prev_grad: None,
        }
    }

    /// Reset all optimizer state; parameter shapes are inferred from the first update.
    pub fn initialize(&mut self, _weights: &[Array2<f64>], _biases: &[Array1<f64>]) {
        self.s_history.clear();
        self.y_history.clear();
        self.rho_history.clear();
        self.prev_params = None;
        self.prev_grad = None;
    }

    /// Flatten weight matrices followed by bias vectors into a single parameter vector.
    fn flatten(weights: &[Array2<f64>], biases: &[Array1<f64>]) -> Vec<f64> {
        let total: usize = weights.iter().map(|w| w.len()).sum::<usize>()
            + biases.iter().map(|b| b.len()).sum::<usize>();
        let mut flat = Vec::with_capacity(total);
        for w in weights {
            flat.extend(w.iter().copied());
        }
        for b in biases {
            flat.extend(b.iter().copied());
        }
        flat
    }

    /// Scatter a flattened parameter vector back into weight matrices and bias vectors.
    fn write_back(flat: &[f64], weights: &mut [Array2<f64>], biases: &mut [Array1<f64>]) {
        let mut offset = 0;
        for w in weights.iter_mut() {
            for value in w.iter_mut() {
                *value = flat[offset];
                offset += 1;
            }
        }
        for b in biases.iter_mut() {
            for value in b.iter_mut() {
                *value = flat[offset];
                offset += 1;
            }
        }
    }

    fn dot(a: &[f64], b: &[f64]) -> f64 {
        a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    }

    /// Perform one L-BFGS step using the provided gradients; updates `weights`
    /// and `biases` in place.
    pub fn update_params(
        &mut self,
        weights: &mut [Array2<f64>],
        biases: &mut [Array1<f64>],
        weight_grads: &[Array2<f64>],
        bias_grads: &[Array1<f64>],
    ) -> NeuralResult<()> {
        let x = Self::flatten(weights, biases);
        let g = Self::flatten(weight_grads, bias_grads);

        // Record the curvature pair (s, y) produced by the previous step.
        if let (Some(prev_x), Some(prev_g)) = (&self.prev_params, &self.prev_grad) {
            let s: Vec<f64> = x.iter().zip(prev_x).map(|(a, b)| a - b).collect();
            let y: Vec<f64> = g.iter().zip(prev_g).map(|(a, b)| a - b).collect();
            let sy = Self::dot(&s, &y);
            if sy > self.epsilon {
                if self.s_history.len() == self.history_size {
                    self.s_history.pop_front();
                    self.y_history.pop_front();
                    self.rho_history.pop_front();
                }
                self.rho_history.push_back(1.0 / sy);
                self.s_history.push_back(s);
                self.y_history.push_back(y);
            }
        }

        // First loop: q starts at the gradient, peeling off the newest corrections.
        let m = self.s_history.len();
        let mut q = g.clone();
        let mut alphas = vec![0.0_f64; m];
        for i in (0..m).rev() {
            let alpha = self.rho_history[i] * Self::dot(&self.s_history[i], &q);
            alphas[i] = alpha;
            for (qj, yj) in q.iter_mut().zip(self.y_history[i].iter()) {
                *qj -= alpha * yj;
            }
        }

        // Initial inverse-Hessian scaling gamma = (sᵀy)/(yᵀy) from the latest pair.
        let gamma = if m > 0 {
            let yy = Self::dot(&self.y_history[m - 1], &self.y_history[m - 1]);
            if yy > 0.0 {
                Self::dot(&self.s_history[m - 1], &self.y_history[m - 1]) / yy
            } else {
                1.0
            }
        } else {
            1.0
        };
        let mut r: Vec<f64> = q.iter().map(|value| value * gamma).collect();

        // Second loop: re-apply the corrections from oldest to newest.
        for i in 0..m {
            let beta = self.rho_history[i] * Self::dot(&self.y_history[i], &r);
            let coeff = alphas[i] - beta;
            for (rj, sj) in r.iter_mut().zip(self.s_history[i].iter()) {
                *rj += coeff * sj;
            }
        }

        // r now holds H_k g_k; descend with d = -r scaled by the step length.
        let new_x: Vec<f64> = x
            .iter()
            .zip(r.iter())
            .map(|(xj, rj)| xj - self.learning_rate * rj)
            .collect();
        Self::write_back(&new_x, weights, biases);

        // Stash the pre-update params/grad for the next curvature pair.
        self.prev_params = Some(x);
        self.prev_grad = Some(g);

        Ok(())
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use scirs2_core::ndarray::{array, Array1, Array2};
    use scirs2_core::numeric::Signed;

    #[test]
    fn test_sgd_solver_initialization() {
        let mut solver =
            SgdSolver::new(0.01, 0.9, false, LearningRateSchedule::Constant, 0.5, 0.01);

        let weights = vec![Array2::zeros((3, 2)), Array2::zeros((2, 1))];
        let biases = vec![Array1::zeros(2), Array1::zeros(1)];

        solver.initialize(&weights, &biases);

        assert_eq!(solver.velocity_weights.len(), 2);
        assert_eq!(solver.velocity_biases.len(), 2);
        assert_eq!(solver.t, 0);
    }

    #[test]
    fn test_adam_solver_initialization() {
        let mut solver = AdamSolver::new(0.001, 0.9, 0.999, 1e-8);

        let weights = vec![Array2::zeros((3, 2)), Array2::zeros((2, 1))];
        let biases = vec![Array1::zeros(2), Array1::zeros(1)];

        solver.initialize(&weights, &biases);

        assert_eq!(solver.m_weights.len(), 2);
        assert_eq!(solver.v_weights.len(), 2);
        assert_eq!(solver.m_biases.len(), 2);
        assert_eq!(solver.v_biases.len(), 2);
        assert_eq!(solver.t, 0);
    }

    #[test]
    fn test_sgd_learning_rate_schedules() {
        let solver = SgdSolver::new(0.01, 0.0, false, LearningRateSchedule::Constant, 0.5, 0.1);
        assert_abs_diff_eq!(solver.get_learning_rate(), 0.01, epsilon = 1e-10);

        let mut solver =
            SgdSolver::new(0.01, 0.0, false, LearningRateSchedule::InvScaling, 0.5, 0.1);
        solver.t = 4;
        let expected_lr = 0.1 / (4.0_f64).powf(0.5);
        assert_abs_diff_eq!(solver.get_learning_rate(), expected_lr, epsilon = 1e-10);
    }

    #[test]
    fn test_parameter_updates_change_weights() {
        let mut solver = SgdSolver::new(0.1, 0.0, false, LearningRateSchedule::Constant, 0.5, 0.01);

        let mut weights = vec![array![[1.0, 2.0], [3.0, 4.0]]];
        let mut biases = vec![array![0.5, -0.5]];
        let weight_grads = vec![array![[0.1, 0.2], [0.3, 0.4]]];
        let bias_grads = vec![array![0.1, -0.1]];

        let original_weights = weights[0].clone();
        let original_biases = biases[0].clone();

        solver.initialize(&weights, &biases);
        solver
            .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
            .expect("operation should succeed");

        // Check that weights and biases have changed (compare element by element)
        let weights_changed = weights[0]
            .iter()
            .zip(original_weights.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(weights_changed);
        let biases_changed = biases[0]
            .iter()
            .zip(original_biases.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(biases_changed);
    }

    #[test]
    fn test_adamw_solver_initialization() {
        let mut solver = AdamWSolver::new(0.001, 0.9, 0.999, 1e-8, 0.01);

        let weights = vec![Array2::zeros((3, 2)), Array2::zeros((2, 1))];
        let biases = vec![Array1::zeros(2), Array1::zeros(1)];

        solver.initialize(&weights, &biases);

        assert_eq!(solver.m_weights.len(), 2);
        assert_eq!(solver.v_weights.len(), 2);
        assert_eq!(solver.weight_decay, 0.01);
        assert_eq!(solver.t, 0);
    }

    #[test]
    fn test_rmsprop_solver_initialization() {
        let mut solver = RMSpropSolver::new(0.001, 0.9, 1e-8, 0.0, false);

        let weights = vec![Array2::zeros((3, 2)), Array2::zeros((2, 1))];
        let biases = vec![Array1::zeros(2), Array1::zeros(1)];

        solver.initialize(&weights, &biases);

        assert_eq!(solver.v_weights.len(), 2);
        assert_eq!(solver.v_biases.len(), 2);
        assert!(!solver.centered);
    }

    #[test]
    fn test_rmsprop_centered_initialization() {
        let mut solver = RMSpropSolver::new(0.001, 0.9, 1e-8, 0.0, true);

        let weights = vec![Array2::zeros((3, 2))];
        let biases = vec![Array1::zeros(2)];

        solver.initialize(&weights, &biases);

        assert!(solver.centered);
        assert_eq!(solver.g_weights.len(), 1);
        assert_eq!(solver.g_biases.len(), 1);
    }

    #[test]
    fn test_nadam_solver_initialization() {
        let mut solver = NadamSolver::new(0.001, 0.9, 0.999, 1e-8);

        let weights = vec![Array2::zeros((3, 2)), Array2::zeros((2, 1))];
        let biases = vec![Array1::zeros(2), Array1::zeros(1)];

        solver.initialize(&weights, &biases);

        assert_eq!(solver.m_weights.len(), 2);
        assert_eq!(solver.v_weights.len(), 2);
        assert_eq!(solver.t, 0);
    }

    #[test]
    fn test_adamw_weight_decay() {
        let mut solver = AdamWSolver::new(0.1, 0.9, 0.999, 1e-8, 0.1);

        let mut weights = vec![array![[1.0, 2.0]]];
        let mut biases = vec![array![0.5]];
        let weight_grads = vec![array![[0.0, 0.0]]]; // Zero gradients to isolate weight decay
        let bias_grads = vec![array![0.0]];

        let original_weight_magnitude = weights[0].iter().map(|x| x.abs()).sum::<f64>();

        solver.initialize(&weights, &biases);
        solver
            .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
            .expect("operation should succeed");

        let new_weight_magnitude = weights[0].iter().map(|x| x.abs()).sum::<f64>();

        // Weight magnitude should decrease due to weight decay
        assert!(new_weight_magnitude < original_weight_magnitude);
    }

    #[test]
    fn test_lars_solver() {
        let mut solver = LarsSolver::new(0.1, 0.9, 0.01, 0.001, 1e-8, 1.0);

        let mut weights = vec![array![[1.0, 2.0], [3.0, 4.0]]];
        let mut biases = vec![array![0.5, -0.5]];
        let weight_grads = vec![array![[0.1, 0.2], [0.3, 0.4]]];
        let bias_grads = vec![array![0.1, -0.1]];

        let original_weights = weights[0].clone();
        let original_biases = biases[0].clone();

        solver.initialize(&weights, &biases);
        solver
            .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
            .expect("operation should succeed");

        // Check that parameters have changed (compare element by element)
        let weights_changed = weights[0]
            .iter()
            .zip(original_weights.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(weights_changed);
        let biases_changed = biases[0]
            .iter()
            .zip(original_biases.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(biases_changed);
    }

    #[test]
    fn test_lamb_solver() {
        let mut solver = LambSolver::new(0.001, 0.9, 0.999, 1e-8, 0.01, 1.0);

        let mut weights = vec![array![[1.0, 2.0], [3.0, 4.0]]];
        let mut biases = vec![array![0.5, -0.5]];
        let weight_grads = vec![array![[0.1, 0.2], [0.3, 0.4]]];
        let bias_grads = vec![array![0.1, -0.1]];

        let original_weights = weights[0].clone();
        let original_biases = biases[0].clone();

        solver.initialize(&weights, &biases);
        solver
            .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
            .expect("operation should succeed");

        // Check that parameters have changed (compare element by element)
        let weights_changed = weights[0]
            .iter()
            .zip(original_weights.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(weights_changed);
        let biases_changed = biases[0]
            .iter()
            .zip(original_biases.iter())
            .any(|(a, b)| (a - b).abs() > 1e-10);
        assert!(biases_changed);
    }

    #[test]
    fn test_lbfgs_solver_updates_parameters() {
        let mut solver = LbfgsSolver::new(0.1, 10, 1e-10);

        let mut weights = vec![array![[1.0, 2.0], [3.0, 4.0]]];
        let mut biases = vec![array![0.5, -0.5]];
        let weight_grads = vec![array![[0.1, 0.2], [0.3, 0.4]]];
        let bias_grads = vec![array![0.1, -0.1]];

        let original_weights = weights[0].clone();

        solver.initialize(&weights, &biases);
        solver
            .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
            .expect("operation should succeed");

        // First step (empty history) is scaled steepest descent: w - lr * g.
        let expected = &original_weights - &weight_grads[0] * 0.1;
        for (a, b) in weights[0].iter().zip(expected.iter()) {
            assert_abs_diff_eq!(*a, *b, epsilon = 1e-12);
        }
    }

    #[test]
    fn test_lbfgs_minimizes_quadratic() {
        // Minimize f(w) = 0.5 * sum((w - target)^2); gradient = w - target.
        // L-BFGS should drive the parameters to `target`.
        let target = array![[2.0, -3.0], [0.5, 4.0]];
        let mut solver = LbfgsSolver::new(1.0, 10, 1e-12);

        let mut weights = vec![Array2::<f64>::zeros((2, 2))];
        let mut biases = vec![Array1::<f64>::zeros(0)];
        solver.initialize(&weights, &biases);

        for _ in 0..50 {
            let grad = &weights[0] - &target;
            let weight_grads = vec![grad];
            let bias_grads = vec![Array1::<f64>::zeros(0)];
            solver
                .update_params(&mut weights, &mut biases, &weight_grads, &bias_grads)
                .expect("operation should succeed");
        }

        for (w, t) in weights[0].iter().zip(target.iter()) {
            assert_abs_diff_eq!(*w, *t, epsilon = 1e-6);
        }
    }

    #[test]
    fn test_gradient_clipping_by_norm() {
        let mut weight_grads = vec![array![[10.0, 20.0], [30.0, 40.0]]];
        let mut bias_grads = vec![array![5.0, -5.0]];

        // Compute original norm
        let original_norm = (weight_grads[0].iter().map(|x| x * x).sum::<f64>()
            + bias_grads[0].iter().map(|x| x * x).sum::<f64>())
        .sqrt();

        let max_norm = 10.0;
        GradientClipper::clip_gradients(
            &mut weight_grads,
            &mut bias_grads,
            GradientClipping::ByNorm(max_norm),
        );

        // Compute new norm
        let new_norm = (weight_grads[0].iter().map(|x| x * x).sum::<f64>()
            + bias_grads[0].iter().map(|x| x * x).sum::<f64>())
        .sqrt();

        assert!(original_norm > max_norm);
        assert_abs_diff_eq!(new_norm, max_norm, epsilon = 1e-6);
    }

    #[test]
    fn test_gradient_clipping_by_value() {
        let mut weight_grads = vec![array![[10.0, -15.0], [20.0, -25.0]]];
        let mut bias_grads = vec![array![12.0, -8.0]];

        let max_value = 5.0;
        GradientClipper::clip_gradients(
            &mut weight_grads,
            &mut bias_grads,
            GradientClipping::ByValue(max_value),
        );

        // Check that all values are within [-max_value, max_value]
        for grad in &weight_grads {
            for &value in grad.iter() {
                assert!(value >= -max_value && value <= max_value);
            }
        }

        for grad in &bias_grads {
            for &value in grad.iter() {
                assert!(value >= -max_value && value <= max_value);
            }
        }
    }

    #[test]
    fn test_learning_rate_schedules() {
        // Test exponential decay
        let mut solver = SgdSolver::new_with_schedule_params(
            1.0,
            0.0,
            false,
            LearningRateSchedule::ExponentialDecay,
            0.5,
            0.1,
            0.9,
            100,
            1000,
            0.0,
            1.0,
            500,
            0.1,
            100,
            1000,
        );
        solver.t = 100;
        let lr = solver.get_learning_rate();
        let expected = 1.0 * 0.9_f64.powf(100.0 / 100.0);
        assert_abs_diff_eq!(lr, expected, epsilon = 1e-6);

        // Test cosine annealing
        let mut solver = SgdSolver::new_with_schedule_params(
            1.0,
            0.0,
            false,
            LearningRateSchedule::CosineAnnealing,
            0.5,
            0.1,
            0.9,
            100,
            1000,
            0.0,
            1.0,
            500,
            0.1,
            100,
            1000,
        );
        solver.t = 500;
        let lr = solver.get_learning_rate();
        let progress = 500.0 / 1000.0;
        let expected = 0.0 + (1.0 - 0.0) * (1.0 + (std::f64::consts::PI * progress).cos()) / 2.0;
        assert_abs_diff_eq!(lr, expected, epsilon = 1e-6);

        // Test warmup
        let mut solver = SgdSolver::new_with_schedule_params(
            1.0,
            0.0,
            false,
            LearningRateSchedule::Warmup,
            0.5,
            0.1,
            0.9,
            100,
            1000,
            0.0,
            1.0,
            500,
            0.1,
            100,
            1000,
        );
        solver.t = 50;
        let lr = solver.get_learning_rate();
        let expected = 1.0 * (50.0 / 100.0);
        assert_abs_diff_eq!(lr, expected, epsilon = 1e-6);
    }
}