use anyhow::*;
use std::ops::*;
use num::Float;
use std::fmt;
use realfft::{RealFftPlanner, num_complex::Complex};
use crossterm::{
    QueueableCommand,
    cursor, style
};
use std::io::{stdout, Write};
 
use crate::{config::FLOAT, deep::Axon};
//use rayon::prelude::*;

/// A dynamically allocated point in N dimensional space. It can be used to describe the shape of a
/// matrix.
///
/// **NOTE:** ConstPoint should be used where possible.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Point<'a>(pub &'a [usize]);

/// A macro for creating points.
/// # Eg.
/// I![const 2,3,4];
/// I![1,2,3];
/// I![some_vec];
#[macro_export]
macro_rules! I{
    (const [$($a: expr),* $(,)?]) => {{
        use exotic::matrix::ConstPoint;
        ConstPoint([$($a),*])
    }};
    (const $($a: expr),* $(,)?) => {{
        use exotic::matrix::ConstPoint;
        ConstPoint([$($a),*])
    }};
    ($a: expr) => {{
        Point(&$a[..])
    }};
    ($($a: expr),* $(,)?) => {{
        Point(&[$($a),*][..])
    }};
}

#[derive(Clone, Copy, Debug)]
pub struct UnitIter<'a> {
    p: Point<'a>,
    n: usize,
    v: usize,
}

impl<'a> Iterator for UnitIter<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        if self.n >= self.p.0.len() {
            return None;
        }

        let this = self.v % self.p.0[self.n];
        self.v = (self.v - this) / self.p.0[self.n];
        self.n += 1;

        Some(this)
    }
}

pub trait Shape<'a> {
    fn data(&self) -> &[usize];
    fn len(&self) -> usize;

    fn unit_index(&self, mut v: usize, buffer: &mut [usize]) {
        for n in 0..buffer.len() {
            let this = v % self.data()[n];
            v = (v - this) / self.data()[n];
            buffer[n] = this;
        }
    }

    fn unit_iter(&'a self, v: usize) -> UnitIter<'a>;

    fn unit_offset(&self, n: usize, offset: Self) -> Option<usize>;

    fn volume(&self) -> usize {
        let mut x = 1;
        for n in self.data().iter() {
            x = x * n
        }
        x
    }

    fn within(&self, other: Self) -> bool;
}

impl<'a> Shape<'a> for Point<'a> {
    fn data(&self) -> &[usize] {
        self.0
    }

    fn len(&self) -> usize {
        self.0.len()
    }

    fn unit_iter(&'a self, v: usize) -> UnitIter<'a> {
        UnitIter {
            p: *self,
            v: v,
            n: 0,
        }
    }

    fn unit_offset(&self, mut n: usize, offset: Self) -> Option<usize> {
        assert!(self.len() == offset.len());
        let mut offset_amount = 1;

        for i in 0..offset.len() {
            if offset.0[i] >= self.0[i] {
                return None;
            }
            n += offset.0[i] * offset_amount;
            offset_amount = offset_amount * self.0[i];
        }
        Some(n)
    }

    fn within(&self, other: Self) -> bool {
        assert!(self.len() == other.len());
        for n in 0..self.len() {
            if self.0[n] >= other.0[n] {
                return false;
            }
        }
        true
    }
}

/// A statically allocated point in N dimensional space. It can be used to describe the shape of a
/// matrix.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ConstPoint<const DIM: usize>(pub [usize; DIM]);

impl<const DIM: usize> ConstPoint<DIM> {
    pub fn unit_offset(&self, mut n: usize, offset: &[usize]) -> Option<usize> {
        assert!(DIM == offset.len());
        let mut offset_amount = 1;

        for i in 0..offset.len() {
            if offset[i] >= self.0[i] {
                return None;
            }
            n += offset[i] * offset_amount;
            offset_amount = offset_amount * self.0[i];
        }
        Some(n)
    }

    pub fn unit_index(&self, mut v: usize, buffer: &mut [usize; DIM]) {
        for n in 0..DIM {
            let this = v % self.data()[n];
            v = (v - this) / self.data()[n];
            buffer[n] = this;
        }
    }
    pub fn const_unit_index(&self, mut v: usize) -> [usize; DIM] {
        let mut buffer = [0; DIM];
        for n in 0..DIM {
            let this = v % self.data()[n];
            v = (v - this) / self.data()[n];
            buffer[n] = this;
        }
        buffer
    }

}
impl ConstPoint<2>{
    pub const fn volume(&self) -> usize{
        self.0[0] * self.0[1]
    }
    pub const fn widen(mut self, w: usize) -> Self{
        self.0[0] += w;
        self
    }
}

