neuralib 0.0.3

use crate::activation::Activation;
use crate::layer::Layer;
use crate::neuron::Neuron;
use crate::training::DataValue;

#[cfg(feature = "serde")]
use {
    rmp_serde::{Deserializer, Serializer},
    serde::{Deserialize, Serialize},
};

/// A neural network
#[derive(Debug, Serialize, Deserialize)]
pub struct NeuralNetwork {
    layers: Vec<Layer>,
    layer_count: usize,
    input_size: usize,
    output_size: usize,
}

impl NeuralNetwork {
    /// Create a new neural network
    ///
    /// Arguments:
    ///
    /// * `layer_sizes` - A slice of usizes containing the size of each layer in the neural network
    /// * `activation_functions` - A Vec of which activation function should be in each layer
    pub fn new(
        layer_sizes: &[usize],
        activation_functions: Vec<Activation>,
    ) -> crate::error::Result<NeuralNetwork> {
        if layer_sizes.len() < 2 {
            return Err(crate::error::NoLayersError {}.into());
        }

        let input_size = layer_sizes[0];
        let layer_sizes = &layer_sizes[1..];

        if layer_sizes.len() != activation_functions.len() {
            // TODO: Make custom error for this
            return Err(crate::error::NoLayersError {}.into());
        }

        // Allocate a vector for the layers
        let mut layers: Vec<Layer> = Vec::with_capacity(layer_sizes.len());

        let mut output_size = 0;

        let mut previous_size = &input_size;
        for (layer_size, activator) in layer_sizes.iter().zip(activation_functions) {
            layers.push(Layer::new(*previous_size, *layer_size, activator));
            previous_size = layer_size;
            output_size = *layer_size;
        }

        Ok(NeuralNetwork {
            layer_count: layers.len(),
            layers,
            input_size,
            output_size,
        })
    }

    /// Run the neural network with specific inputs
    ///
    /// Arguments:
    ///
    /// * `inputs` - A slice of f64s to be used as input to the network
    pub fn activate(&mut self, inputs: &[f64]) -> crate::error::Result<Vec<f64>> {
        if inputs.len() != self.input_size {
            return Err(crate::error::InputSizeError {
                inputted: inputs.len(),
                expected: self.input_size,
                chain_depth: "NeuralNetwork".to_owned(),
            }
            .into());
        }

        // We have to feed each layer's output into the next layer's input

        let mut next_in = inputs.to_vec();

        for layer in &mut self.layers {
            // All the sizes *should* be correct
            next_in = layer
                .activate(&next_in)
                .expect("Length was already checked. This should not fail. (Network)")
        }

        Ok(next_in)
    }

    /// Get the number of layers in this neural network
    pub fn get_layer_count(&self) -> usize {
        self.layer_count
    }

    #[allow(dead_code)]
    fn get_layer(&self, idx: usize) -> Option<&Layer> {
        self.layers.get(idx)
    }

    fn get_layer_mut(&mut self, idx: usize) -> Option<&mut Layer> {
        self.layers.get_mut(idx)
    }

    /// Calculate the loss of the network with a DataValue
    ///
    /// Arguments:
    ///
    /// * `value` - A reference to a DataValue
    pub fn loss_with_value(&mut self, value: &DataValue) -> crate::error::Result<f64> {
        if value.expected_output.len() != self.output_size {
            return Err(crate::error::InputSizeError {
                inputted: value.expected_output.len(),
                expected: self.output_size,
                chain_depth: "NeuralNetwork".to_owned(),
            }
            .into());
        }

        let output = self.activate(&value.input)?;

        let mut loss = 0.0;

        for (actual, expected) in output.iter().zip(value.expected_output.iter()) {
            loss += Neuron::loss(actual, expected);
        }

        Ok(loss)
    }

    /// Calculate the average loss for a slice of DataValues.
    /// This method should be preferred over `loss_with_value`
    ///
    /// Arguments:
    ///
    /// * `values` - A slice of DataValues to test
    pub fn loss(&mut self, values: &[DataValue]) -> crate::error::Result<f64> {
        let mut total_loss = 0.0;

        let value_length = values.len();

        for value in values {
            total_loss += self.loss_with_value(value)?;
        }

        Ok(total_loss / (value_length as f64))
    }

    fn apply_gradients(&mut self, learn_rate: f64) {
        for layeridx in 0..self.get_layer_count() {
            let layer = self.get_layer_mut(layeridx).unwrap();
            for neuronidx in 0..layer.get_neuron_count() {
                let neuron = layer.get_neuron_mut(neuronidx).unwrap();
                neuron.apply_gradients(learn_rate);
            }
        }
    }

    /// Train the network on some data
    ///
    /// Arguments:
    ///
    /// * `training_data` - The data to train the network on in a slice of DataValues
    /// * `learn_rate` - How fast the network should try to learn
    pub fn learn(
        &mut self,
        training_data: &[DataValue],
        learn_rate: f64,
    ) -> crate::error::Result<()> {
        for value in training_data {
            self.update_all_gradients(value)?;
        }

        self.apply_gradients(learn_rate / (training_data.len() as f64));

        Ok(())
    }

    pub fn learn_randomly(
        &mut self,
        training_data: &[DataValue],
        learn_rate: f64,
        amount: usize,
    ) -> crate::error::Result<()> {
        use rand::seq::SliceRandom;
        let mut rand_split = training_data.to_vec();

        // Shuffle the data
        rand_split.shuffle(&mut rand::rng());
        // Get the split
        self.learn(&rand_split[..amount], learn_rate)
    }

