//! # Neural Network Layers
//!
//! This module defines the core layer types used in the Smyl library:
//! `SignalLayer`, `WeightLayer`, `BiasLayer`, and `SynapseLayer`.
//!
//! ## Signal Layer
//!
//! A `SignalLayer` is a 1xN matrix representing the output or input of a neural
//! network layer.
//!
//! ## Weight Layer
//!
//! A `WeightLayer` is an MxN matrix representing the weights connecting the
//! previous layer to the current layer.
//!
//! ## Bias Layer
//!
//! A `BiasLayer` is a 1xN matrix representing the biases of a neural network
//! layer.
//!
//! ## Synapse Layer
//!
//! A `SynapseLayer` is a combination of a `WeightLayer` and a `BiasLayer`,
//! along with an `ActivationFunction` that is applied to the layer's output.
//!
//! ### Usage
//!
//! ```rust
//! use smyl::network::LayerGradient;
//! use smyl::prelude::*;
//!
//! // Create a new synapse layer with ReLU activation
//! let mut layer = SynapseLayer::<f64, ReLU>::zero(2, 3);
//!
//! // Forward propagation
//! let inputs = Matrix::from([[1.0, 2.0]]);
//! let outputs = layer.forward(inputs.clone());
//!
//! // Backward propagation
//! let output_gradient = Matrix::from([[0.1, 0.2, 0.3]]);
//! let mut layer_gradient = LayerGradient::zero(2, 3);
//! let input_gradient = layer.backward(&inputs, outputs, output_gradient, &mut layer_gradient);
//! ```

use std::fmt::{Debug, Display};
use std::marker::PhantomData;
use std::ops::{Add, AddAssign, Div, Mul, Sub, SubAssign};

use crate::maths::activation::Function;
use crate::maths::matrix::Matrix;

/// A struct used to compute average of multiple gradients calculated for each
/// examples
#[derive(Debug, Clone)]
pub struct LayerGradient<T>
where
    T: Copy,
{
    weights:       Matrix<T>,
    biases:        Matrix<T>,
    example_count: i32,
}