impl<'a, const DIM: usize> Shape<'a> for ConstPoint<DIM> {
    fn data(&self) -> &[usize] {
        &self.0
    }

    fn len(&self) -> usize {
        DIM
    }

    fn unit_iter(&'a self, v: usize) -> UnitIter<'a> {
        UnitIter {
            p: I!(self.0),
            v: v,
            n: 0,
        }
    }

    fn unit_offset(&self, mut n: usize, offset: ConstPoint<DIM>) -> Option<usize> {
        let mut offset_amount = 1;

        for i in 0..DIM {
            if offset.0[i] >= self.0[i] {
                return None;
            }
            n += offset.0[i] * offset_amount;
            offset_amount = offset_amount * self.0[i];
        }
        Some(n)
    }

    fn within(&self, other: ConstPoint<DIM>) -> bool {
        for n in 0..DIM {
            if self.0[n] >= other.0[n] {
                return false;
            }
        }
        true
    }
}

/// A simple wrapper around a slice that describes an N dimensional shape for it, and implements
/// functions used for matrix math.
#[derive(Clone, Copy, PartialEq)]
pub struct Matrix<'a, T: Float, const DIM: usize> {
    pub data: &'a [T],
    pub shape: ConstPoint<DIM>,
}

impl<'a, const DIM: usize, const LEN: usize> Axon<LEN> for Matrix<'a, FLOAT, DIM>{
    fn zero() -> Self{
        Matrix{
            data: NULL,
            shape: ConstPoint([0; DIM])
        }
    }
}

impl<'a, T: Float, const DIM: usize> Index<&[usize]> for Matrix<'a, T, DIM> {
    type Output = T;

    fn index<'b>(&'b self, i: &[usize]) -> &'b T {
        assert!(i.len() == DIM);
        let mut p = 0;
        let mut k = 1;

        for n in 0..DIM {
            if i[n] > self.shape.0[n] {
                panic!("{:?} out of bounds: {:?}", i, self.shape)
            }
            p += i[n] * k;
            k = k * self.shape.0[n]
        }
        &self.data[p]
    }
}

impl<'a, T: Float, const DIM: usize> Index<[usize; DIM]> for Matrix<'a, T, DIM> {
    type Output = T;

    fn index<'b>(&'b self, i: [usize; DIM]) -> &'b T {
        let mut p = 0;
        let mut k = 1;

        for n in 0..DIM {
            if i[n] > self.shape.0[n] {
                panic!("{:?} out of bounds: {:?}", i, self.shape)
            }
            p += i[n] * k;
            k = k * self.shape.0[n]
        }
        &self.data[p]
    }
}

impl <'a, T: Float, const DIM: usize> Index<usize> for Matrix<'a, T, DIM>{
    type Output = T;

    fn index(&self, n: usize) -> &T{
        &self.data[n]
    }
}

pub struct MatrixDimIter<'a, T: Float,  const DIM: usize> {
    source: &'a Matrix<'a, T, DIM>,
    n: usize,
    shell_volume: usize,
}

impl<'a, T: Float, const DIM: usize> Iterator for MatrixDimIter<'a, T, DIM> {
    type Item = &'a [T];
    fn next(&mut self) -> Option<&'a [T]> {
        if self.n >= self.source.shape.0[DIM - 1] {
            return None;
        }
        let tmp = &self.source.data[self.n * self.shell_volume..(self.n + 1) * self.shell_volume];
        self.n += 1;
        Some(tmp)
    }
}

impl<'a, T: Float, const DIM: usize> Matrix<'a, T, DIM> {
    pub fn iter_outer(&'a self) -> MatrixDimIter<'a, T, DIM> {
        MatrixDimIter {
            source: self,
            n: 0,
            shell_volume: I![self.shape.0[0..DIM - 1]].volume(),
        }
    }

    pub fn dim(&self) -> usize {
        DIM
    }

    pub fn possible(&self) -> Result<()> {
        if self.shape.volume() == self.data.len() {
            Ok(())
        } else {
            Err(anyhow!(
                "Imposible shape. Vector of length {} does not fit into {:?}",
                self.data.len(),
                self.shape.volume()
            ))
        }
    }

    pub fn volume(&self) -> usize {
        self.data.len()
    }
}

impl<'a, T: Float + fmt::Display> fmt::Debug for Matrix<'a, T, 2> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut longest_decimal = 5;

