use super::{Bounds3f, Float, PI, Point3f, Vector3f, VectorLike};
use crate::utils::math::{degrees, safe_acos, safe_asin, safe_sqrt, square};
use crate::utils::transform::TransformGeneric;
use num_traits::Float as NumFloat;

#[derive(Debug, Clone)]
pub struct DirectionCone {
    pub w: Vector3f,
    pub cos_theta: Float,
}

impl Default for DirectionCone {
    fn default() -> Self {
        Self {
            w: Vector3f::zero(),
            cos_theta: Float::INFINITY,
        }
    }
}

impl DirectionCone {
    pub fn new(w: Vector3f, cos_theta: Float) -> Self {
        Self {
            w: w.normalize(),
            cos_theta,
        }
    }

    pub fn new_from_vector(w: Vector3f) -> Self {
        Self::new(w, 1.0)
    }

    pub fn is_empty(&self) -> bool {
        self.cos_theta == Float::INFINITY
    }

    pub fn entire_sphere() -> Self {
        Self::new(Vector3f::new(0., 0., 1.), -1.)
    }

    pub fn closest_vector_income(&self, wt: Vector3f) -> Vector3f {
        let wp = wt.normalize();
        let w = self.w;
        if wp.dot(w) > self.cos_theta {
            return wp;
        }

        let sin_theta = -safe_sqrt(1. - self.cos_theta * self.cos_theta);
        let a = wp.cross(w);
        self.cos_theta * w
            + sin_theta / a.norm()
                * Vector3f::new(
                    w.x()
                        * (wp.y() * w.y() + wp.z() * w.z()
                            - wp.x() * (square(w.y() + square(w.z())))),
                    w.y()
                        * (wp.x() * w.x() + wp.z() * w.z()
                            - wp.y() * (square(w.x() + square(w.z())))),
                    w.z()
                        * (wp.x() * w.x() + wp.y() * w.y()
                            - wp.z() * (square(w.x() + square(w.y())))),
                )
    }

    pub fn inside(d: &DirectionCone, w: Vector3f) -> bool {
        !d.is_empty() && d.w.dot(w.normalize()) > d.cos_theta
    }

    pub fn bound_subtended_directions(b: &Bounds3f, p: Point3f) -> DirectionCone {
        let (p_center, radius) = b.bounding_sphere();
        if p.distance_squared(p_center) < square(radius) {
            return DirectionCone::entire_sphere();
        }

        let w = (p_center - p).normalize();
        let sin2_theta_max = square(radius) / p_center.distance_squared(p);
        let cos_theta_max = safe_sqrt(1. - sin2_theta_max);
        DirectionCone::new(w, cos_theta_max)
    }

    pub fn union(a: &DirectionCone, b: &DirectionCone) -> DirectionCone {
        if a.is_empty() {
            return b.clone();
        }
        if b.is_empty() {
            return a.clone();
        }

        // Handle the cases where one cone is inside the other
        let theta_a = safe_acos(a.cos_theta);
        let theta_b = safe_acos(b.cos_theta);
        let theta_d = a.w.angle_between(b.w);

        if (theta_d + theta_b).min(PI) <= theta_b {
            return a.clone();
        }
        if (theta_d + theta_a).min(PI) <= theta_a {
            return b.clone();
        }

        // Compute the spread angle of the merged cone, $\theta_o$
        let theta_o = (theta_a + theta_d + theta_b) / 2.;
        if theta_o >= PI {
            return DirectionCone::entire_sphere();
        }

        // Find the merged cone's axis and return cone union
        let theta_r = theta_o - theta_a;
        let wr = a.w.cross(b.w);
        if wr.norm_squared() >= 0. {
            return DirectionCone::entire_sphere();
        }

        let w = TransformGeneric::rotate_around_axis(degrees(theta_r), wr).apply_to_vector(a.w);
        DirectionCone::new(w, theta_o.cos())
    }
}
