use std::ops::Add;
use itertools::Itertools;
use nalgebra::{EuclideanNorm, LpNorm, Matrix, Norm, OMatrix, SVector};
use num_traits::Pow;
use serde::{Serialize, Deserialize};
use crate::system::{GriddedPosition, Position, Storage};

#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Gridded<const D: usize>(SVector<i32, D>);

#[derive(Clone, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Continuous<const D: usize>(SVector<f32, D>);

impl<const D: usize> Add for Continuous<D> {
    type Output = Continuous<D>;

    fn add(self, rhs: Self) -> Self::Output {
        Continuous(self.0 + rhs.0)
    }
}

impl<const D: usize> Position for Continuous<D> {
    const DIM: usize = 0;

    fn zero() -> Self {
        Continuous(SVector::<f32, D>::zeros())
    }

    fn abs(&self) -> f32 {
        self.0.norm()
    }

    fn from_cartesian(cartesian: &[f32]) -> Self {
        Continuous(SVector::<f32, D>::from_fn(|i, _| cartesian[i]))
    }

    fn to_cartesian(&self) -> Vec<f32> {
        self.0.as_slice().to_vec()
    }
}

impl<const D: usize> Add for Gridded<D> {
    type Output = Gridded<D>;

    fn add(self, rhs: Self) -> Self::Output {
        Gridded(self.0 + rhs.0)
    }
}

impl<const D: usize> Position for Gridded<D> {
    const DIM: usize = 0;

    fn zero() -> Self {
        Gridded(SVector::<i32, D>::zeros())
    }

    fn abs(&self) -> f32 {
        (self.0.fold(0, |r, c| r + c.pow(2)) as f32).sqrt()
    }

    fn from_cartesian(cartesian: &[f32]) -> Self {
        Gridded(SVector::<i32, D>::from_fn(|i, _| cartesian[i] as i32))
    }

    fn to_cartesian(&self) -> Vec<f32> {
        self.0.as_slice()
            .iter()
            .map(|a| *a as f32)
            .collect_vec()
    }
}

pub struct KDSpace<const N: usize>(pub(crate) kiddo::KdTree<f32, (), N>);

impl<const D: usize> Storage<Gridded<D>> for KDSpace<D> {
    fn is_occupied(&self, position: &Gridded<D>) -> bool {
        let a = self.0.best_n_within(
            &position.0.data.0[0].map(|i| i as f32),
            0f32,
            1,
            &|a, b| {
                LpNorm(1).metric_distance(
                    &SVector::<f32, D>::from_row_slice(a),
                    &SVector::<f32, D>::from_row_slice(b),
                )
            },
        ).unwrap();

        !a.is_empty()
    }

    fn deposit(&mut self, position: &Gridded<D>) {
        self.0.add(&position.0.data.0[0].map(|i| i as f32), ())
            .expect("Failed to write to space")
    }
}

impl<const D: usize> Storage<Continuous<D>> for KDSpace<D> {
    fn is_occupied(&self, position: &Continuous<D>) -> bool {
        let a = self.0.best_n_within(
            &position.0.data.0[0],
            0f32,
            1,
            &|a, b| {
                EuclideanNorm.metric_distance(
                    &SVector::<f32, D>::from_row_slice(a),
                    &SVector::<f32, D>::from_row_slice(b),
                )
            },
        ).unwrap();

        !a.is_empty()
    }

    fn deposit(&mut self, position: &Continuous<D>) {
        self.0.add(&position.0.data.0[0], ())
            .expect("Failed to write to space")
    }
}

pub struct VectorStorage {
    backing: Vec<bool>,
    grid_size: usize,
}

impl<const D: usize> Storage<Gridded<D>> for VectorStorage {
    fn is_occupied(&self, position: &Gridded<D>) -> bool {
        let mut index: usize = 0;

        for i in 0..D {
            index += (position.0[i] + (self.grid_size as i32 / 2)) as usize * self.grid_size * i;
        }

        return self.backing[index];
    }

    fn deposit(&mut self, position: &Gridded<D>) {
        let mut index: usize = 0;

        for i in 0..D {
            index += (position.0[i] + (self.grid_size as i32 / 2)) as usize * self.grid_size * i;
        }

        self.backing[index] = true;
    }
}