use anyhow::*;
use rand::*;
use rayon::prelude::*;
use crate::config::*;
use std::ops::*;

use std::fmt::{self, Display};
use crossterm::{
    QueueableCommand,
    style
};
use std::io::{stdout, Write};
 
#[macro_export]
macro_rules! axon{
    (ref $data: expr) => {{
        AxonRef::Ref($data)
    }};
    ($a: expr, $b: expr) => {{
        AxonChain($a, $b)
    }}
}

//pub mod conv;
//pub mod par;
pub mod recurent;
pub mod activation;

// Number Array Trait

pub trait Axon<const LEN: usize>: Copy + Sync + Index<usize, Output = FLOAT> {
    fn zero() -> Self;
    fn to_array(&self) -> [FLOAT; LEN]{
        let mut tmp = [0.; LEN];
        for n in 0..LEN{
            tmp[n] = self[n]
        }
        tmp
    }
    fn to_ref<'a>(&'a self) -> AxonRef<'a, Self, LEN>{
        AxonRef::Ref(self)
    }
    fn to_fake(self) -> AxonRef<'static, Self, LEN>{
        AxonRef::Fake(self)
    }
}

#[derive(Copy, Clone)]
pub struct AxonChain<A: Axon<A_>, B: Axon<B_>, const A_: usize, const B_: usize>(
    A,
    B,
);

impl<A: Axon<A_>, B: Axon<B_>, const A_: usize, const B_: usize> Index<usize>
    for AxonChain<A, B, A_, B_>
{
    type Output = FLOAT;
    fn index(&self, n: usize) -> &FLOAT {
        if n < A_ {
            &self.0[n]
        } else {
            debug_assert!(n < A_ + B_, "Index out of chain: {} >= {} + {}", n, A_, B_);
            &self.1[n-A_]
        }
    }
}

#[derive(Clone, Copy)]
pub enum AxonRef<'a, T: Axon<LEN>, const LEN: usize>{
    Ref(&'a T),
    Fake(T),
}

 
impl<'a, T: Axon<LEN>, const LEN: usize> Display for AxonRef<'a, T, LEN>{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut stdout = stdout();

        let magnitude = (0..LEN).map(|n|self[n]).fold(FLOAT::NAN, FLOAT::max);

        for x in 0..LEN{
            let val = self[x] / magnitude * 255.;
            let color = match val.is_sign_positive(){
                true  => style::Color::Rgb{r:0, g:0, b: val as u8},
                false => style::Color::Rgb{r: val as u8, g:0, b:0},
            };
            stdout.queue(style::SetBackgroundColor(color)).unwrap();
            if val as u8 == 0{
                stdout.queue(style::ResetColor).unwrap();
            }
            stdout.queue(style::Print("  ")).unwrap();
        }

        stdout.queue(style::ResetColor).unwrap();
        stdout.flush().unwrap();

        write!(f, "")?;
        Ok(())
    }
}

impl<'a, T: Axon<LEN>, const LEN: usize> Deref for AxonRef<'a, T, LEN> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        match self{
            AxonRef::Ref(r) => r,
            AxonRef::Fake(f) => f,
        }
    }
}

impl<'a, T: Axon<LEN>, const LEN: usize> Index<usize> for AxonRef<'a, T, LEN>{
    type Output = FLOAT;
    fn index(&self, n: usize) -> &FLOAT {
        match self{
            AxonRef::Ref(r) => &r[n],
            AxonRef::Fake(f) => &f[n],
        }
    }
}

impl<A: Axon<A_>, B: Axon<B_>, const A_: usize, const B_: usize>
    Axon<{ A_ + B_ }> for AxonChain<A, B, A_, B_>{
    fn zero() -> Self{AxonChain(A::zero(), B::zero())}
}

impl<const LEN: usize> Axon<LEN> for [FLOAT; LEN] {
    fn zero() -> Self{[0.;LEN]}
}

