crystal_ball 0.3.0

use std::convert::TryFrom;
use std::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};

use nanorand::tls::TlsWyRand;

use crate::color::Color;
use crate::math::{Point3, Vec2, Vec4};
use crate::util::random_float;

/// A 3-dimensional vector.
#[derive(Copy, Clone, Default, Debug, PartialEq)]
pub struct Vec3 {
    pub x: f64,
    pub y: f64,
    pub z: f64,
}

impl Vec3 {
    pub const X: Vec3 = Vec3 {
        x: 1.0,
        y: 0.0,
        z: 0.0,
    };

    pub const Y: Vec3 = Vec3 {
        x: 0.0,
        y: 1.0,
        z: 0.0,
    };

    pub const Z: Vec3 = Vec3 {
        x: 0.0,
        y: 0.0,
        z: 1.0,
    };

    pub const ZERO: Vec3 = Vec3 {
        x: 0.0,
        y: 0.0,
        z: 0.0,
    };

    /// Create a new [`Vec3`].
    pub fn new(x: f64, y: f64, z: f64) -> Self {
        Vec3 { x, y, z }
    }

    /// Create a new [`Vec3`] with all elements set to `value`.
    pub fn splat(value: f64) -> Self {
        Vec3 {
            x: value,
            y: value,
            z: value,
        }
    }

    /// Generate a [`Vec3`] where each component is a uniform random number between `min` and `max`.
    pub fn random(rng: &mut TlsWyRand, min: f64, max: f64) -> Self {
        Vec3 {
            x: random_float(rng, min, max),
            y: random_float(rng, min, max),
            z: random_float(rng, min, max),
        }
    }

    /// Generate a random [`Vec3`] on the surface of the unit sphere
    /// with uniform (i.e. isotropic) distribution.
    pub fn random_unit_vector(rng: &mut TlsWyRand) -> Self {
        Self::random_in_unit_sphere(rng).normalize()
    }

    /// Generate a random [`Vec3`] inside the unit sphere with uniform distribution.
    pub fn random_in_unit_sphere(rng: &mut TlsWyRand) -> Self {
        loop {
            let random_vec3 = Self::random(rng, -1.0, 1.0);
            if random_vec3.magnitude() < 1.0 {
                return random_vec3;
            }
        }
    }

    /// Reflect the `self` with respect to the given normal.
    pub fn reflect(self, normal: Vec3) -> Self {
        self - 2.0 * Self::dot(self, normal) * normal
    }

    /// Refract the `self` with respect to the given normal using n1 and n2 as the indices of refraction.
    pub fn refract(self, normal: Vec3, n1: f64, n2: f64, rng: &mut TlsWyRand) -> Self {
        let n1_inverse = 1.0 / n1;

        let incidence_angle = (-Self::dot(self, normal)).acos();

        // Total internal reflection.
        if incidence_angle >= (n2 * n1_inverse).asin()
            || Self::reflection_coefficient(incidence_angle, n1, n2) > random_float(rng, 0.0, 1.0)
        {
            return self.reflect(normal);
        }

        let refraction_angle = (n1 * incidence_angle.sin() / n2).asin();

        self - normal * (n2 * refraction_angle.cos() - n1 * incidence_angle.cos()) * n1_inverse
    }

    /// Calculate the reflection coefficient using Schlick's approximation.
    fn reflection_coefficient(incidence_angle: f64, n1: f64, n2: f64) -> f64 {
        let normal_reflectivity = ((n1 - n2) / (n1 + n2)).powi(2);
        normal_reflectivity + (1.0 - normal_reflectivity) * (1.0 - incidence_angle.cos()).powi(5)
    }

    /// Convert the [`Vec2`] to a [`Point2`](crate::math::Point2).
    pub fn to_point3(&self) -> Point3 {
        (*self).into()
    }

    /// Calculate the vector's magnitude (length).
    #[doc(alias = "length")]
    pub fn magnitude(&self) -> f64 {
        Self::dot(*self, *self).sqrt()
    }