        for y in 0..self.shape.0[1] {
            for n in &self.data[y * self.shape.0[0]..(y + 1) * self.shape.0[0]] {
                let len = format!("{:.2}", n).len() + 2;
                if len > longest_decimal {
                    longest_decimal = len
                }
            }
        }

        write!(f, "     ")?;
        for n in 0..self.shape.0[0] {
            write!(f, "c[{}]  {}", n, " ".repeat(longest_decimal - 5))?;
        }

        for y in 0..self.shape.0[1] {
            write!(f, "\nr[{}]   ", y)?;
            for n in &self.data[y * self.shape.0[0]..(y + 1) * self.shape.0[0]] {
                let n = format!("{:.2}", n);

                write!(f, "{} {}", n, " ".repeat(longest_decimal - n.len()))?
            }
        }
        Ok(())
    }
}

impl<'a> Matrix<'a, FLOAT, 2> {
    pub fn heatmap(&self, magnitude: FLOAT) -> Result<()>{
        let mut stdout = stdout();
        stdout.queue(cursor::MoveToNextLine(1))?.queue(style::Print("|"))?;

        for y in 0..self.shape.0[1]{
            for x in 0..self.shape.0[0]{
                let val = self[[x, y]] / 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))?;
                if val as u8 == 0{
                    stdout.queue(style::ResetColor)?;
                }
                stdout.queue(style::Print("  "))?;
            }
            stdout.queue(style::Print("|\n|"))?;
        }
        stdout.queue(style::ResetColor)?.queue(style::Print("-\n"))?;

        stdout.flush()?;
        Ok(())
    }

    pub fn lazy_flat_n_pad<const LEN: usize>(self, padding: usize) -> LazyFlatnPad2d<'a, LEN>{
        assert!((self.shape.0[0] + padding) * self.shape.0[1] == LEN);
        LazyFlatnPad2d(self, padding)
    }
}

pub fn conv_real<const LEN: usize>(a: &mut [FLOAT; LEN], b: &mut [FLOAT; LEN]) -> [FLOAT; LEN]
 where [FLOAT; LEN/2+1]: Sized
{  
    let mut planner = RealFftPlanner::<FLOAT>::new();
    let fft = planner.plan_fft_forward(LEN);
 
    let mut a_buff = [Complex{re: 0., im: 0.}; LEN/2+1];
    fft.process(&mut a[..], &mut a_buff[..]).unwrap();

    let mut b_buff = [Complex{re: 0., im: 0.}; LEN/2+1];
    fft.process(&mut b[..], &mut b_buff[..]).unwrap();

    drop(fft);

    let mut conv_buff = [Complex{re: 0., im: 0.}; LEN/2+1];
    for n in 0..LEN/2+1{
        conv_buff[n] = a_buff[n] * b_buff[n];
    }

    let fft = planner.plan_fft_inverse(LEN);
    let mut res = [0.; LEN];
    fft.process(&mut conv_buff[..], &mut res[..]).unwrap();
    res.iter_mut().for_each(|n| *n /= LEN as FLOAT);
    res
} 

pub fn conv_flat<const A: usize, const B: usize>(a: &[FLOAT; A], b: &[FLOAT; B]) -> [FLOAT; A+B-1]
    where [FLOAT; (A+B-1)/2+1]: Sized
{  
    let pad = |v: &[FLOAT]| -> [FLOAT; A+B-1]{
        let mut buff = [0.; A+B-1];
        for n in 0..v.len(){
            buff[n] = v[n]
        }
        buff
    };

    let mut planner = RealFftPlanner::<FLOAT>::new();
    let fft = planner.plan_fft_forward(A+B-1);
 
    let mut a_pad = pad(&a[..]);
    let mut a_buff = [Complex{re: 0., im: 0.}; (A+B-1)/2+1];
    fft.process(&mut a_pad[..], &mut a_buff[..]).unwrap();
    drop(a_pad);

    let mut b_pad = pad(&b[..]);
    let mut b_buff = [Complex{re: 0., im: 0.}; (A+B-1)/2+1];
    fft.process(&mut b_pad[..], &mut b_buff[..]).unwrap();
    drop(b_pad);

    drop(fft);

    let mut conv_buff = [Complex{re: 0., im: 0.}; (A+B-1)/2+1];
    for n in 0..(A+B-1)/2+1{
        conv_buff[n] = a_buff[n] * b_buff[n];
    }

    let fft = planner.plan_fft_inverse(A+B-1);
    let mut res = [0.; A+B-1];
    fft.process(&mut conv_buff[..], &mut res[..]).unwrap();
    res.iter_mut().for_each(|n| *n /= (A+B-1) as FLOAT);
    res
}

