use std::array;
use crate::algo::compute::compute;
use crate::algo::fix::{apply_error, calc_error};
use rand::Rng;
use std::ops::{Div, Mul};

pub(super) struct Net<const InputLayerSize: usize, const OutputLayerSize: usize> {
    inner_layer_size: usize,
    pub input_data: [f64; InputLayerSize],
    i_to_h_weights: Vec<[f64; InputLayerSize]>,
    hidden_potentials: Vec<f64>,
    hidden_data: Vec<f64>,
    h_to_o_weights: [Vec<f64>; OutputLayerSize],
    output_potentials: [f64; OutputLayerSize],
    pub output_data: [f64; OutputLayerSize],
    output_errors: [f64; OutputLayerSize],
    hidden_errors: Vec<f64>,
}

impl<const InputLayerSize: usize, const OutputLayerSize: usize>
    Net<InputLayerSize, OutputLayerSize>
{
    pub(crate) fn new() -> Self {
        return Self {
            inner_layer_size: 1,
            input_data: [0.0; InputLayerSize],
            i_to_h_weights: Vec::from([[0.0; InputLayerSize]]),
            hidden_potentials: Vec::from([0.0]),
            hidden_data: Vec::from([0.0]),
            h_to_o_weights: array::from_fn(|_| Vec::from([0.0])),
            output_potentials: [0.0; OutputLayerSize],
            output_data: [0.0; OutputLayerSize],
            output_errors: [0.0; OutputLayerSize],
            hidden_errors: Vec::from([0.0]),
        }
    }


    fn _sigmoid(x: f64) -> f64 {
        return 1.0.div(1.0 + (-x).exp());
    }

    fn _sigmoidDerivative(x: f64) -> f64 {
        return x.mul(1.0 - x);
    }

    pub fn compute(&mut self) {
        compute(
            &self.input_data,
            self.i_to_h_weights.as_slice(),
            self.hidden_potentials.as_mut_slice(),
            self.hidden_data.as_mut_slice(),
            Self::_sigmoid,
        );

        compute(
            self.hidden_data.as_slice(),
            &self.h_to_o_weights,
            &mut self.output_potentials,
            &mut self.output_data,
            Self::_sigmoid,
        )
    }

    pub fn fix(&mut self, expected: &[f64; OutputLayerSize], n: f64) {
        for (i, (a, e)) in self.output_data.iter().zip(expected).enumerate() {
            self.output_errors[i] = e - a;
        }

        calc_error(
            &self.output_errors,
            &self.h_to_o_weights,
            self.hidden_errors.as_mut_slice(),
        );

        apply_error(
            n,
            &self.output_errors,
            &mut self.h_to_o_weights,
            self.hidden_data.as_slice(),
            &self.output_potentials,
            Self::_sigmoidDerivative,
        );
        apply_error(
            n,
            self.hidden_errors.as_slice(),
            self.i_to_h_weights.as_mut_slice(),
            &self.input_data,
            self.hidden_potentials.as_slice(),
            Self::_sigmoidDerivative,
        );
    }

    pub fn hidden_layer_size(&self) -> usize {
        return self.inner_layer_size;
    }

    pub fn resize_hidden_layer(&mut self, new_size: usize) {
        assert!(new_size > 0);
        self.inner_layer_size = new_size;
        self.i_to_h_weights.resize(new_size, [0.0; InputLayerSize]);
        self.hidden_potentials.resize(new_size, 0.0);
        self.hidden_data.resize(new_size, 0.0);
        for w in self.h_to_o_weights.iter_mut() {
            w.resize(new_size, 0.0);
        }
        self.hidden_errors.resize(new_size, 0.0);
    }

    pub fn set_random_weights(&mut self, rng: &mut impl Rng) {
        for ww in self.i_to_h_weights.iter_mut() {
            for w in ww.iter_mut() {
                *w = rng.random()
            }
        }
        for ww in self.h_to_o_weights.iter_mut() {
            for w in ww.iter_mut() {
                *w = rng.random()
            }
        }
    }
}