    /// Calculate the square of the vector's magnitude (length).
    #[doc(alias = "length_squared")]
    pub fn magnitude_squared(&self) -> f64 {
        Self::dot(*self, *self)
    }

    /// Return the absolute value (synonym [`magnitude`](Self::magnitude)).
    pub fn abs(&self) -> f64 {
        self.magnitude()
    }

    /// Return the unit vector parallel to self.
    ///
    /// # Panics
    ///
    /// Panics if `self` cannot be normalized.
    pub fn normalize(&self) -> Self {
        assert_ne!(self.magnitude(), 0.0, "Can't normalize zero vector");

        *self / self.magnitude()
    }

    /// Calculate the dot product between two [`Vec3`]s.
    pub fn dot(vec_a: Vec3, vec_b: Vec3) -> f64 {
        vec_a.x * vec_b.x + vec_a.y * vec_b.y + vec_a.z * vec_b.z
    }

    /// Calculate the cross product between two [`Vec3`]s.
    pub fn cross(vec_a: Vec3, vec_b: Vec3) -> Self {
        Vec3 {
            x: vec_a.y * vec_b.z - vec_a.z * vec_b.y,
            y: vec_a.z * vec_b.x - vec_a.x * vec_b.z,
            z: vec_a.x * vec_b.y - vec_a.y * vec_b.x,
        }
    }

    /// Create a new [`Vec4`] using the x, y and z values of `self` and the given `w` value.
    pub fn extend(&self, w: f64) -> Vec4 {
        Vec4::new(self.x, self.y, self.z, w)
    }

    /// Create a new [`Vec2`] discarding the `z` value.
    pub fn xyz(&self) -> Vec2 {
        Vec2::new(self.x, self.y)
    }
}

impl Add<Vec3> for Vec3 {
    type Output = Vec3;

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

impl AddAssign<Vec3> for Vec3 {
    fn add_assign(&mut self, rhs: Vec3) {
        *self = *self + rhs;
    }
}

impl Sub<Vec3> for Vec3 {
    type Output = Vec3;

    fn sub(self, rhs: Vec3) -> Self::Output {
        Vec3 {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
            z: self.z - rhs.z,
        }
    }
}

impl SubAssign<Vec3> for Vec3 {
    fn sub_assign(&mut self, rhs: Vec3) {
        *self = *self - rhs;
    }
}

impl Mul<f64> for Vec3 {
    type Output = Vec3;

    fn mul(self, rhs: f64) -> Self::Output {
        Vec3 {
            x: self.x * rhs,
            y: self.y * rhs,
            z: self.z * rhs,
        }
    }
}

impl MulAssign<f64> for Vec3 {
    fn mul_assign(&mut self, rhs: f64) {
        *self = *self * rhs;
    }
}

impl Mul<Vec3> for f64 {
    type Output = Vec3;

    fn mul(self, rhs: Vec3) -> Self::Output {
        rhs * self
    }
}

impl Div<f64> for Vec3 {
    type Output = Vec3;

    fn div(self, rhs: f64) -> Self::Output {
        let rhs_inverse = 1.0 / rhs;

        Vec3 {
            x: self.x * rhs_inverse,
            y: self.y * rhs_inverse,
            z: self.z * rhs_inverse,
        }
    }
}

impl DivAssign<f64> for Vec3 {
    fn div_assign(&mut self, rhs: f64) {
        *self = *self / rhs;
    }
}

impl Neg for Vec3 {
    type Output = Vec3;

