use crate::float::Float;
use crate::linear_algebra::Vector;

use crate::transforms;
use crate::transforms::{ReluTransform, SoftmaxTransform};

#[derive(Debug, Clone)]
pub struct NNPoint {
    label: usize,
    normalized_image: Vec<Float>,
}

impl NNPoint {
    pub fn new(label: u8, normalized_image: Vec<Float>) -> Self {
        Self {
            label: label as usize,
            normalized_image,
        }
    }
}

#[derive(Debug)]
pub struct NeuralNetworkInTraining {
    neurons_per_layer: Vec<usize>,

    inputs: Vec<Vector>, // Each input will start with 1
    output: Vector,
    input_gradients: Vec<Vector>,

    transforms: Vec<ReluTransform>,
    output_transform: SoftmaxTransform,
}

#[derive(Debug, Copy, Clone)]
pub struct NeuralNetworkParameters {
    pub epochs: usize,
    pub batch_size: usize,
    pub learning_rate: Float,
}

impl NeuralNetworkParameters {
    fn show(&self) -> String {
        format!(
            "epoch = {}, batch = {}, rate = {}",
            self.epochs, self.batch_size, self.learning_rate
        )
    }
}

impl NeuralNetworkInTraining {
    pub fn new(neurons_per_layer: Vec<usize>) -> Self {
        // e.g. neurons_per_layer = [699, 79, 49, 19, 14, 10]

        // By a layer here we mean a collection of neurons that's between two neighbouring Transforms or the initial input or final input neurons.
        // Note that there are N + 1 layers where N is the number of Transforms used.
        if neurons_per_layer.len() < 2 {
            todo!()
        }

        let neurons_per_layer_except_last: Vec<usize> = neurons_per_layer
            .clone()
            .into_iter()
            .rev()
            .skip(1)
            .rev()
            .collect();
        let neurons_per_layer_except_first: Vec<usize> =
            neurons_per_layer.clone().into_iter().skip(1).collect();

        // Convention: the first component should always be 1.0 - this allows the first
        // column of the weight matrix to be interpreted as bias.
        let inputs: Vec<Vector> = neurons_per_layer_except_last
            .iter()
            .map(|neuron_count| {
                let mut v = Vector::zero(neuron_count + 1);
                v[0] = 1.0;
                v
            })
            .collect();

        let input_gradients: Vec<Vector> = neurons_per_layer_except_first
            .iter()
            .map(|neuron_count| Vector::zero(*neuron_count))
            .collect();

        let transforms: Vec<ReluTransform> = neurons_per_layer_except_last
            .iter()
            .zip(neurons_per_layer_except_last.iter().skip(1))
            .map(|(input_neuron_count, output_neuron_count)| {
                ReluTransform::new(*input_neuron_count + 1, *output_neuron_count)
            })
            .collect();

        let neurons_in_last_layer = *neurons_per_layer.last().unwrap();
        let neurons_in_next_to_last_layer = *neurons_per_layer_except_last.last().unwrap();

        let output = Vector::zero(neurons_in_last_layer);
        let output_transform =
            SoftmaxTransform::new(neurons_in_next_to_last_layer + 1, neurons_in_last_layer);

        Self {
            neurons_per_layer,

            inputs,
            output,
            input_gradients,

            transforms,
            output_transform,
        }
    }

    // You need to initialize inputs[0] before use.
    fn output_mut(&mut self) {
        // The following iterates over pairs of neighbouring inputs where we
        // have a mutable reference to both of them.
        // With Rust borrow-checking rules this can't be done directly,
        // so I had to resort to `split_at_mut(k)`
        for i in 0..self.inputs.len() - 1 {
            let (left_inputs, right_inputs) = self.inputs.split_at_mut(i + 1);

            let transform = &mut self.transforms[i];
            let input = &left_inputs[i][..];
            let output = &mut right_inputs[0][1..];
            transform.output_mut(input, output);
        }
        self.output_transform.output_mut(
            &self.inputs[self.inputs.len() - 1][..],
            &mut self.output[..],
        );
    }

    // Initialize input_gradients[-1] with error gradient before use.
    fn update_weights_mut(&mut self) {
        let last_index = self.input_gradients.len() - 1;

        // TODO: Last layer is no different. This should be part of the same loop
        {
            let transform = &mut self.output_transform;

            transform.potential_gradient_mut(&self.input_gradients[last_index][..]);
            transform
                .gradient_with_respect_to_input_mut(&mut self.input_gradients[last_index - 1][..]);
            transform.add_gradient_with_respect_to_weights_mut(&self.inputs[last_index][..]);
        }

        for i in (0..self.input_gradients.len() - 2).rev() {
            let (left_input_gradient, right_input_gradient) =
                self.input_gradients.split_at_mut(i + 1);

            let transform = &mut self.transforms[i + 1];
            let left_grad = &mut left_input_gradient[i][..];
            let right_grad = &right_input_gradient[0][..];
            let input = &self.inputs[i + 1][..];

            transform.potential_gradient_mut(right_grad);
            transform.gradient_with_respect_to_input_mut(left_grad);
            transform.add_gradient_with_respect_to_weights_mut(input);
        }

        {
            let transform = &mut self.transforms[0];
            transform.potential_gradient_mut(&self.input_gradients[0][..]);
            transform.add_gradient_with_respect_to_weights_mut(&self.inputs[0][..]);
            // Note that we are not computing gradient with respect to input, since this is the
            // first layer and we don't care about changes to input, only to weights.
        }
    }

