use crate::deep::*;
use std::ops::Index;
use rayon::prelude::*;

#[derive(Clone, Copy)]
pub struct DenseLayer<const I_LEN: usize, const O_LEN: usize>
where Neuron<{I_LEN + O_LEN}>: Sized{
    neurons: [Neuron<{I_LEN + O_LEN}>; O_LEN],
    pub buffer: [FLOAT; O_LEN],
    lr: FLOAT,
}

struct MemoryAxon<'a, T: Index<usize, Output = FLOAT>>(&'a [FLOAT], &'a T);
impl<'a, T: Index<usize, Output = FLOAT>> Index<usize> for MemoryAxon<'a, T>{
    type Output = FLOAT;
    fn index(&self, n: usize) -> &FLOAT{
        if n < self.0.len(){
            &self.0[n]
        }else{
            &self.1[n-self.0.len()]
        }
    }
}

impl<const I_LEN: usize, const O_LEN: usize> DenseLayer<I_LEN, O_LEN>
where Neuron<{I_LEN + O_LEN}>: Sized{

    pub fn random(lr: FLOAT) -> Self{
        let mut buffer = [Neuron::<{I_LEN + O_LEN}>::random(); O_LEN];
        buffer.iter_mut().for_each(|n|*n = Neuron::<{I_LEN + O_LEN}>::const_xavier::<O_LEN>());
        Self{
            neurons: buffer,
            buffer: [0.; O_LEN],
            lr,
        }
    }

    pub fn memorize<T: Axon<I_LEN>>(&mut self, i: T) -> Result<()>{
        let mut buffer = [0.; O_LEN];

        for (n, o) in buffer.iter_mut().enumerate(){
            let bruh = MemoryAxon(&self.buffer, &i);
            *o = self.neurons[n].predict(bruh)
        }

        self.buffer = buffer;
        Ok(())
    }

    pub fn memorize_par(&mut self, i: impl Axon<I_LEN> + Sync) -> Result<()>{
        let mut buffer = [0.; O_LEN];

        buffer.par_iter_mut().enumerate().for_each(|(n, o)|{
            let bruh = MemoryAxon(&self.buffer, &i);
            *o = self.neurons[n].predict(bruh)
        });

        self.buffer = buffer;
        Ok(())
    }

    pub fn forget(&mut self){
        self.buffer = [0.; O_LEN];
    }
}

impl<const I_LEN: usize, const O_LEN: usize> LayerTy for DenseLayer<I_LEN, O_LEN>
where Neuron<{I_LEN + O_LEN}>: Sized{
    type Gradient = [FLOAT; I_LEN];
    type Output   = [FLOAT; O_LEN];
}

impl<const I_LEN: usize, const O_LEN: usize> Layer<I_LEN, O_LEN> for DenseLayer<I_LEN, O_LEN>
where Neuron<{I_LEN + O_LEN}>: Sized{
    fn predict(&self, i: impl Axon<I_LEN>) -> Result<[FLOAT; O_LEN]>{
        let mut buffer = [0.; O_LEN];

        for n in 0..O_LEN{
            buffer[n] = self.neurons[n].predict(MemoryAxon(&self.buffer, &i))
        }

        Ok(buffer)
    }

    fn backpropagate(&mut self, i: impl Axon<I_LEN>, gradient: impl Axon<O_LEN>) -> Result<[FLOAT; I_LEN]>{
        let mut buffer = [0.; I_LEN];

        for n in 0..O_LEN{
            let d = self.neurons[n]
                .backpropagate(MemoryAxon(&self.buffer, &i), gradient[n], self.lr)
                .context(format!("Error in neuron {}", n))?;
            for n in 0..I_LEN{ 
                buffer[n] += d[n]
            }
        }

        Ok(buffer)
    }
}