    fn neg(self) -> Self::Output {
        Vec3 {
            x: -self.x,
            y: -self.y,
            z: -self.z,
        }
    }
}

impl From<[f64; 3]> for Vec3 {
    fn from(s: [f64; 3]) -> Self {
        Vec3 {
            x: s[0],
            y: s[1],
            z: s[2],
        }
    }
}

impl From<(f64, f64, f64)> for Vec3 {
    fn from(t: (f64, f64, f64)) -> Self {
        Vec3 {
            x: t.0,
            y: t.1,
            z: t.2,
        }
    }
}

impl From<Point3> for Vec3 {
    fn from(p: Point3) -> Self {
        Vec3 {
            x: p.x,
            y: p.y,
            z: p.z,
        }
    }
}

impl From<Color> for Vec3 {
    fn from(c: Color) -> Self {
        Vec3 {
            x: c.r as f64,
            y: c.g as f64,
            z: c.b as f64,
        }
    }
}

impl TryFrom<Vec<f64>> for Vec3 {
    type Error = &'static str;

    fn try_from(v: Vec<f64>) -> Result<Self, Self::Error> {
        if v.len() != 3 {
            Err("Vec3 can only be build from a vector of length 3.")
        } else {
            Ok(Vec3 {
                x: v[0],
                y: v[1],
                z: v[2],
            })
        }
    }
}

impl TryFrom<&[f64]> for Vec3 {
    type Error = &'static str;

    fn try_from(s: &[f64]) -> Result<Self, Self::Error> {
        if s.len() != 3 {
            Err("Vec3 can only be build from a slice of length 3.")
        } else {
            Ok(Vec3 {
                x: s[0],
                y: s[1],
                z: s[2],
            })
        }
    }
}

impl Index<usize> for Vec3 {
    type Output = f64;