    fn update_all_gradients(&mut self, value: &DataValue) -> crate::error::Result<()> {
        if value.expected_output.len() != self.output_size {
            return Err(crate::error::InputSizeError {
                inputted: value.expected_output.len(),
                expected: self.output_size,
                chain_depth: "NeuralNetwork".to_owned(),
            }
            .into());
        }

        // Prep the network
        self.activate(&value.input)?;

        let output_layer = self
            .get_layer_mut(self.get_layer_count() - 1)
            .expect("Length was already checked. This should not fail. (Network)");
        output_layer.update_gradients_output(&value.expected_output);

        for layeridx in (0..self.get_layer_count()).rev().skip(1) {
            // Fun borrow checker shenanigans
            let (up_to_current, past_current) = self
                .layers
                .split_at_mut_checked(layeridx + 1)
                .expect("Length was already checked. This should not fail. (Network)");
            let current_layer = up_to_current
                .get_mut(layeridx)
                .expect("Length was already checked. This should not fail. (Network)");
            let next_layer = past_current
                .first()
                .expect("Length was already checked. This should not fail. (Network)");
            current_layer.update_gradients_hidden(next_layer);
        }

        Ok(())
    }

    #[cfg(feature = "serde")]
    pub fn save(&self, file: &mut impl std::io::Write) -> std::io::Result<()> {
        let mut buf = Vec::new();
        self.serialize(&mut Serializer::new(&mut buf)).unwrap();

        file.write_all(&buf)
    }

    #[cfg(feature = "serde")]
    pub fn from_saved(file: impl std::io::Read) -> Result<Self, rmp_serde::decode::Error> {
        //let mut buf = Vec::new();

        //file.read_to_end(&mut buf).unwrap();

        //Self::deserialize(&mut Deserializer::new(&buf[..]))
        Self::deserialize(&mut Deserializer::new(file))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn methods() {
        let mut network =
            NeuralNetwork::new(&[2, 2, 2], vec![Activation::Sigmoid, Activation::Step]).unwrap();

        network.activate(&[0.0, 0.0]).unwrap();

        assert_eq!(network.get_layer_count(), 2);
    }

    #[test]
    fn errors() {
        assert!(NeuralNetwork::new(&[], vec![]).is_err());
        assert!(NeuralNetwork::new(&[0], vec![]).is_err());
        assert!(NeuralNetwork::new(&[0, 0], vec![]).is_err());
        assert!(NeuralNetwork::new(&[0, 1], vec![]).is_err());
        assert!(NeuralNetwork::new(&[], vec![Activation::Linear]).is_err());
        assert!(NeuralNetwork::new(&[0], vec![Activation::Linear]).is_err());

        let mut network = NeuralNetwork::new(&[1, 1], vec![Activation::Linear]).unwrap();
        assert!(network.activate(&[]).is_err());
    }

    #[test]
    fn train_works() {
        let mut network =
            NeuralNetwork::new(&[2, 2, 1], vec![Activation::Sigmoid, Activation::Sigmoid]).unwrap();

        let data: Vec<DataValue> = vec![
            DataValue {
                input: vec![0.0, 0.0],
                expected_output: vec![0.0],
            },
            DataValue {
                input: vec![1.0, 0.0],
                expected_output: vec![1.0],
            },
            DataValue {
                input: vec![0.0, 1.0],
                expected_output: vec![1.0],
            },
            DataValue {
                input: vec![1.0, 1.0],
                expected_output: vec![0.0],
            },
        ];

        let starting_loss = network.loss(&data).unwrap();

        let _ = network.learn(&data, 0.5);

        let new_loss = network.loss(&data).unwrap();

        // Make sure the network can actually learn
        assert!(starting_loss > new_loss);
    }

    #[test]
    fn xor_train() {
        // Try 3 times bc it could fail if a bad network is made
        for _ in 0..3 {
            let mut network =
                NeuralNetwork::new(&[2, 2, 1], vec![Activation::Sigmoid, Activation::Sigmoid])
                    .unwrap();

            let data: Vec<DataValue> = vec![
                DataValue {
                    input: vec![0.0, 0.0],
                    expected_output: vec![0.0],
                },
                DataValue {
                    input: vec![1.0, 0.0],
                    expected_output: vec![1.0],
                },
                DataValue {
                    input: vec![0.0, 1.0],
                    expected_output: vec![1.0],
                },
                DataValue {
                    input: vec![1.0, 1.0],
                    expected_output: vec![0.0],
                },
            ];

            for _ in 0..50000 {
                let _ = network.learn(&data, 0.5);
            }

            // Assert that it learned the general theme
            // Could fail randomly...
            if (network.activate(&[0.0, 0.0]).unwrap()[0].abs() < 0.25)
                && ((1.0 - network.activate(&[1.0, 0.0]).unwrap()[0]).abs() < 0.25)
                && ((1.0 - network.activate(&[0.0, 1.0]).unwrap()[0]).abs() < 0.25)
                && (network.activate(&[1.0, 1.0]).unwrap()[0].abs() < 0.25)
            {
                return;
            }
        }
        panic!("Could not train a working network after 3 tries")
    }
}