pub fn activate(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

pub fn deriv_activate(x: f64) -> f64 {
    let a = activate(x);
    a * (1.0 - a)
}

pub struct NeuralNetwork {
    layer_count: usize,
    layers: Vec<Layer>,
    layer_data: Vec<LayerData>,
}

impl NeuralNetwork {
    pub fn new(
        input: usize,
        output: usize,
        hidden_layer_sizes: &[usize],
        rng: &mut impl Rng,
    ) -> Self {
        let mut layers = Vec::with_capacity(hidden_layer_sizes.len());
        let mut layer_data = Vec::with_capacity(hidden_layer_sizes.len());
        let mut prev_size = input;

        let mut i = 1;

        for &size in hidden_layer_sizes {
            layers.push(Layer::new(prev_size, size, rng, i));
            layer_data.push(LayerData::new(prev_size, size, i));
            prev_size = size;
            i += 1;
        }

        layers.push(Layer::new(prev_size, output, rng, i));
        layer_data.push(LayerData::new(prev_size, output, i));

        NeuralNetwork {
            layer_count: layers.len(),
            layers,
            layer_data,
        }
    }
}

pub struct Layer {
    index: usize,
    inputs: usize,
    outputs: usize,
    weights: Vec<f64>,
    biases: Vec<f64>,
}

impl Layer {
    pub fn new(inputs: usize, outputs: usize, rng: &mut impl Rng, index: usize) -> Self {
        let distr = Uniform::new(0., 1.);
        Layer {
            index,
            inputs,
            outputs,
            weights: (0..inputs * outputs)
                .into_iter()
                .map(|_| rng.sample(distr))
                .collect(),
            biases: vec![0.0; outputs],
        }
    }

    pub fn think(&self, inputs: &[f64], output: &mut LayerData) {
        output
            .weighted_inputs
            .par_iter_mut()
            .zip(output.activations.par_iter_mut())
            .enumerate()
            .for_each(|(output_index, (weighted_input, activation))| {
                *weighted_input = self.biases[output_index];
                self.weight_slice(output_index).iter().for_each(|weight| {
                    *weighted_input += weight * inputs[output_index];
                });
                *activation = activate(*weighted_input);
            });
    }

    pub fn weight(&self, input: usize, output: usize) -> f64 {
        self.weights[output * self.inputs + input]
    }

    pub fn weight_slice(&self, output: usize) -> &[f64] {
        &self.weights[output * self.inputs..(output + 1) * self.inputs]
    }
}

pub struct LayerData {
    index: usize,
    /// z_n = a_{n-1} * w_n + b_n
    weighted_inputs: Vec<f64>,
    /// a_n = Activation(z_n)
    activations: Vec<f64>,
    /// ∂c/∂a_n * ∂a_n/∂z_n
    backprop: Vec<f64>,
    grad_w: Vec<f64>,
    grad_b: Vec<f64>,
}

impl LayerData {
    pub fn new(inputs: usize, outputs: usize, index: usize) -> Self {
        LayerData {
            index,
            weighted_inputs: vec![0.; outputs],
            activations: vec![0.; outputs],
            backprop: vec![0.; outputs],
            grad_w: vec![0.; outputs * inputs],
            grad_b: vec![0.; outputs],
        }
    }

    pub fn flatten_index(&self, input: usize, output: usize) -> usize {
        output * self.activations.len() + input
    }
}

pub struct DataPoint {
    pub inputs: Vec<f64>,
    pub targets: Vec<f64>,
}

fn deriv_cross_entropy(activations: f64, targets: f64) -> f64 {
    (-activations + targets) / (activations * (1. - activations))
}

impl NeuralNetwork {
    pub fn learn(&mut self, data: &[DataPoint], batch_size: usize, learn_rate: f64) {
        for batch in data.chunks(batch_size) {
            self.learn_batch(batch, learn_rate);
        }
    }

    fn learn_batch(&mut self, batch: &[DataPoint], learn_rate: f64) {
        batch.iter().for_each(|data| {
            self.think(data);
            self.calc_backprop_cache(data);
            self.calc_grad(data);
        });

        self.apply_and_clear_grad(learn_rate);
    }

    fn think(&mut self, data: &DataPoint) {
        self.layers[0].think(&data.inputs, &mut self.layer_data[0]);
        for i in 1..self.layer_count {
            let [prev, cur, ..] = &mut self.layer_data[i - 1..] else {
                unreachable!()
            };
            self.layers[i].think(&prev.activations, cur);
        }
    }

    fn calc_backprop_cache(&mut self, data: &DataPoint) {
        self.layer_data[self.layer_count - 1].calc_backprop_cache_output(&data.targets);
        for i in (0..self.layer_count - 1).rev() {
            let [cur, next, ..] = &mut self.layer_data[i..] else {
                unreachable!()
            };
            cur.calc_backprop_cache_hidden(&self.layers[i + 1], &next);
        }
    }

    fn calc_grad(&mut self, data: &DataPoint) {
        self.layer_data[0].calc_grad(&data.inputs);
        for i in 1..self.layer_count {
            let [prev, cur, ..] = &mut self.layer_data[i - 1..] else {
                unreachable!()
            };
            cur.calc_grad(&prev.activations);
        }
    }

    fn apply_and_clear_grad(&mut self, learn_rate: f64) {
        self.layers
            .par_iter_mut()
            .zip(self.layer_data.par_iter_mut())
            .for_each(|(layer, data)| {
                data.apply_and_clear_grad(layer, learn_rate);
            });
    }
}

impl LayerData {
    pub fn calc_backprop_cache_output(&mut self, targets: &[f64]) {
        self.backprop
            .par_iter_mut()
            .enumerate()
            .for_each(|(neuron_index, backprop)| {
                let cost_deriv =
                    deriv_cross_entropy(self.activations[neuron_index], targets[neuron_index]);
                let acti_deriv = super::deriv_activate(self.weighted_inputs[neuron_index]);
                *backprop = cost_deriv * acti_deriv;
            });
    }

    pub fn calc_backprop_cache_hidden(&mut self, next: &Layer, next_data: &LayerData) {
        self.backprop
            .par_iter_mut()
            .enumerate()
            .for_each(|(neuron_index, backprop)| {
                *backprop = next_data
                    .backprop
                    .iter()
                    .map(|next_backprop| *next_backprop * next.weights[neuron_index])
                    .sum::<f64>()
                    * super::deriv_activate(self.weighted_inputs[neuron_index]);
            });
    }

    pub fn calc_grad(&mut self, inputs: &[f64]) {
        for output in 0..self.activations.len() {
            for input in 0..inputs.len() {
                self.grad_w[output * inputs.len() + input] += inputs[input] * self.backprop[output];
                self.grad_b[output] += self.backprop[output];
            }
        }
    }

    pub fn apply_and_clear_grad(&mut self, layer: &mut Layer, learn_rate: f64) {
        self.grad_w
            .par_iter_mut()
            .zip(layer.weights.par_iter_mut())
            .for_each(|(grad, weight)| {
                *weight -= learn_rate * *grad;
                *grad = 0.;
            });

        self.grad_b
            .par_iter_mut()
            .zip(layer.biases.par_iter_mut())
            .for_each(|(grad, bias)| {
                *bias -= learn_rate * *grad;
                *grad = 0.;
            });
    }
}