pub mod draw;
pub mod nms;
pub mod roi;

use nalgebra::{Point, Point2, Point3, SVector, SimdPartialOrd, SimdValue};
pub trait Num: num::Num + Copy + core::fmt::Debug + 'static {}
impl<T: num::Num + Copy + core::fmt::Debug + 'static> Num for T {}

/// An axis aligned bounding box in `D` dimensions, defined by the minimum vertex and a size vector.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct AxisAlignedBoundingBox<T: Num, const D: usize> {
    /// The point of the bounding box closest to the origin
    point: Point<T, D>,
    /// The size of the bounding box in each dimension
    size: SVector<T, D>,
}

pub type Aabb<T, const D: usize> = AxisAlignedBoundingBox<T, D>;
pub type Aabb2<T> = AxisAlignedBoundingBox<T, 2>;
pub type Aabb3<T> = AxisAlignedBoundingBox<T, 3>;

impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
    pub fn new(point: Point<T, D>, size: SVector<T, D>) -> Self {
        Self { point, size }
    }

    pub fn from_min_max_vertices(point1: Point<T, D>, point2: Point<T, D>) -> Self
    where
        T: core::ops::SubAssign,
    {
        let size = point2 - point1;
        Self::new(point1, SVector::from(size))
    }

    /// Only considers the points closest and furthest from origin
    /// Points which are rotated along in the z axis (in 2d) are not considered
    pub fn from_vertices(points: [Point<T, D>; 4]) -> Option<Self>
    where
        T: core::ops::SubAssign,
        T: PartialOrd,
    {
        // find the closest and farthest points from the origin
        let min = points
            .iter()
            .reduce(|acc, p| (acc > p).then_some(p).unwrap_or(acc))?;
        let max = points
            .iter()
            .reduce(|acc, p| (acc < p).then_some(p).unwrap_or(acc))?;
        Some(Self::from_min_max_vertices(*min, *max))
    }

    pub fn size(&self) -> SVector<T, D> {
        self.size
    }

    pub fn center(&self) -> Point<T, D>
    where
        T: core::ops::AddAssign,
        T: core::ops::DivAssign,
    {
        self.point + self.size / (T::one() + T::one())
    }

    pub fn padding(mut self, padding: T) -> Self
    where
        T: core::ops::AddAssign,
        T: core::ops::DivAssign,
        T: core::ops::SubAssign,
    {
        self.size.iter_mut().for_each(|s| {
            *s += padding;
        });
        let two = T::one() + T::one();
        self.point
            .coords
            .iter_mut()
            .for_each(|c| *c -= padding / two);
        self
    }

    pub fn translate(&mut self, translation: SVector<T, D>)
    where
        T: core::ops::AddAssign,
    {
        self.point += translation;
    }

    pub fn min_vertex(&self) -> Point<T, D> {
        self.point
    }

    pub fn max_vertex(&self) -> Point<T, D>
    where
        T: core::ops::AddAssign,
    {
        self.point + self.size
    }

    pub fn contains_point(&self, point: &Point<T, D>) -> bool
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
    {
        let min = self.min_vertex();
        let max = self.max_vertex();

        *point >= min && *point <= max
    }

    pub fn scale(self, vector: SVector<T, D>) -> Self
    where
        T: core::ops::MulAssign,
        T: core::ops::DivAssign,
        T: core::ops::SubAssign,
    {
        let two = T::one() + T::one();
        let new_size = self.size.component_mul(&vector);
        let new_point = self.point.coords - new_size / two;
        Self {
            point: Point::from(new_point),
            size: new_size,
        }
    }

    pub fn contains_bbox(&self, other: &Self) -> bool
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
    {
        let self_min = self.min_vertex();
        let self_max = self.max_vertex();
        let other_min = other.min_vertex();
        let other_max = other.max_vertex();

        other_min >= self_min && other_max <= self_max
    }

    pub fn clamp(&self, other: &Self) -> Option<Self>
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
        T: nalgebra::SimdPartialOrd,
        T: nalgebra::SimdValue,
    {
        if other.contains_bbox(self) {
            return Some(*self);
        }
        self.intersection(other)
    }

    pub fn union(&self, other: &Self) -> Self
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
        T: nalgebra::SimdValue,
        T: nalgebra::SimdPartialOrd,
    {
        let self_min = self.min_vertex();
        let self_max = self.max_vertex();
        let other_min = other.min_vertex();
        let other_max = other.max_vertex();
        let min = self_min.inf(&other_min);
        let max = self_max.sup(&other_max);
        Self::from_min_max_vertices(min, max)
    }

    pub fn intersection(&self, other: &Self) -> Option<Self>
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
        T: nalgebra::SimdPartialOrd,
        T: nalgebra::SimdValue,
    {
        let self_min = self.min_vertex();
        let self_max = self.max_vertex();
        let other_min = other.min_vertex();
        let other_max = other.max_vertex();

        if self_max < other_min || other_max < self_min {
            return None; // No intersection
        }

        let min = self_min.sup(&other_min);
        let max = self_max.inf(&other_max);
        Some(Self::from_min_max_vertices(
            Point::from(min),
            Point::from(max),
        ))
    }

    pub fn denormalize(&self, factor: nalgebra::SVector<T, D>) -> Self
    where
        T: core::ops::MulAssign,
        T: core::ops::AddAssign,
        // nalgebra::constraint::ShapeConstraint:
        //     nalgebra::constraint::DimEq<nalgebra::Const<D>, nalgebra::Const<D>>,
    {
        Self {
            point: (self.point.coords.component_mul(&factor)).into(),
            size: self.size.component_mul(&factor),
        }
    }

    pub fn try_cast<T2>(&self) -> Option<Aabb<T2, D>>
    where
        // T: num::NumCast,
        T2: Num + simba::scalar::SubsetOf<T>,
    {
        Some(Aabb {
            point: Point::from(self.point.coords.try_cast::<T2>()?),
            size: self.size.try_cast::<T2>()?,
        })
    }

    // pub fn as_<T2>(&self) -> Option<Aabb<T2, D>>
    // where
    //     T2: Num + simba::scalar::SubsetOf<T>,
    // {
    //     Some(Aabb {
    //         point: Point::from(self.point.coords.as_()),
    //         size: self.size.as_(),
    //     })
    // }
    pub fn measure(&self) -> T
    where
        T: core::ops::MulAssign,
    {
        self.size.product()
    }

    pub fn iou(&self, other: &Self) -> Option<T>
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
        T: nalgebra::SimdPartialOrd,
        T: nalgebra::SimdValue,
        T: core::ops::MulAssign,
    {
        let intersection = self.intersection(other)?;
        let union = self.union(other);
        Some(intersection.measure() / union.measure())
    }
}