impl<'a,T: Axon<LEN>, const LEN: usize> Axon<LEN> for AxonRef<'a, T, LEN>{
    fn zero() -> Self{AxonRef::Fake(T::zero())}
}

// Layer trait

pub trait LayerTy{
    type Gradient;
    type Output;
}
 
pub trait Layer<const I_LEN: usize, const O_LEN: usize>: LayerTy where
 <Self as LayerTy>::Gradient: Axon<I_LEN>,
 <Self as LayerTy>::Output: Axon<O_LEN>,
{
    fn predict(&self, i: impl Axon<I_LEN>) -> Result<Self::Output>;
    fn backpropagate(
        &mut self,
        i: impl Axon<I_LEN>,
        gradient: impl Axon<O_LEN>,
    ) -> Result<Self::Gradient>;
}

#[derive(Clone, Copy, Debug)]
pub struct Neuron<const LEN: usize> {
    w: [FLOAT; LEN],
    b: FLOAT,
}

impl<const LEN: usize> Neuron<LEN> {
    pub fn random() -> Self {
        let mut buffer = [0.; LEN];
        buffer
            .iter_mut()
            .for_each(|n| *n = random());
        Self {
            w: buffer,
            b: random(),
        }
    }

    pub fn const_xavier<const OUT: usize>() -> Self {
        let xavier = ||random::<FLOAT>() * (2. / (LEN as FLOAT + OUT as FLOAT)).sqrt();
        let mut buffer = [0.; LEN];
        buffer
            .iter_mut()
            .for_each(|n| *n = xavier());
        Self {
            w: buffer,
            b: xavier(),
        }
    }

    pub fn zero() -> Self {
        let buffer = [0.; LEN];
        Self {
            w: buffer,
            b: 0.,
        }
    }

    pub fn predict<T: Index<usize, Output = FLOAT>>(&self, i: T) -> FLOAT {
        let mut sum = 0.;
        for n in 0..LEN - 1 {
            sum += self.w[n] * i[n]
        }
        sum.mul_add(i[LEN - 1], self.b)
    }

    pub fn backpropagate<T: Index<usize, Output = FLOAT>>(&mut self, i: T, delta: FLOAT, lr: FLOAT) -> Result<[FLOAT; LEN]> {
        let delta = round(delta);

        self.b = self.b - delta * lr;
        if self.b.is_infinite() {
            bail!("Exploding gradient in bias")
        }

        let mut gradient = self.w;

        for n in 0..LEN {
            gradient[n] = round(gradient[n] * delta * i[n]);
            self.w[n] -= delta * i[n] * lr;
            if self.w[n].is_infinite() {
                bail!("Exploding gradient in weight {}\n input = {}\n Δ = {}", n, i[n], delta)
            }
        }

        Ok(gradient)
    }
}

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

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

impl<const I_LEN: usize, const O_LEN: usize> LayerTy for DenseLayer<I_LEN, O_LEN>{
    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>
{
    fn predict(&self, i: impl Axon<I_LEN>) -> Result<[FLOAT; O_LEN]> {
        let mut buffer = [0.; O_LEN];

        buffer.iter_mut().enumerate().for_each(|(n, o)|{
            *o = self.neurons[n].predict(axon!(ref &i))
        });

        Ok(buffer)
    }

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

        self.neurons.iter_mut().zip(matrix.iter_mut()).enumerate().map(|(n, (neuron, o))| {
            *o = neuron
                .backpropagate(axon!(ref &i), gradient[n], lr)
                .context(format!("Error in neuron {}", n))?;
            Ok(())
       }).collect::<Result<()>>()?;

        
        let mut buffer = [0.; I_LEN];
        buffer.iter_mut().enumerate().for_each(|(x, o)|{
            for y in 0..O_LEN{
                *o += matrix[y][x]
            }
        });

        Ok(buffer)
    }
}