/// The convolve function is broken, and should therefore at the moment not be used. LEN should not
/// be equal to big.volume()
pub fn convolve<'a, T: Float, const DIM: usize, const LEN: usize>(
    _big: Matrix<'a, T, DIM>,
    _small: Matrix<'a, T, DIM>,
) -> [T; LEN] {
    unimplemented!();
}

#[macro_export]
macro_rules! matrix {
    ($data: expr, $shape: expr $(,)?) => {{
        use exotic::matrix::*;
        assert!($shape.volume() == $data.len());
        Matrix {
            data: $data,
            shape: $shape,
        }
    }};
}

#[macro_export]
macro_rules! try_matrix {
    ($data: expr, $shape: expr $(,)?) => {{
        use exotic::{anyhow, matrix::*, Result};
        if $shape.volume() == $data.len() {
            Ok(Matrix {
                data: $data,
                shape: $shape,
            })
        } else {
            Err(anyhow!(
                "Data of len {} does not fit into shape {:?}, expected {}",
                $data.len(),
                $shape,
                $shape.volume()
            ))
        }
    }};
}

const NULL: &[FLOAT; 0] = &[];

#[derive(Clone, Copy)]
pub struct LazyFlatnPad2d<'a, const LEN: usize>(pub Matrix<'a, FLOAT, 2>, pub usize);

impl<'a, const LEN: usize> Axon<LEN> for LazyFlatnPad2d<'a, LEN>{
    fn zero() -> Self{
        //debug_assert!(LEN % 2 == 0, "LazyFlatnPad2d LEN should always be even, as it is padded evenly");
        Self(Axon::<LEN>::zero(), 0)
    }
}

impl<'a, const LEN: usize> Index<usize> for LazyFlatnPad2d<'a, LEN>{
    type Output = FLOAT;
    fn index(&self, n: usize) -> &FLOAT{
        debug_assert!(n < LEN, "Index out of FlatnPad2d");
        let x = n % (self.0.shape.0[0] + self.1);
        if x < self.0.shape.0[0]{
            let y = (n-x) / (self.0.shape.0[0] + self.1);
            &self.0.data[y * self.0.shape.0[0] + x]
        }else{
            &0.
        }
    }
}

impl<'a, const LEN: usize> LazyFlatnPad2d<'a, LEN>{
    pub fn heatmap(&self, magnitude: FLOAT) -> Result<()>{
        Matrix{
            data: &self.to_array(),
            shape: ConstPoint([self.0.shape.0[0] + self.1, self.0.shape.0[1]])
        }
            .heatmap(magnitude)?;
        Ok(())
    }
}

#[derive(Copy, Clone)]
pub struct MatrixSlice<'a, const DIM: usize, const LEN: usize>{
    pub data: &'a [FLOAT],
    pub shape: ConstPoint<DIM>,
    pub slice: ConstPoint<DIM>
}

impl<'a, const DIM: usize, const LEN: usize> Index<usize> for MatrixSlice<'a, DIM, LEN>{
    type Output = FLOAT;
    fn index(&self, n: usize) -> &FLOAT{
        let mut index = 0;
        let mut ofset = 1;
        for (i, j) in self.slice.unit_iter(n).enumerate(){
            index += j * ofset;
            ofset += self.shape.0[i]
        }
        &self.data[index]
    }
}

impl<'a, const DIM: usize, const LEN: usize> Axon<LEN> for MatrixSlice<'a, DIM, LEN>{
    fn zero() -> Self{
        MatrixSlice{
            data: NULL,
            shape: ConstPoint([0; DIM]),
            slice: ConstPoint([0; DIM]),
        }
    }
}