    fn forward_and_backwards_mut(&mut self, point: &NNPoint) {
        self.inputs[0][1..].copy_from_slice(&point.normalized_image[..]);
        self.output_mut();

        let last_index = self.input_gradients.len() - 1;
        // WARNING: Use proper error function
        // transforms::gradient_error_mut(&self.output[..], desired_output, &mut self.input_gradients[last_index][..]);
        transforms::cross_entropy_derivative_simple(
            &self.output[..],
            point.label,
            &mut self.input_gradients[last_index][..],
        );
        self.update_weights_mut()
    }

    fn iterate_over_batch_mut(&mut self, learning_rate: Float, batch: &[NNPoint]) {
        // iterates over the batch, while updating gradient of weights.
        for point in batch {
            self.forward_and_backwards_mut(point);
        }

        // Update the current weights by the opposite of the epsilon / batch_size *weight_gradients
        let batch_size = batch.len() as Float;
        let epsilon = -learning_rate / batch_size;

        for transform in &mut self.transforms {
            transform
                .weight
                .add_to_self_scaled_mut(epsilon, &transform.weight_gradient);
            // Resets the weight gradient to zero so it can be used in next batch.
            transform.weight_gradient.zero_mut();
        }
    }

    fn iterate_over_epoch(
        &mut self,
        training_set: &mut [NNPoint],
        batch_size: usize,
        learning_rate: Float,
    ) {
        use rand::seq::SliceRandom;
        use rand::thread_rng;
        training_set.shuffle(&mut thread_rng()); // Shuffling is linear in the size of the slice

        for batch in training_set.chunks(batch_size) {
            self.iterate_over_batch_mut(learning_rate, batch);
        }
    }

    pub fn train(
        &mut self,
        parameters: NeuralNetworkParameters,
        training_set: &mut [NNPoint],
        testing_set: Option<&[NNPoint]>,
    ) {
        fn test(
            nn: &mut NeuralNetworkInTraining,
            parameters: NeuralNetworkParameters,
            testing_set: Option<&[NNPoint]>,
            accuracy_per_epoch: &mut Vec<f32>,
        ) {
            if let Some(testing_set) = testing_set {
                let accuracy = nn.accuracy(testing_set);
                accuracy_per_epoch.push(accuracy);
                println!();
                println!("{}", parameters.show());
                println!("{:?}", nn.neurons_per_layer);
                println!("{:?}", accuracy_per_epoch);
            }
        }

        fn next_learning_rate(initial_learning_rate: Float, epoch: usize) -> Float {
            // initial_learning_rate / (1.0 + epoch as Float / 30.0)
            // initial_learning_rate * (0.1 as Float).powf(epoch as Float / 50.0)
            initial_learning_rate * (0.1 as Float).powf(epoch as Float / 20.0)
        }

        use std::time::Instant;
        let now = Instant::now();

        let number_of_epochs = parameters.epochs;
        let batch_size = parameters.batch_size;
        let mut learning_rate = parameters.learning_rate;
        let mut accuracy_per_epoch = Vec::with_capacity(number_of_epochs + 1);

        test(self, parameters, testing_set, &mut accuracy_per_epoch);
        for epoch in 0..number_of_epochs {
            println!("Epoch {}/{}", epoch + 1, number_of_epochs);
            println!("Current learning rate = {}", learning_rate);
            self.iterate_over_epoch(training_set, batch_size, learning_rate);
            test(self, parameters, testing_set, &mut accuracy_per_epoch);

            learning_rate = next_learning_rate(parameters.learning_rate, epoch);

            let elapsed = now.elapsed();
            let total_seconds = elapsed.as_secs();
            let minutes = total_seconds / 60;
            let seconds = total_seconds % 60;
            println!("Duration of training: {} min {} sec", minutes, seconds);
        }
    }

    pub fn show_accuracy_on(&mut self, testing_set: &[NNPoint], dataset_name: &str) {
        println!(
            "{} dataset accuracy: {:?}",
            dataset_name,
            self.accuracy(testing_set)
        );
    }

    pub fn output_label(&self) -> usize {
        let mut state: Option<(usize, Float)> = None;
        for (i, y) in self.output.iter().enumerate() {
            match state {
                Some((_, max_so_far)) => {
                    if *y > max_so_far {
                        state = Some((i, *y))
                    }
                }
                None => state = Some((i, *y)),
            }
        }
        match state {
            Some((label, _)) => label,
            None => {
                todo!()
            }
        }
    }

    pub fn accuracy(&mut self, testing_set: &[NNPoint]) -> f32 {
        let mut num_of_correct_classifications = 0;
        for point in testing_set {
            self.inputs[0][1..].copy_from_slice(&point.normalized_image[..]);
            self.output_mut(); // This doesn't change the weights

            let neural_network_produced_label = self.output_label();
            if point.label == neural_network_produced_label {
                num_of_correct_classifications += 1
            }
            // println!("max-label: {}, desired-label: {}, prob-distr: {:?}", neural_network_produced_label, point.label, self.output);
        }
        num_of_correct_classifications as f32 / testing_set.len() as f32
    }

    pub fn test(&mut self, dataset: &[NNPoint]) -> Vec<usize> {
        dataset
            .iter()
            .map(|point| {
                self.inputs[0][1..].copy_from_slice(&point.normalized_image[..]);
                self.output_mut();
                self.output_label()
            })
            .collect()
    }

    pub fn test_distributions(&mut self, dataset: &[NNPoint]) -> Vec<Vec<Float>> {
        dataset
            .iter()
            .map(|point| {
                self.inputs[0][1..].copy_from_slice(&point.normalized_image[..]);
                self.output_mut();
                self.output.clone().to_vec()
            })
            .collect()
    }
}