impl<T: Num> Aabb2<T> {
    pub fn from_x1y1x2y2(x1: T, y1: T, x2: T, y2: T) -> Self
    where
        T: core::ops::SubAssign,
    {
        let point1 = Point2::new(x1, y1);
        let point2 = Point2::new(x2, y2);
        Self::from_min_max_vertices(point1, point2)
    }
    pub fn new_2d(point1: Point2<T>, point2: Point2<T>) -> Self
    where
        T: core::ops::SubAssign,
    {
        let size = point2.coords - point1.coords;
        Self::new(point1, SVector::from(size))
    }
    pub fn x1y1(&self) -> Point2<T> {
        self.point
    }

    pub fn x2y2(&self) -> Point2<T>
    where
        T: core::ops::AddAssign,
    {
        self.point + self.size
    }

    pub fn x2y1(&self) -> Point2<T>
    where
        T: core::ops::AddAssign,
    {
        Point2::new(self.point.x + self.size.x, self.point.y)
    }

    pub fn x1y2(&self) -> Point2<T>
    where
        T: core::ops::AddAssign,
    {
        Point2::new(self.point.x, self.point.y + self.size.y)
    }

    pub fn x1(&self) -> T {
        self.point.x
    }

    pub fn y1(&self) -> T {
        self.point.y
    }

    pub fn x2(&self) -> T
    where
        T: core::ops::AddAssign,
    {
        self.point.x + self.size.x
    }

    pub fn y2(&self) -> T
    where
        T: core::ops::AddAssign,
    {
        self.point.y + self.size.y
    }

    pub fn corners(&self) -> [Point2<T>; 4]
    where
        T: core::ops::AddAssign,
    {
        [self.x1y1(), self.x2y1(), self.x2y2(), self.x1y2()]
    }

    pub fn area(&self) -> T
    where
        T: core::ops::MulAssign,
    {
        self.measure()
    }
}

impl<T: Num> Aabb3<T> {
    pub fn new_3d(point1: Point3<T>, point2: Point3<T>) -> Self
    where
        T: core::ops::SubAssign,
    {
        let size = point2.coords - point1.coords;
        Self::new(point1, SVector::from(size))
    }

    pub fn volume(&self) -> T
    where
        T: core::ops::MulAssign,
    {
        self.measure()
    }
}

#[test]
fn test_bbox_new() {
    use nalgebra::{Point2, Vector2};

    let point1 = Point2::new(1.0, 2.0);
    let point2 = Point2::new(4.0, 6.0);
    let bbox = AxisAlignedBoundingBox::new_2d(point1, point2);

    assert_eq!(bbox.min_vertex(), point1);
    assert_eq!(bbox.size(), Vector2::new(3.0, 4.0));
    assert_eq!(bbox.center(), Point2::new(2.5, 4.0));
}

#[test]
fn test_bounding_box_center_2d() {
    use nalgebra::{Point2, Vector2};

    let point = Point2::new(1.0, 2.0);
    let size = Vector2::new(3.0, 4.0);
    let bbox = AxisAlignedBoundingBox::new(point, size);

    assert_eq!(bbox.min_vertex(), point);
    assert_eq!(bbox.size(), size);
    assert_eq!(bbox.center(), Point2::new(2.5, 4.0));
}

