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

use nalgebra::{Point, Point2, SVector, Vector2};
pub trait Num:
    num::Num
    + core::ops::AddAssign
    + core::ops::SubAssign
    + core::ops::MulAssign
    + core::ops::DivAssign
    + core::cmp::PartialOrd
    + core::cmp::PartialEq
    + nalgebra::SimdPartialOrd
    + nalgebra::SimdValue
    + Copy
    + core::fmt::Debug
    + 'static
{
}
impl<
    T: num::Num
        + core::ops::AddAssign
        + core::ops::SubAssign
        + core::ops::MulAssign
        + core::ops::DivAssign
        + core::cmp::PartialOrd
        + core::cmp::PartialEq
        + nalgebra::SimdPartialOrd
        + nalgebra::SimdValue
        + 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> {
    // Panics if max < min
    pub fn new(min_point: Point<T, D>, max_point: Point<T, D>) -> Self {
        if max_point >= min_point {
            Self::from_min_max_vertices(min_point, max_point)
        } else {
            panic!("max_point must be greater than or equal to min_point");
        }
    }
    pub fn try_new(min_point: Point<T, D>, max_point: Point<T, D>) -> Option<Self> {
        if max_point < min_point {
            return None;
        }
        Some(Self::from_min_max_vertices(min_point, max_point))
    }
    pub fn new_point_size(point: Point<T, D>, size: SVector<T, D>) -> Self {
        Self { point, size }
    }

    pub fn from_min_max_vertices(min: Point<T, D>, max: Point<T, D>) -> Self {
        let size = max - min;
        Self::new_point_size(min, 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 component_clamp(&self, min: T, max: T) -> Self
    where
        T: PartialOrd,
    {
        let mut this = *self;
        this.point.iter_mut().for_each(|x| {
            *x = nalgebra::clamp(*x, min, max);
        });
        this.size.iter_mut().for_each(|x| {
            *x = nalgebra::clamp(*x, min, max);
        });
        this
    }

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

    pub fn union(&self, other: &Self) -> T
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: core::ops::MulAssign,
        T: PartialOrd,
        T: nalgebra::SimdValue,
        T: nalgebra::SimdPartialOrd,
    {
        self.measure() + other.measure()
            - Self::intersection(self, other)
                .map(|x| x.measure())
                .unwrap_or(T::zero())
    }

    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 inter_min = self.min_vertex().sup(&other.min_vertex());
        let inter_max = self.max_vertex().inf(&other.max_vertex());
        Self::try_new(inter_min, inter_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) -> T
    where
        T: core::ops::AddAssign,
        T: core::ops::SubAssign,
        T: PartialOrd,
        T: nalgebra::SimdPartialOrd,
        T: nalgebra::SimdValue,
        T: core::ops::MulAssign,
    {
        let lhs_min = self.min_vertex();
        let lhs_max = self.max_vertex();
        let rhs_min = other.min_vertex();
        let rhs_max = other.max_vertex();

        let inter_min = lhs_min.sup(&rhs_min);
        let inter_max = lhs_max.inf(&rhs_max);
        if inter_max >= inter_min {
            let intersection = Aabb::new(inter_min, inter_max).measure();
            intersection / (self.measure() + other.measure() - intersection)
        } else {
            return T::zero();
        }
    }
}

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::new(point1, point2)
    }

    pub fn from_xywh(x: T, y: T, w: T, h: T) -> Self {
        let point = Point2::new(x, y);
        let size = Vector2::new(w, h);
        Self::new_point_size(point, 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 volume(&self) -> T
    where
        T: core::ops::MulAssign,
    {
        self.measure()
    }
}

#[cfg(test)]
mod boudning_box_tests {
    use super::*;
    use nalgebra::*;

    #[test]
    fn test_bbox_new() {
        let point1 = Point2::new(1.0, 2.0);
        let point2 = Point2::new(4.0, 6.0);
        let bbox = AxisAlignedBoundingBox::new(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_intersection_and_merge() {
        let point1 = Point2::new(1, 5);
        let point2 = Point2::new(3, 2);
        let size1 = Vector2::new(3, 4);
        let size2 = Vector2::new(1, 3);

        let this = Aabb2::new_point_size(point1, size1);
        let other = Aabb2::new_point_size(point2, size2);
        let inter = this.intersection(&other);
        let merged = this.merge(&other);
        assert_ne!(inter, Some(merged))
    }

    #[test]
    fn test_bounding_box_center_2d() {
        let point = Point2::new(1.0, 2.0);
        let size = Vector2::new(3.0, 4.0);
        let bbox = AxisAlignedBoundingBox::new_point_size(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() {
        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(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() {
        let point = Point2::new(1.0, 2.0);
        let size = Vector2::new(3.0, 4.0);
        let bbox = AxisAlignedBoundingBox::new_point_size(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() {
        let point = Point2::new(1.0, 1.0);
        let size = Vector2::new(3.0, 4.0);
        let bbox = AxisAlignedBoundingBox::new_point_size(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() {
        let point1 = Point2::new(1.0, 2.0);
        let point2 = Point2::new(4.0, 6.0);
        let bbox = AxisAlignedBoundingBox::new(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() {
        let point1 = Point2::new(1.0, 2.0);
        let point2 = Point2::new(4.0, 6.0);
        let bbox1 = AxisAlignedBoundingBox::new(point1, point2);

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

        let union_bbox = bbox1.merge(&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() {
        let point1 = Point2::new(1.0, 2.0);
        let point2 = Point2::new(4.0, 6.0);
        let bbox1 = AxisAlignedBoundingBox::new(point1, point2);

        let point3 = Point2::new(3.0, 5.0);
        let point4 = Point2::new(5.0, 7.0);
        let bbox2 = AxisAlignedBoundingBox::new(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() {
        let point1 = Point2::new(2, 3);
        let point2 = Point2::new(5, 4);
        let bbox = AxisAlignedBoundingBox::new(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)
    }

    #[test]
    fn test_iou_identical_boxes() {
        let a = Aabb2::from_x1y1x2y2(1.0, 2.0, 4.0, 6.0);
        let b = Aabb2::from_x1y1x2y2(1.0, 2.0, 4.0, 6.0);
        assert_eq!(a.iou(&b), 1.0);
    }

    #[test]
    fn test_iou_non_overlapping_boxes() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 1.0, 1.0);
        let b = Aabb2::from_x1y1x2y2(2.0, 2.0, 3.0, 3.0);
        assert_eq!(a.iou(&b), 0.0);
    }

    #[test]
    fn test_iou_partial_overlap() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 2.0, 2.0);
        let b = Aabb2::from_x1y1x2y2(1.0, 1.0, 3.0, 3.0);
        // Intersection area = 1, Union area = 7
        assert!((a.iou(&b) - 1.0 / 7.0).abs() < 1e-6);
    }

    #[test]
    fn test_iou_one_inside_another() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 4.0, 4.0);
        let b = Aabb2::from_x1y1x2y2(1.0, 1.0, 3.0, 3.0);
        // Intersection area = 4, Union area = 16
        assert!((a.iou(&b) - 0.25).abs() < 1e-6);
    }

    #[test]
    fn test_iou_edge_touching() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 1.0, 1.0);
        let b = Aabb2::from_x1y1x2y2(1.0, 0.0, 2.0, 1.0);
        assert_eq!(a.iou(&b), 0.0);
    }

    #[test]
    fn test_iou_corner_touching() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 1.0, 1.0);
        let b = Aabb2::from_x1y1x2y2(1.0, 1.0, 2.0, 2.0);
        assert_eq!(a.iou(&b), 0.0);
    }

    #[test]
    fn test_iou_zero_area_box() {
        let a = Aabb2::from_x1y1x2y2(0.0, 0.0, 0.0, 0.0);
        let b = Aabb2::from_x1y1x2y2(0.0, 0.0, 1.0, 1.0);
        assert_eq!(a.iou(&b), 0.0);
    }

    #[test]
    fn test_specific_values() {
        let box1 = Aabb2::from_xywh(0.69482, 0.6716774, 0.07493961, 0.14968264);
        let box2 = Aabb2::from_xywh(0.41546485, 0.70290875, 0.06197411, 0.08818436);
        assert!(box1.iou(&box2) >= 0.0);
    }
}