impl<T> LayerGradient<T>
where
    T: Copy,
{
    #[inline]
    pub fn expect_dimensions<L>(&self, previous_size: usize, size: usize, label: L)
    where
        L: Display,
    {
        self.weights
            .expect_dimensions(previous_size, size, format!("{label}: weights"));
        self.biases
            .expect_dimensions(1, size, format!("{label}: biases"));
    }

    /// Constructs a new layer gradient filled with zeros
    #[must_use]
    #[inline]
    pub fn zero(previous_size: usize, size: usize) -> Self
    where
        T: Default,
    {
        Self {
            weights:       Matrix::zero(previous_size, size),
            biases:        Matrix::zero(1, size),
            example_count: 0,
        }
    }

    /// Add a gradient for an individual example into the average computation
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::network::LayerGradient;
    /// use smyl::prelude::Matrix;
    ///
    /// let mut gradient = LayerGradient::<i32>::zero(2, 3);
    /// gradient.add(
    ///     Matrix::from([[1, 2, 3], [4, 5, 6]]),
    ///     Matrix::from([[7, 8, 9]]),
    /// );
    /// ```
    pub fn add(&mut self, weight_gradient: Matrix<T>, bias_gradient: Matrix<T>)
    where
        T: AddAssign<T> + Copy,
    {
        weight_gradient.expect_dimensions_of(&self.weights, "weight_gradient");
        bias_gradient.expect_dimensions_of(&self.biases, "bias_gradient");

        self.weights += weight_gradient;
        self.biases += bias_gradient;
        self.example_count += 1;
    }

    /// Compute the average gradient from the accumulated examples
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::network::LayerGradient;
    /// use smyl::prelude::Matrix;
    ///
    /// let mut gradient = LayerGradient::<f64>::zero(2, 3);
    /// gradient.add(
    ///     Matrix::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
    ///     Matrix::from([[7.0, 8.0, 9.0]]),
    /// );
    ///
    /// let mut weights = Matrix::zero(2, 3);
    /// let mut biases = Matrix::zero(1, 3);
    ///
    /// gradient.apply(&mut weights, &mut biases, -1.0);
    ///
    /// assert_eq!(
    ///     (weights, biases),
    ///     (
    ///         Matrix::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
    ///         Matrix::from([[7.0, 8.0, 9.0]])
    ///     )
    /// );
    ///
    /// gradient.add(
    ///     Matrix::from([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]),
    ///     Matrix::from([[9.0, 10.0, 11.0]]),
    /// );
    ///
    /// dbg!(&gradient);
    ///
    /// let mut weights = Matrix::zero(2, 3);
    /// let mut biases = Matrix::zero(1, 3);
    ///
    /// gradient.apply(&mut weights, &mut biases, -1.0);
    ///
    /// assert_eq!(
    ///     (weights, biases),
    ///     (
    ///         Matrix::from([[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]]),
    ///         Matrix::from([[8.0, 9.0, 10.0]])
    ///     )
    /// );
    ///
    /// gradient.add(
    ///     Matrix::from([[-1.0, 0.0, 1.0], [2.0, 3.0, 4.0]]),
    ///     Matrix::from([[5.0, 6.0, 7.0]]),
    /// );
    ///
    /// let mut weights = Matrix::zero(2, 3);
    /// let mut biases = Matrix::zero(1, 3);
    ///
    /// gradient.apply(&mut weights, &mut biases, -1.0);
    ///
    /// assert_eq!(
    ///     (weights, biases),
    ///     (
    ///         Matrix::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
    ///         Matrix::from([[7.0, 8.0, 9.0]])
    ///     )
    /// );
    /// ```
    #[inline]
    pub fn apply(&self, weights: &mut Matrix<T>, biases: &mut Matrix<T>, learning_rate: T)
    where
        T: Div<T, Output = T> + Debug + Mul<T, Output = T> + SubAssign<T> + Copy + From<i32>,
    {
        if self.example_count != 0 {
            *weights -= &self.weights * (learning_rate / T::from(self.example_count));
            *biases -= &self.biases * (learning_rate / T::from(self.example_count));
        }
    }

    #[inline]
    pub fn empty(&mut self)
    where
        T: Default,
    {
        self.weights.empty();
        self.biases.empty();
        self.example_count = 0;
    }
}

/// A struct representing a synapse layer in a neural network.
///
/// A `SynapseLayer` combines a `WeightLayer` and a `BiasLayer`, along with an
/// `ActivationFunction` that is applied to the layer's output.
#[derive(Debug, Clone)]
pub struct SynapseLayer<T, A>
where
    T: Copy + Clone + Default + Debug,
    A: Function<T>,
{
    weights:    Matrix<T>,
    biases:     Matrix<T>,
    activation: PhantomData<A>,
}

