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 P2 {
    x: f32,
    y: f32,
}

impl P2 {
    fn as_arr(&self) -> [f32; 2] {
        [self.x, self.y]
    }

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

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

        P2 { x, y }
    }
}

impl Add for P2 {
    type Output = P2;

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

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

impl Position for P2 {
    type Cartesian = [f32; 2];

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

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

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

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

impl Storage<P2> for ContinuousStorage {
    fn is_occupied(&self, position: &P2) -> 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: &P2) {
        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<P2, ContinuousStorage> for ContinuousSticker {
    fn should_stick<R: Rng>(&self, _rng: &mut R, space: &ContinuousStorage, position: &P2) -> 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<P2> for ContinuousWalker {
    fn walk<R: Rng>(&self, rng: &mut R, position: &P2) -> P2 {
        position.clone() + P2::random_with_radius(rng, self.walk_step)
    }
}
