use anyhow::anyhow;
use rand::Rng;
use crate::system::{GriddedPosition, Position, Storage};
use crate::system::spaces::square_grid::{Grid2D, Grid3D};
use crate::system::sticker::Sticker;

fn taxicab_grid2(a: &[f32; 2], b: &[f32; 2]) -> f32 {
    (a[0] - b[0]).abs() + (a[1] - b[1]).abs()
}

fn taxicab_grid3(a: &[f32; 3], b: &[f32; 3]) -> f32 {
    (a[0] - b[0]).abs() + (a[1] - b[1]).abs() + (a[2] - b[2]).abs()
}

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

impl<const N: usize> KDSpace<N> {
    pub fn new() -> KDSpace<N> {
        KDSpace { inner: kiddo::KdTree::new() }
    }
}

impl Storage<Grid2D> for KDSpace<2> {
    fn is_occupied(&self, position: &Grid2D) -> bool {
        let a = self.inner.best_n_within(&[position.x as f32, position.y as f32], 0f32, 1, &taxicab_grid2).unwrap();
        !a.is_empty()
    }

    fn deposit(&mut self, position: &Grid2D) {
        self.inner.add(&[position.x as f32, position.y as f32], ())
            .expect("Failed to write to space")
    }
}

impl Storage<Grid3D> for KDSpace<3> {
    fn is_occupied(&self, position: &Grid3D) -> bool {
        let a = self.inner.best_n_within(&[position.x as f32, position.y as f32, position.z as f32], 0f32, 1, &taxicab_grid3).unwrap();
        !a.is_empty()
    }

    fn deposit(&mut self, position: &Grid3D) {
        self.inner.add(&[position.x as f32, position.y as f32, position.z as f32], ())
            .expect("Failed to write to space")
    }
}

pub struct KDProbabilisticSticking {
    pub(crate) stick_probability: f32,
}

impl KDProbabilisticSticking {
    fn new(stick_probability: f32) -> anyhow::Result<KDProbabilisticSticking> {
        return if 0f32 < stick_probability && stick_probability <= 1f32 {
            Ok(KDProbabilisticSticking { stick_probability })
        } else {
            Err(anyhow!("Sticking probability outside of (0, 1] range."))
        }
    }
}

impl Sticker<Grid2D, KDSpace<2>> for KDProbabilisticSticking {
    fn should_stick<R: Rng>(&self, rng: &mut R, space: &KDSpace<2>, position: &Grid2D) -> bool {
        let a = space.inner.best_n_within(&[position.x as f32, position.y as f32], 1f32, Grid2D::NEIGHBOURS as usize, &taxicab_grid2)
            .unwrap();

        if a.len() == 0 {
            return false;
        }

        let q = 1f32 - self.stick_probability;
        let a = q.powi(a.len() as i32);
        rng.gen_range(0f32..1f32) > a
    }
}

impl Sticker<Grid3D, KDSpace<3>> for KDProbabilisticSticking {
    fn should_stick<R: Rng>(&self, rng: &mut R, space: &KDSpace<3>, position: &Grid3D) -> bool {
        let a = space.inner.best_n_within(&[position.x as f32, position.y as f32, position.z as f32], 1f32, Grid2D::NEIGHBOURS as usize, &taxicab_grid3)
            .unwrap();

        let q = (1f32 - self.stick_probability);
        let a = q.powi(a.len() as i32);
        rng.gen_range(0f32..1f32) > a
    }
}