impl<T, A> SynapseLayer<T, A>
where
    T: Copy + Clone + Default + Debug,
    A: Function<T>,
{
    #[must_use]
    #[inline]
    pub const fn size(&self) -> usize { self.weights.width() }

    #[must_use]
    #[inline]
    pub const fn previous_size(&self) -> usize { self.weights.height() }

    #[must_use]
    #[inline]
    pub fn zero(previous_size: usize, size: usize) -> Self {
        Self::new(Matrix::zero(previous_size, size), Matrix::zero(1, size))
    }

    #[cfg(feature = "random")]
    #[must_use]
    pub fn random<R>(previous_size: usize, size: usize, rng: &mut R) -> Self
    where
        R: rand::Rng + ?Sized,
        rand::distributions::Standard: rand::distributions::Distribution<T>,
    {
        let mut layer = Self::zero(previous_size, size);
        rng.fill(&mut layer);
        layer
    }

    /// Creates a new `SynapseLayer` with the given weights and biases.
    #[must_use]
    pub fn new(weights: Matrix<T>, biases: Matrix<T>) -> Self {
        biases.expect_dimensions(1, weights.width(), "biases");

        Self {
            weights,
            biases,
            activation: PhantomData,
        }
    }

    /// Returns a reference to the weights of the synapse layer.
    #[must_use]
    #[inline]
    pub const fn weights(&self) -> &Matrix<T> { &self.weights }

    /// Returns a reference to the biases of the synapse layer.
    #[must_use]
    #[inline]
    pub const fn biases(&self) -> &Matrix<T> { &self.biases }

    /// Performs forward propagation through the synapse layer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::prelude::*;
    ///
    /// let weights = Matrix::from([[1.0, 0.3, 0.6], [0.5, 0.0, 0.2]]);
    /// let biases = Matrix::from([[0.2, 0.8, 0.5]]);
    ///
    /// let mut layer = SynapseLayer::<f64, ReLU>::new(weights, biases);
    /// let inputs = Matrix::from([[1.0, 2.0]]);
    /// let outputs = layer.forward(inputs);
    /// assert_eq!(outputs, Matrix::from([[2.2, 1.1, 1.5]]));
    /// ```
    #[must_use]
    pub fn forward(&self, input: Matrix<T>) -> Matrix<T>
    where
        T: Add<T, Output = T> + AddAssign + Mul<T, Output = T> + std::iter::Sum,
    {
        input.expect_dimensions(1, self.previous_size(), "input");

        input.dot(&self.weights).add(&self.biases).activate::<A>()
    }

    /// Apply the average of multiple gradients computed from
    /// `SynapseLayer::backward`
    ///
    /// ```rust
    /// use smyl::maths::matrix::Matrix;
    /// use smyl::prelude::{LayerGradient, ReLU, SynapseLayer};
    ///
    /// let learning_rate = 2.0;
    /// let mut gradient = LayerGradient::zero(2, 3);
    ///
    /// gradient.add(
    ///     Matrix::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
    ///     Matrix::from([[7.0, 8.0, 9.0]]),
    /// );
    ///
    /// let mut layer = SynapseLayer::<f64, ReLU>::zero(2, 3);
    /// layer.apply_gradient(&gradient, learning_rate);
    ///
    /// assert_eq!(
    ///     layer.weights(),
    ///     &Matrix::from([[-2.0, -4.0, -6.0], [-8.0, -10.0, -12.0]])
    /// );
    /// assert_eq!(layer.biases(), &Matrix::from([[-14.0, -16.0, -18.0]]));
    ///
    /// gradient.add(
    ///     Matrix::from([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]),
    ///     Matrix::from([[9.0, 10.0, 11.0]]),
    /// );
    ///
    /// let mut layer = SynapseLayer::<f64, ReLU>::zero(2, 3);
    /// layer.apply_gradient(&gradient, learning_rate);
    ///
    /// assert_eq!(
    ///     layer.weights(),
    ///     &Matrix::from([[-4.0, -6.0, -8.0], [-10.0, -12.0, -14.0]])
    /// );
    /// assert_eq!(layer.biases(), &Matrix::from([[-16.0, -18.0, -20.0]]));
    ///
    /// gradient.add(
    ///     Matrix::from([[-4.0, -6.0, -8.0], [-10.0, -12.0, -14.0]]),
    ///     Matrix::from([[-16.0, -18.0, -20.0]]),
    /// );
    ///
    /// let mut layer = SynapseLayer::<f64, ReLU>::zero(2, 3);
    /// layer.apply_gradient(&gradient, learning_rate);
    ///
    /// assert_eq!(
    ///     layer.weights(),
    ///     &Matrix::from([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
    /// );
    /// assert_eq!(layer.biases(), &Matrix::from([[0.0, 0.0, 0.0]]));
    /// ```
    pub fn apply_gradient(&mut self, layer_gradient: &LayerGradient<T>, learning_rate: T)
    where
        T: Div<T, Output = T> + SubAssign<T> + Mul<T, Output = T> + From<i32>,
    {
        layer_gradient.expect_dimensions(self.previous_size(), self.size(), "layer_gradient");

        layer_gradient.apply(&mut self.weights, &mut self.biases, learning_rate);
    }

    /// Performs backward propagation through the synapse layer.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use smyl::network::LayerGradient;
    /// use smyl::prelude::*;
    ///
    /// let weights = Matrix::from([[1.0, 0.3, 0.6], [0.5, 0.0, 0.2]]);
    /// let biases = Matrix::from([[0.2, 0.8, 0.5]]);
    ///
    /// let mut layer = SynapseLayer::<f64, ReLU>::new(weights, biases);
    /// let inputs = Matrix::from([[1.0, 2.0]]);
    /// let outputs = layer.forward(inputs.clone());
    /// assert_eq!(outputs, Matrix::from([[2.2, 1.1, 1.5]]));
    ///
    /// let mut layer_gradient = LayerGradient::zero(2, 3);
    /// let input_gradient = layer.backward(
    ///     &inputs,
    ///     outputs,
    ///     Matrix::from([[-0.4, 0.5, -0.8]]),
    ///     &mut layer_gradient,
    /// );
    /// let mut weights = Matrix::zero(2, 3);
    /// let mut biases = Matrix::zero(1, 3);
    /// layer_gradient.apply(&mut weights, &mut biases, 2.0);
    /// assert_eq!(weights, Matrix::from([[0.8, -1.0, 1.6], [1.6, -2.0, 3.2]]));
    /// assert_eq!(biases, Matrix::from([[0.8, -1.0, 1.6]]));
    /// assert_eq!(
    ///     input_gradient,
    ///     Matrix::from([[-0.73, -0.36000000000000004]])
    /// );
    /// ```
    pub fn backward(
        &self,
        input: &Matrix<T>,
        output: Matrix<T>,
        output_gradient: Matrix<T>,
        layer_gradient: &mut LayerGradient<T>,
    ) -> Matrix<T>
    where
        T: Add<T, Output = T>
            + AddAssign<T>
            + Sub<T, Output = T>
            + Mul<T, Output = T>
            + std::iter::Sum,
    {
        input.expect_dimensions(1, self.previous_size(), "input");
        output.expect_dimensions(1, self.size(), "output");
        output_gradient.expect_dimensions(1, self.size(), "output_gradient");
        layer_gradient.expect_dimensions(self.previous_size(), self.size(), "layer_gradient");

        let activation_derivative = output.activate_derivative::<A>();
        let pre_activation_gradient = output_gradient.hadamard(&activation_derivative);

        let mut weight_gradient = Matrix::zero(self.previous_size(), self.size());

        for j in 0..self.previous_size() {
            for i in 0..self.size() {
                weight_gradient.set(j, i, input.get(0, j) * pre_activation_gradient.get(0, i));
            }
        }

        let bias_gradient = pre_activation_gradient.clone();

        layer_gradient.add(weight_gradient, bias_gradient);

        pre_activation_gradient.dot(&self.weights.clone().transpose())
    }
}

#[cfg(feature = "serde")]
impl<T, A> serde::Serialize for SynapseLayer<T, A>
where
    T: Default + Clone + Copy + Debug + serde::Serialize,
    A: Function<T>,
{
    /// Serializes a `SynapseLayer` using the Serde library.
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use serde::ser::SerializeStruct;

        let mut layer = serializer.serialize_struct("SynapseLayer", 3)?;
        layer.serialize_field("weights", &self.weights)?;
        layer.serialize_field("biases", &self.biases)?;
        layer.serialize_field("activation", std::any::type_name::<A>())?;
        layer.end()
    }
}

#[cfg(feature = "random")]
impl<T, A> rand::Fill for SynapseLayer<T, A>
where
    T: Debug + Default + Copy,
    A: Function<T>,
    rand::distributions::Standard: rand::distributions::Distribution<T>,
{
    fn try_fill<R: rand::Rng + ?Sized>(&mut self, rng: &mut R) -> Result<(), rand::Error> {
        rng.try_fill(&mut self.weights)?;
        rng.try_fill(&mut self.biases)
    }
}