#[test]
fn test_bounding_box_center_3d() {
    use nalgebra::{Point3, Vector3};

    let point = Point3::new(1.0, 2.0, 3.0);
    let size = Vector3::new(4.0, 5.0, 6.0);
    let bbox = AxisAlignedBoundingBox::new(point, size);

    assert_eq!(bbox.min_vertex(), point);
    assert_eq!(bbox.size(), size);
    assert_eq!(bbox.center(), Point3::new(3.0, 4.5, 6.0));
}

#[test]
fn test_bounding_box_padding_2d() {
    use nalgebra::{Point2, Vector2};

    let point = Point2::new(1.0, 2.0);
    let size = Vector2::new(3.0, 4.0);
    let bbox = AxisAlignedBoundingBox::new(point, size);

    let padded_bbox = bbox.padding(1.0);
    assert_eq!(padded_bbox.min_vertex(), Point2::new(0.5, 1.5));
    assert_eq!(padded_bbox.size(), Vector2::new(4.0, 5.0));
}

#[test]
fn test_bounding_box_scaling_2d() {
    use nalgebra::{Point2, Vector2};

    let point = Point2::new(1.0, 1.0);
    let size = Vector2::new(3.0, 4.0);
    let bbox = AxisAlignedBoundingBox::new(point, size);

    let padded_bbox = bbox.scale(Vector2::new(2.0, 2.0));
    assert_eq!(padded_bbox.min_vertex(), Point2::new(-2.0, -3.0));
    assert_eq!(padded_bbox.size(), Vector2::new(6.0, 8.0));
}

#[test]
fn test_bounding_box_contains_2d() {
    use nalgebra::Point2;

    let point1 = Point2::new(1.0, 2.0);
    let point2 = Point2::new(4.0, 6.0);
    let bbox = AxisAlignedBoundingBox::new_2d(point1, point2);

    assert!(bbox.contains_point(&Point2::new(2.0, 3.0)));
    assert!(!bbox.contains_point(&Point2::new(5.0, 7.0)));
}

#[test]
fn test_bounding_box_union_2d() {
    use nalgebra::{Point2, Vector2};

    let point1 = Point2::new(1.0, 2.0);
    let point2 = Point2::new(4.0, 6.0);
    let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2);

    let point3 = Point2::new(3.0, 5.0);
    let point4 = Point2::new(7.0, 8.0);
    let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4);

    let union_bbox = bbox1.union(&bbox2);
    assert_eq!(union_bbox.min_vertex(), Point2::new(1.0, 2.0));
    assert_eq!(union_bbox.size(), Vector2::new(6.0, 6.0));
}

#[test]
fn test_bounding_box_intersection_2d() {
    use nalgebra::{Point2, Vector2};

    let point1 = Point2::new(1.0, 2.0);
    let point2 = Point2::new(4.0, 6.0);
    let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2);

    let point3 = Point2::new(3.0, 5.0);
    let point4 = Point2::new(5.0, 7.0);
    let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4);

    let intersection_bbox = bbox1.intersection(&bbox2).unwrap();
    assert_eq!(intersection_bbox.min_vertex(), Point2::new(3.0, 5.0));
    assert_eq!(intersection_bbox.size(), Vector2::new(1.0, 1.0));
}

#[test]
fn test_bounding_box_contains_point() {
    use nalgebra::Point2;

    let point1 = Point2::new(2, 3);
    let point2 = Point2::new(5, 4);
    let bbox = AxisAlignedBoundingBox::new_2d(point1, point2);
    use itertools::Itertools;
    for (i, j) in (0..=10).cartesian_product(0..=10) {
        if bbox.contains_point(&Point2::new(i, j)) {
            if !(2..=5).contains(&i) && !(3..=4).contains(&j) {
                panic!(
                    "Point ({}, {}) should not be contained in the bounding box",
                    i, j
                );
            }
        } else {
            if (2..=5).contains(&i) && (3..=4).contains(&j) {
                panic!(
                    "Point ({}, {}) should be contained in the bounding box",
                    i, j
                );
            }
        }
    }
}

#[test]
fn test_bounding_box_clamp_box_2d() {
    let bbox1 = Aabb2::from_x1y1x2y2(1, 1, 4, 4);
    let bbox2 = Aabb2::from_x1y1x2y2(2, 2, 3, 3);
    let clamped = bbox2.clamp(&bbox1).unwrap();
    assert_eq!(bbox2, clamped);
    let clamped = bbox1.clamp(&bbox2).unwrap();
    assert_eq!(bbox2, clamped);

    let bbox1 = Aabb2::from_x1y1x2y2(4, 5, 7, 8);
    let bbox2 = Aabb2::from_x1y1x2y2(5, 4, 8, 7);
    let clamped = bbox1.clamp(&bbox2).unwrap();
    let expected = Aabb2::from_x1y1x2y2(5, 5, 7, 7);
    assert_eq!(clamped, expected)
}