    fn index(&self, index: usize) -> &Self::Output {
        match index {
            0 => &self.x,
            1 => &self.y,
            2 => &self.z,
            _ => panic!(
                "index out of bounds: the len is 3 but the index is {}",
                index
            ),
        }
    }
}

impl IndexMut<usize> for Vec3 {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        match index {
            0 => &mut self.x,
            1 => &mut self.y,
            2 => &mut self.z,
            _ => panic!(
                "index out of bounds: the len is 3 but the index is {}",
                index
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use assert_approx_eq::assert_approx_eq;
    use nanorand::tls_rng;

    use crate::util::EPSILON_F64;

    use super::Vec3;

    #[test]
    fn vec3_add() {
        assert_approx_eq!(
            Vec3::new(2.0, 1.0, 0.0) + Vec3::new(1.0, 1.0, 1.0),
            Vec3::new(3.0, 2.0, 1.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::new(5.72, 2.5, 8.824) + Vec3::new(8.7, 5.987, 0.12),
            Vec3::new(14.42, 8.487, 8.944),
            EPSILON_F64
        );
        let mut vec_a = Vec3::new(7.0, 2.5, 3.2);
        vec_a += Vec3::new(1.2, 9.23, 6.2);
        assert_approx_eq!(vec_a, Vec3::new(8.2, 11.73, 9.4), EPSILON_F64)
    }

    #[test]
    fn vec3_sub() {
        assert_approx_eq!(
            Vec3::new(2.0, 1.0, 0.0) - Vec3::new(1.0, 1.0, 1.0),
            Vec3::new(1.0, 0.0, -1.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::new(5.72, 2.5, 8.824) - Vec3::new(8.7, 5.987, 0.12),
            Vec3::new(-2.98, -3.487, 8.704),
            EPSILON_F64
        );
        let mut vec_a = Vec3::new(7.0, 2.5, 3.2);
        vec_a -= Vec3::new(1.2, 9.23, 6.2);
        assert_approx_eq!(vec_a, Vec3::new(5.8, -6.73, -3.0), EPSILON_F64)
    }

    #[test]
    fn vec3_mul() {
        assert_approx_eq!(
            Vec3::new(2.0, 1.0, 0.0) * 2.0,
            Vec3::new(4.0, 2.0, 0.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            2.5 * Vec3::new(8.7, 5.987, 0.12),
            Vec3::new(21.75, 14.9675, 0.3),
            EPSILON_F64
        );
        let mut vec_a = Vec3::new(7.0, 2.5, 3.2);
        vec_a *= -2.0;
        assert_approx_eq!(vec_a, Vec3::new(-14.0, -5.0, -6.4), EPSILON_F64);
    }

    #[test]
    fn vec3_div() {
        assert_approx_eq!(
            Vec3::new(2.0, 1.0, 0.0) / 2.0,
            Vec3::new(1.0, 0.5, 0.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::new(8.7, 5.987, 0.12) / 2.5,
            Vec3::new(3.48, 2.3948, 0.048),
            EPSILON_F64
        );
        let mut vec_a = Vec3::new(7.0, 2.5, 3.2);
        vec_a /= -2.0;
        assert_approx_eq!(vec_a, Vec3::new(-3.5, -1.25, -1.6), EPSILON_F64);
    }

    #[test]
    fn vec3_neg() {
        assert_approx_eq!(
            -Vec3::new(2.0, 1.0, 0.0),
            Vec3::new(-2.0, -1.0, 0.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            -Vec3::new(8.7, 5.987, 0.12),
            Vec3::new(-8.7, -5.987, -0.12),
            EPSILON_F64
        );
    }

    #[test]
    fn vec3_random() {
        let mut rng = tls_rng();
        assert_approx_eq!(
            Vec3::random_unit_vector(&mut rng).magnitude(),
            1.0,
            EPSILON_F64
        );
        assert!(Vec3::random_in_unit_sphere(&mut rng).magnitude() < 1.0);
    }

    #[test]
    fn vec3_reflect() {
        assert_approx_eq!(
            Vec3::new(2.0, 1.0, 0.0).reflect(Vec3::new(0.0, 1.0, 0.0)),
            Vec3::new(2.0, -1.0, 0.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::new(8.7, 5.987, 0.12).reflect(Vec3::new(1.0, 1.0, 0.0)),
            Vec3::new(-20.674, -23.387, 0.12),
            EPSILON_F64
        );
    }

    #[test]
    fn vec3_normalize() {
        assert_approx_eq!(
            (-Vec3::new(2.0, 1.0, 0.0)).normalize().magnitude(),
            1.0,
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::new(8.7, 5.987, 0.12).normalize().magnitude(),
            1.0,
            EPSILON_F64
        );
    }

    #[test]
    #[should_panic(expected = "Can't normalize zero vector")]
    fn vec3_normalize_panic() {
        Vec3::new(0.0, 0.0, 0.0).normalize();
    }

    #[test]
    fn vec3_magnitude() {
        assert_approx_eq!(Vec3::new(42.0, 0.0, 0.0).magnitude(), 42.0, EPSILON_F64);
        assert_approx_eq!(Vec3::new(-3.0, -4.0, 0.0).magnitude(), 5.0, EPSILON_F64);
        assert_approx_eq!(Vec3::new(2.0, -2.0, 1.0).abs(), 3.0, EPSILON_F64);
    }

    #[test]
    fn vec3_dot() {
        assert_approx_eq!(
            Vec3::dot(Vec3::new(2.0, 1.0, 0.0), Vec3::new(1.0, 1.0, 1.0)),
            3.0,
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::dot(Vec3::new(5.72, 2.5, 8.824), Vec3::new(8.7, 5.987, 0.12)),
            65.79038,
            EPSILON_F64
        );
    }

    #[test]
    fn vec3_cross() {
        assert_approx_eq!(
            Vec3::cross(Vec3::new(2.0, 1.0, 0.0), Vec3::new(1.0, 1.0, 1.0)),
            Vec3::new(1.0, -2.0, 1.0),
            EPSILON_F64
        );
        assert_approx_eq!(
            Vec3::cross(Vec3::new(5.72, 2.5, 8.824), Vec3::new(2.0, 1.0, 0.0)),
            Vec3::new(-8.824, 17.648, 0.72),
            EPSILON_F64
        );
    }
}