/*
    Appellation: features <module>
    Contrib: @FL03
*/

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[cfg_attr(
    feature = "serde",
    derive(serde_derive::Deserialize, serde_derive::Serialize)
)]
pub struct NeuralFeatures {
    pub(crate) attention: usize,
    // the number of features of the input layer
    pub(crate) inputs: usize,
    // the number of features of the hidden layer
    pub(crate) hidden: usize,
    // the number of hidden layers
    pub(crate) layers: usize,
    // the number of features of the read/write heads
    pub(crate) memory: usize,
    // the number of features of the output layer
    pub(crate) outputs: usize,
}

impl NeuralFeatures {
    pub fn new(
        attention: usize,
        memory: usize,
        inputs: usize,
        outputs: usize,
        hidden: usize,
        layers: usize,
    ) -> Self {
        Self {
            attention,
            memory,
            inputs,
            hidden,
            layers,
            outputs,
        }
    }
    /// returns a copy of the attention size
    pub const fn attention(&self) -> usize {
        self.attention
    }
    /// returns a copy of the number of features within the read/write heads
    pub const fn codex(&self) -> usize {
        self.memory
    }
    /// returns a copy of the hidden size
    pub const fn hidden(&self) -> usize {
        self.hidden
    }
    /// returns a copy of the input size
    pub const fn inputs(&self) -> usize {
        self.inputs
    }
    /// returns a copy of the number of hidden layers
    pub const fn layers(&self) -> usize {
        self.layers
    }
    /// returns a copy of the output size
    pub const fn outputs(&self) -> usize {
        self.outputs
    }
    /// the shape of the attention layer (n_outputs, n_attention)
    pub const fn dim_attention(&self) -> (usize, usize) {
        // TODO: Consider if this should be (self.attention, self.hidden)
        (self.outputs, self.attention)
    }
    /// the shape of the input layer (inputs, n_hidden)
    pub const fn dim_input(&self) -> (usize, usize) {
        (self.inputs, self.hidden)
    }
    /// the shape of the hidden layer (hidden, hidden)
    pub const fn dim_hidden(&self) -> (usize, usize) {
        (self.hidden, self.hidden)
    }
    /// the shape of the memory matrix (n_memory, n_attention)
    pub const fn dim_memory(&self) -> (usize, usize) {
        (self.memory, self.attention)
    }
    /// the shape of the output layer (n_hidden, n_out)
    pub const fn dim_output(&self) -> (usize, usize) {
        (self.hidden, self.outputs)
    }
    /// set the number of features for attention
    pub fn set_attention(&mut self, attention: usize) {
        self.attention = attention;
    }
    /// set the number of features for the read/write heads
    pub fn set_codex(&mut self, codex: usize) {
        self.memory = codex;
    }
    /// set the number of features for then input layer
    pub fn set_inputs(&mut self, inputs: usize) {
        self.inputs = inputs;
    }
    /// set the number of features for the hidden layer
    pub fn set_hidden(&mut self, hidden: usize) {
        self.hidden = hidden;
    }
    /// set the number of features for the output layer
    pub fn set_outputs(&mut self, outputs: usize) {
        self.outputs = outputs;
    }
    /// consumes the current instance; returns another with the given number of features for
    /// attention
    pub fn with_attention(self, attention: usize) -> Self {
        Self { attention, ..self }
    }
    /// consumes the current instance; returns another with the given number of features for the
    /// read/write heads
    pub fn with_codex(self, codex: usize) -> Self {
        Self {
            memory: codex,
            ..self
        }
    }
    /// consumes the current instance; returns another with the given number of input features
    pub fn with_inputs(self, inputs: usize) -> Self {
        Self { inputs, ..self }
    }
    /// consumes the current instance; returns another with the given number of hidden features
    pub fn with_hidden(self, hidden: usize) -> Self {
        Self { hidden, ..self }
    }
    /// consumes the current instance; returns another with the given number of output features
    pub fn with_outputs(self, outputs: usize) -> Self {
        Self { outputs, ..self }
    }
}

impl Default for NeuralFeatures {
    fn default() -> Self {
        Self {
            attention: 20,
            memory: 128,
            inputs: 5, // (alphabet (3) + states (2))
            hidden: 50,
            layers: 1,
            outputs: 10,
        }
    }
}
