/*
    Appellation: controller <module>
    Contrib: @FL03
*/
use super::NeuralFeatures;
use crate::surface::Params;
use cnc::{ReLU, Sigmoid};
use ndarray::{Array1, Array2, Ix2, ScalarOperand};
use num_traits::Float;

/// A controller for the neural Turing machine
#[derive(Clone, Debug)]
pub struct NeuralController<T = f64> {
    pub(crate) features: NeuralFeatures,
    pub(crate) attention: Array2<T>,
    pub(crate) hidden: Params<T, Ix2>,
    pub(crate) input: Params<T, Ix2>,
    pub(crate) output: Params<T, Ix2>,
}

impl<T> NeuralController<T> {
    pub fn new() -> Self
    where
        T: Clone + Default,
    {
        Self::from_features(NeuralFeatures::default())
    }
    /// create a new controller from the given features
    pub fn from_features(features: NeuralFeatures) -> Self
    where
        T: Clone + Default,
    {
        Self {
            features,
            attention: Array2::default(features.dim_attention()),
            hidden: Params::default(features.dim_hidden()),
            input: Params::default(features.dim_input()),
            output: Params::default(features.dim_output()),
        }
    }
    /// returns a copy of the features of the controller
    pub const fn features(&self) -> NeuralFeatures {
        self.features
    }
    /// returns a mutable reference to the features of the controller
    pub fn features_mut(&mut self) -> &mut NeuralFeatures {
        &mut self.features
    }
    /// returns an immutable reference to the attention weights
    pub const fn attention(&self) -> &Array2<T> {
        &self.attention
    }
    /// returns an mutable reference to the attention weights
    pub fn attention_mut(&mut self) -> &mut Array2<T> {
        &mut self.attention
    }
    /// returns an immutable reference to the hidden layer(s)
    pub const fn hidden(&self) -> &Params<T, Ix2> {
        &self.hidden
    }
    /// returns a mutable reference to the hidden layer(s)
    pub fn hidden_mut(&mut self) -> &mut Params<T, Ix2> {
        &mut self.hidden
    }
    /// returns an immutable reference to the hidden weights
    pub const fn hidden_weights(&self) -> &Array2<T> {
        self.hidden().weights()
    }
    /// returns an mutable reference to the hidden weights
    pub fn hidden_weights_mut(&mut self) -> &mut Array2<T> {
        self.hidden_mut().weights_mut()
    }
    /// returns an immutable reference to the input layer
    pub const fn input(&self) -> &Params<T, Ix2> {
        &self.input
    }
    /// returns a mutable reference to the input layer
    pub fn input_mut(&mut self) -> &mut Params<T, Ix2> {
        &mut self.input
    }
    /// returns an immutable reference to the input bias
    pub const fn input_bias(&self) -> &Array1<T> {
        self.input().bias()
    }
    /// returns a mutable reference to the input bias
    pub fn input_bias_mut(&mut self) -> &mut Array1<T> {
        self.input_mut().bias_mut()
    }
    /// returns an immutable reference to the input weights
    pub const fn input_weights(&self) -> &Array2<T> {
        self.input().weights()
    }
    /// returns an mutable reference to the input weights
    pub fn input_weights_mut(&mut self) -> &mut Array2<T> {
        self.input_mut().weights_mut()
    }
    /// returns an immutable reference to the output layer
    pub const fn output(&self) -> &Params<T, Ix2> {
        &self.output
    }
    /// returns a mutable reference to the output layer
    pub fn output_mut(&mut self) -> &mut Params<T, Ix2> {
        &mut self.output
    }
    /// returns an immutable reference to the output bias
    pub const fn output_bias(&self) -> &Array1<T> {
        self.output().bias()
    }
    /// returns a mutable reference to the output bias
    pub fn output_bias_mut(&mut self) -> &mut Array1<T> {
        self.output_mut().bias_mut()
    }
    /// returns an immutable reference to the output weights
    pub const fn output_weights(&self) -> &Array2<T> {
        self.output().weights()
    }
    /// returns an mutable reference to the output weights
    pub fn output_weights_mut(&mut self) -> &mut Array2<T> {
        self.output_mut().weights_mut()
    }
    /// consumes the controller to initialize the various parameters with random values
    #[cfg(feature = "rand")]
    pub fn init(self) -> Self
    where
        T: Float + num_traits::FromPrimitive + rand_distr::uniform::SampleUniform,
        rand_distr::StandardNormal: rand_distr::Distribution<T>,
    {
        use crate::surface::Initialize;

        // initialize the hidden layer(s)
        let hidden = Params::glorot_normal(self.dim_hidden());
        // Controller input weights (alphabet + state -> hidden)
        let input = Params::glorot_normal(self.dim_input());
        // initialize the output layers
        let output = Params::glorot_normal(self.dim_output());
        // return a new instance with the initialized layers
        Self {
            attention: Array2::glorot_normal(self.dim_attention()),
            hidden,
            input,
            output,
            ..self
        }
    }
}

impl<T> NeuralController<T>
where
    T: ScalarOperand,
{
    /// forward input through the controller network
    pub fn forward(&self, input: &Array1<T>) -> cnc::Result<Array1<T>>
    where
        T: Float + ScalarOperand,
    {
        // forward the input through the input layer; activate using relu
        let mut output = self.input.forward(input)?.relu();
        // forward the input through the hidden layer(s); activate using relu
        output = self.hidden.forward(&output)?.relu();
        // forward the input through the output layer; activate using sigmoid
        output = self.output.forward(&output)?.sigmoid();

        Ok(output)
    }
}

impl<T> Default for NeuralController<T>
where
    T: Clone + Default,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T> core::ops::Deref for NeuralController<T> {
    type Target = NeuralFeatures;

    fn deref(&self) -> &Self::Target {
        &self.features
    }
}
