//! Parameter-update rules (SGD, Adam, etc.) that adjust model weights from
//! computed gradients during training.

use serde::{Serialize, Deserialize};

use crate::error::{TuphiumError, TuphiumResult};

/// A parameter-update rule applied during training.
pub trait Optimizer {
    /// Updates `params` in place from their `grads`.
    ///
    /// `params` and `grads` are paired element-by-element; entries past the
    /// shorter of the two are left untouched.
    fn update(&mut self, params: &mut [f32], grads: &[f32]) -> TuphiumResult<()>;
}

/// Stochastic Gradient Descent with a fixed learning rate `lr`.
#[derive(Serialize, Deserialize)]
pub struct SGD {
    /// Learning rate: the step size applied along the negative gradient.
    pub lr: f32,
}

impl Optimizer for SGD {
    /// Applies the SGD step `p -= lr * g` to every parameter.
    fn update(&mut self, params: &mut [f32], grads: &[f32]) -> TuphiumResult<()> {
        if params.len() != grads.len() {
            return Err(TuphiumError::LengthMismatch {
                operation: "SGD::update",
                expected: params.len(),
                actual: grads.len(),
            });
        }
        if !self.lr.is_finite() {
            return Err(TuphiumError::NonFiniteValue {
                operation: "SGD::update",
                value: self.lr,
            });
        }
        if let Some(&value) = params
            .iter()
            .chain(grads.iter())
            .find(|value| !value.is_finite())
        {
            return Err(TuphiumError::NonFiniteValue {
                operation: "SGD::update",
                value,
            });
        }
        for (param, grad) in params.iter_mut().zip(grads.iter()) {
            *param -= self.lr * grad;
        }
        Ok(())
    }
}
