//! Activation functions for neural-network layers

use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize, Clone)]
pub enum Activation {
    Sigmoid,
    ReLU
}

/// Applies an activation function to a single scalar.
///
/// `Sigmoid` maps `z` to `1 / (1 + e^-z)`, squashing it into the open range
/// `(0, 1)`. `ReLU` returns `z` when it is positive and `0` otherwise.
///
/// # Examples
///
/// ```
/// use tuphium::activations::{activate_scalar, Activation};
///
/// assert_eq!(activate_scalar(-3.0, &Activation::ReLU), 0.0);
/// assert_eq!(activate_scalar(2.0, &Activation::ReLU), 2.0);
/// ```
pub fn activate_scalar(input: f32, act: &Activation) -> f32 {
    match act {
        Activation::Sigmoid => 1.0 / (1.0 + (-input).exp()),
        Activation::ReLU => if input > 0.0 { input } else { 0.0 },
    }
}

/// Applies an activation function element-wise to a vector.
///
/// Each element is transformed independently with [`activate_scalar`], so the
/// returned vector has the same length as the input.
///
/// # Examples
///
/// ```
/// use tuphium::activations::{activate, Activation};
///
/// let out = activate(vec![-1.0, 0.5], &Activation::ReLU);
/// assert_eq!(out, vec![0.0, 0.5]);
/// ```
pub fn activate(inputs: Vec<f32>, act: &Activation) -> Vec<f32> {
    inputs
        .iter()
        .map(|&input| activate_scalar(input, act))
        .collect()
}

/// Returns the derivative of the activation, expressed in terms of its output.
///
/// Both activations are differentiated with respect to their pre-activation
/// input but written using the already-computed `output`, which is what the
/// backward pass has available: `Sigmoid` derivative is `output * (1 - output)`
/// and `ReLU` derivative is `1` when `output > 0` and `0` otherwise.
///
/// # Examples
///
/// ```
/// use tuphium::activations::{derivative, Activation};
///
/// assert_eq!(derivative(0.0, &Activation::ReLU), 0.0);
/// assert_eq!(derivative(4.0, &Activation::ReLU), 1.0);
/// ```
pub fn derivative(output: f32, act: &Activation) -> f32 {
    match act {
        Activation::Sigmoid => output * (1.0 - output),
        Activation::ReLU => if output > 0.0 { 1.0 } else { 0.0 }
    }
}
