use std::f32::consts::PI;
use std::ops::Add;
use kiddo::distance::squared_euclidean;
use rand::Rng;
use serde::Serialize;
use crate::system::sticker::Sticker;
use crate::system::{Position, Storage};
use crate::system::walker::Walker;

#[derive(Serialize, Debug, Clone)]
pub struct P3 {
    x: f32,
    y: f32,
    z: f32,
}

impl P3 {
    fn as_arr(&self) -> [f32; 3] {
        [self.x, self.y, self.z]
    }

    pub fn random_with_radius<R: Rng>(rng: &mut R, radius: f32) -> P3 {
        let theta = rng.gen_range(0f32..1.0) * 2.0 * PI;
        let phi = rng.gen_range(0f32..1.0) * 2.0 * PI;

        let (x, y, z) = (
            radius * theta.sin() * phi.cos(),
            radius * theta.sin() * phi.sin(),
            radius * theta.cos()
        );

        P3 { x, y, z}
    }
}

impl Add for P3 {
    type Output = P3;

    fn add(self, rhs: Self) -> Self::Output {
        P3 {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

pub struct ContinuousStorage {
    pub inner: kiddo::KdTree<f32, (), 3>,
    pub ball_radius_sq: f32,
}

impl Position for P3 {
    type Cartesian = [f32; 3];

    fn zero() -> Self {
        P3 { x: 0f32, y: 0f32, z: 0f32 }
    }

    fn abs(&self) -> f32 {
        (self.x.powi(2) + self.y.powi(2) + self.z.powi(2)).powf(0.5)
    }

    fn from_cartesian(cartesian: Self::Cartesian) -> Self {
        P3 { x: cartesian[0], y: cartesian[1], z: cartesian[3] }
    }

    fn to_cartesian(&self) -> Self::Cartesian {
        [self.x, self.y, self.z]
    }
}

impl Storage<P3> for ContinuousStorage {
    fn is_occupied(&self, position: &P3) -> bool {
        let (dist_sq, _) = self.inner.nearest_one(&position.as_arr(), &squared_euclidean).unwrap();

        // Is the distance of this point to the next one less than twice the ball radius
        dist_sq < 2.0 * self.ball_radius_sq
    }

    fn deposit(&mut self, position: &P3) {
        self.inner.add(&position.as_arr(), ()).expect("Failed to write to space")
    }
}

pub struct ContinuousSticker {
    /// INVARIANT: THIS SHOULD BE GREATER THAN THE BALL_RADIUS_SQ value
    pub range_sq: f32,
}

impl Sticker<P3, ContinuousStorage> for ContinuousSticker {
    fn should_stick<R: Rng>(&self, _rng: &mut R, space: &ContinuousStorage, position: &P3) -> bool {
        let (a, _) = space.inner.nearest_one(&position.as_arr(), &squared_euclidean).unwrap();
        a < self.range_sq
    }
}

pub struct ContinuousWalker {
    pub walk_step: f32
}

impl Walker<P3> for ContinuousWalker {
    fn walk<R: Rng>(&self, rng: &mut R, position: &P3) -> P3 {
        position.clone() + P3::random_with_radius(rng, self.walk_step)
    }
}