use crate::error::TuphiumResult;
use crate::optimizers::Optimizer;

/// A trainable building block of a network (a layer or a whole model).
///
/// Implementors cache whatever they need from [`Module::forward`] so that a
/// later [`Module::backward`] can compute gradients. A typical training step is
/// `forward` → loss → `backward` → [`Module::step`] → [`Module::zero_grad`].
pub trait Module {
    /// Runs the forward pass on `input` and returns this module's output.
    ///
    /// Implementors usually cache `input` and the output for the backward pass.
    fn forward(&mut self, input: &[f32]) -> TuphiumResult<Vec<f32>>;

    /// Runs the backward pass.
    ///
    /// Takes the gradient of the loss with respect to this module's output
    /// (`grad_output`), accumulates the parameter gradients internally, and
    /// returns the gradient with respect to this module's input so it can be
    /// chained to the previous module.
    fn backward(&mut self, grad_output: &[f32]) -> TuphiumResult<Vec<f32>>;

    /// Returns mutable references to every trainable parameter.
    fn params(&mut self) -> Vec<&mut f32>;

    /// Returns the currently accumulated gradients, aligned with [`Module::params`].
    fn grads(&self) -> Vec<f32>;

    /// Clears the accumulated gradients. Defaults to a no-op.
    fn zero_grad(&mut self) {}

    /// Applies one optimizer update to the parameters. Defaults to a no-op.
    fn step(&mut self, _opt: &mut dyn Optimizer) -> TuphiumResult<()> {
        Ok(())
    }
}
