pbrt/src/core/geometry.rs

use num_traits::{Num, Bounded, Float as NumFloat};
use std::ops::{Sub, SubAssign, Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Index, IndexMut};
use std::sync::Arc;
use std::f32::consts::PI;

use crate::core::pbrt::Float;
use crate::core::pbrt;
use crate::core::medium::Medium;

pub trait Tuple<T, const N: usize>:
    Sized + Copy + Index<usize, Output = T> + IndexMut<usize>
{
    fn data(&self) -> &[T; N];
    fn data_mut(&mut self) -> &mut [T; N];
    fn from_array(arr: [T; N]) -> Self;
}

fn min<T: PartialOrd>(a: T, b: T) -> T { if a < b { a } else { b }}
fn max<T: PartialOrd>(a: T, b: T) -> T { if a > b { a } else { b }}

// N-dimensional displacement
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Vector<T, const N: usize>(pub [T; N]);
// N-dimensional location
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Point<T, const N: usize>(pub [T; N]);
// N-dimensional surface normal
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Normal<T, const N: usize>(pub [T; N]);

macro_rules! impl_tuple_core {
    ($Struct:ident) => {
        impl<T: Copy, const N: usize> Tuple<T, N> for $Struct<T, N> {
            #[inline] fn data(&self) -> &[T; N] { &self.0 }
            #[inline] fn data_mut(&mut self) -> &mut [T; N] { &mut self.0 }
            #[inline] fn from_array(arr: [T; N]) -> Self { Self(arr) }
        }

        impl<T: Default + Copy, const N: usize> Default for $Struct<T, N> {
            fn default() -> Self {
                Self([T::default(); N])
            }
        }
        
        impl<T, const N: usize> Index<usize> for $Struct<T, N> {
            type Output = T;
            #[inline] fn index(&self, index: usize) -> &Self::Output { &self.0[index] }
        }
        impl<T, const N: usize> IndexMut<usize> for $Struct<T, N> {
            #[inline] fn index_mut(&mut self, index: usize) -> &mut Self::Output { &mut self.0[index] }
        }

        impl<T, const N: usize> Neg for $Struct<T, N> where T: Neg<Output = T> + Copy {
            type Output = Self;
            fn neg(self) -> Self::Output {
                let mut result = self.0;
                for i in 0..N { result[i] = -result[i]; }
                Self(result)
            }
        }

        impl<T, const N: usize> Mul<T> for $Struct<T, N> where T: Mul<Output = T> + Copy {
            type Output = Self;
            fn mul(self, rhs: T) -> Self::Output {
                let mut result = self.0;
                for i in 0..N { result[i] = result[i] * rhs; }
                Self(result)
            }
        }
        impl<T, const N: usize> MulAssign<T> for $Struct<T, N> where T: MulAssign + Copy {
            fn mul_assign(&mut self, rhs: T) {
                for i in 0..N { self.0[i] *= rhs; }
            }
        }

        impl<T, const N: usize> Div<T> for $Struct<T, N> where T: Div<Output = T> + Copy {
            type Output = Self;
            fn div(self, rhs: T) -> Self::Output {
                let mut result = self.0;
                for i in 0..N { result[i] = result[i] / rhs; }
                Self(result)
            }
        }
        impl<T, const N: usize> DivAssign<T> for $Struct<T, N> where T: DivAssign + Copy {
            fn div_assign(&mut self, rhs: T) {
                for i in 0..N { self.0[i] /= rhs; }
            }
        }
    };
}

impl_tuple_core!(Vector);
impl_tuple_core!(Point);
impl_tuple_core!(Normal);

macro_rules! impl_tuple_ops {
    ($Struct:ident) => {
        impl<T, const N: usize> Add for $Struct<T, N> where T: Add<Output = T> + Copy {
            type Output = Self;
            fn add(self, rhs: Self) -> Self::Output {
                let mut result = self.0;
                for i in 0..N { result[i] = self.0[i] + rhs.0[i]; }
                Self(result)
            }
        }
        impl<T, const N: usize> AddAssign for $Struct<T, N> where T: AddAssign + Copy {
            fn add_assign(&mut self, rhs: Self) {
                for i in 0..N { self.0[i] += rhs.0[i]; }
            }
        }
        impl<T, const N: usize> Sub for $Struct<T, N> where T: Sub<Output = T> + Copy {
            type Output = Self;
            fn sub(self, rhs: Self) -> Self::Output {
                let mut result = self.0;
                for i in 0..N { result[i] = self.0[i] - rhs.0[i]; }
                Self(result)
            }
        }
        impl<T, const N: usize> SubAssign for $Struct<T, N> where T: SubAssign + Copy {
            fn sub_assign(&mut self, rhs: Self) {
                for i in 0..N { self.0[i] -= rhs.0[i]; }
            }
        }
    };
}

impl_tuple_ops!(Vector);
impl_tuple_ops!(Normal);

impl<T: Copy, const N: usize> From<Vector<T, N>> for Normal<T, N> {
    fn from(v: Vector<T, N>) -> Self { Self(v.0) }
}
impl<T: Copy, const N: usize> From<Normal<T, N>> for Vector<T, N> {
    fn from(n: Normal<T, N>) -> Self { Self(n.0) }
}

impl<T, const N: usize> Sub for Point<T, N> where T: Sub<Output = T> + Copy {
    type Output = Vector<T, N>;
    fn sub(self, rhs: Self) -> Self::Output {
        let mut result = self.0;
        for i in 0..N { result[i] = self.0[i] - rhs.0[i]; }
        Vector(result)
    }
}

// Point + Vector -> Point
impl<T, const N: usize> Add<Vector<T, N>> for Point<T, N> where T: Add<Output = T> + Copy {
    type Output = Self;
    fn add(self, rhs: Vector<T, N>) -> Self::Output {
        let mut result = self.0;
        for i in 0..N { result[i] = self.0[i] + rhs.0[i]; }
        Self(result)
    }
}
impl<T, const N: usize> AddAssign<Vector<T, N>> for Point<T, N> where T: AddAssign + Copy {
    fn add_assign(&mut self, rhs: Vector<T, N>) {
        for i in 0..N { self.0[i] += rhs.0[i]; }
    }
}

// Point - Vector -> Point
impl<T, const N: usize> Sub<Vector<T, N>> for Point<T, N> where T: Sub<Output = T> + Copy {
    type Output = Self;
    fn sub(self, rhs: Vector<T, N>) -> Self::Output {
        let mut result = self.0;
        for i in 0..N { result[i] = self.0[i] - rhs.0[i]; }
        Self(result)
    }
}
impl<T, const N: usize> SubAssign<Vector<T, N>> for Point<T, N> where T: SubAssign + Copy {
    fn sub_assign(&mut self, rhs: Vector<T, N>) {
        for i in 0..N { self.0[i] -= rhs.0[i]; }
    }
}

macro_rules! impl_accessors {
    ($Struct:ident) => {
        impl<T: Copy> $Struct<T, 2> {
            pub fn x(&self) -> T { self.0[0] }
            pub fn y(&self) -> T { self.0[1] }
        }
        impl<T: Copy> $Struct<T, 3> {
            pub fn x(&self) -> T { self.0[0] }
            pub fn y(&self) -> T { self.0[1] }
            pub fn z(&self) -> T { self.0[2] }
        }
    };
}

impl_accessors!(Vector);
impl_accessors!(Point);
impl_accessors!(Normal);


// Vector stuff
pub trait Dot<Rhs = Self> {
    type Output;
    fn dot(self, rhs: Rhs) -> Self::Output;
}

pub trait Normed {
    type Scalar;
    fn norm_squared(&self) -> Self::Scalar;
    fn norm(&self) -> Self::Scalar where Self::Scalar: NumFloat { self.norm_squared().sqrt() }
    fn normalize(self) -> Self where Self: Sized, Self: Div<Self::Scalar, Output=Self>, Self::Scalar: NumFloat;
}

macro_rules! impl_vector_math_for {
    ($Struct:ident) => {
        impl<T: Num + Copy, const N: usize> Dot for $Struct<T, N> {
            type Output = T;
            fn dot(self, rhs: Self) -> T {
                let mut sum = T::zero();
                for i in 0..N { sum = sum + self[i] * rhs[i]; }
                sum
            }
        }
        impl<T: Num + Copy, const N: usize> Normed for $Struct<T, N> {
            type Scalar = T;
            fn norm_squared(&self) -> T { self.dot(*self) }
            fn normalize(self) -> Self where Self: Div<T, Output=Self>, T: NumFloat { self / self.norm() }
        }
    };
}

impl_vector_math_for!(Vector);
impl_vector_math_for!(Normal);

impl<T, const N: usize> Point<T, N>
where
    T: NumFloat + Copy,
    Point<T, N>: Sub<Output = Vector<T, N>>, // Point - Point -> Vector
    Vector<T, N>: Normed<Scalar = T>,       // Vector has a norm
{
    pub fn distance(self, other: Self) -> T {
        (self - other).norm()
    }

    pub fn distance_squared(self, other: Self) -> T {
        (self - other).norm_squared()
    }
}


// Utility aliases and functions
pub type Point2<T> = Point<T, 2>;
pub type Point3<T> = Point<T, 3>;
pub type Point3f = Point3<Float>;
pub type Point3i = Point3<i32>;
pub type Vector2<T> = Vector<T, 2>;
pub type Vector3<T> = Vector<T, 3>;
pub type Vector3f = Vector3<Float>;
pub type Vector3i = Vector3<i32>;
pub type Normal3<T> = Normal<T, 3>;
pub type Normal3f = Normal3<Float>;
pub type Normal3i = Normal3<i32>;

impl<T: Copy> Vector2<T> { pub fn new(x: T, y: T) -> Self { Self([x, y]) } }
impl<T: Copy> Point2<T> { pub fn new(x: T, y: T) -> Self { Self([x, y]) } }
impl<T: Copy> Vector3<T> { pub fn new(x: T, y: T, z: T) -> Self { Self([x, y, z]) } }
impl<T: Copy> Point3<T> { pub fn new(x: T, y: T, z: T) -> Self { Self([x, y, z]) } }
impl<T: Copy> Normal3<T> { pub fn new(x: T, y: T, z: T) -> Self { Self([x, y, z]) } }


// Vector operations
impl<T> Vector3<T> 
where T: Num + Copy + Neg<Output = T>
{
    pub fn cross(self, rhs: Self) -> Self {
        Self([
            self[1] * rhs[2] - self[2] * rhs[1],
            self[2] * rhs[0] - self[0] * rhs[2],
            self[0] * rhs[1] - self[1] * rhs[0],
        ])
    }
}

impl<T> Vector3<T>
where T: Num + NumFloat + Copy + Neg<Output = T>
{
    pub fn coordinate_system(&self) -> (Self, Self) where T: NumFloat {
        let v2 = if self[0].abs() > self[1].abs() {
            Self::new(-self[2], T::zero(), self[0]) / (self[0] * self[0] + self[2] * self[2]).sqrt()
        } else {
            Self::new(T::zero(), self[2], -self[1]) / (self[1] * self[1] + self[2] * self[2]).sqrt()
        };
        (v2, self.cross(v2))
    }
}

impl<T> Normal3<T> 
where T: Num + PartialOrd + Copy + Neg<Output = T> 
{
    pub fn face_forward(self, v: Vector3<T>) -> Self {
        if self.dot(v.into()) < T::zero() { -self } else { self }
    }
}

// SPHERICAL GEOMETRY
pub fn spherical_direction<T: NumFloat>(sin_theta: Float, cos_theta: Float, phi: Float) -> Vector3f {
    Vector3f::new(sin_theta * phi.cos(), sin_theta * phi.sin(), cos_theta)
}

pub fn spherical_triangle_area<T: NumFloat>(a: Vector3f, b: Vector3f, c: Vector3f) -> Float {
    (2.0 * (a.dot(b.cross(c))).atan2(1.0 + a.dot(b) + a.dot(c) + b.dot(c))).abs()
}

pub fn spherical_theta<T: NumFloat>(v: Vector3f) -> Float {
    pbrt::clamp_t(v.z(), -1.0, 1.0).acos()
}

pub fn spherical_phi<T: NumFloat>(v: Vector3f) -> Float {
    let p = v.y().atan2(v.x());
    if p < 0.0 {
        p + 2.0 * PI
    } else {
        p
    }
}

// AABB bouding boxes
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Bounds<T, const N: usize> {
    pub p_min: Point<T, N>,
    pub p_max: Point<T, N>,
}

impl<T, const N: usize> Bounds<T, N> 
where 
    T: Num + PartialOrd + Copy
{
    pub fn from_point(p: Point<T, N>) -> Self {
        Self { p_min: p, p_max: p }
    }

    pub fn from_points(p1: Point<T, N>, p2: Point<T, N>) -> Self {
        let mut p_min_arr = [T::zero(); N];
        let mut p_max_arr = [T::zero(); N];

        for i in 0..N {
            if p1[i] < p2[i] {
                p_min_arr[i] = p1[i];
                p_max_arr[i] = p2[i];
            } else {
                p_min_arr[i] = p2[i];
                p_max_arr[i] = p1[i];
            }
        }

        Self { p_min: Point(p_min_arr), p_max: Point(p_max_arr) }
    }

    pub fn union_point(self, p: Point<T, N>) -> Self {
        let mut p_min = self.p_min;
        let mut p_max = self.p_max;
        for i in 0..N {
            p_min[i] = min(p_min[i], p[i]);
            p_max[i] = max(p_max[i], p[i]);
        }
        Self { p_min, p_max }
    }

    pub fn union(self, b2: Self) -> Self {
        let mut p_min = self.p_min;
        let mut p_max = self.p_max;
        for i in 0..N {
            p_min[i] = min(p_min[i], b2.p_min[i]);
            p_max[i] = max(p_max[i], b2.p_max[i]);
        }
        Self { p_min, p_max }
    }
    
    pub fn diagonal(&self) -> Vector<T, N> {
        self.p_max - self.p_min
    }
    
    pub fn volume(&self) -> T {
        let d = self.diagonal();
        d.0.iter().fold(T::one(), |acc, &val| acc * val)
    }

    pub fn lerp(&self, t: Point<T, N>) -> Point<T, N>
    {
        let mut results_arr = [T::zero(); N];
        for i in 0..N {
            results_arr[i] = pbrt::lerp(t[i], self.p_min[i], self.p_max[i])
        }

        Point(results_arr)
    }

    pub fn max_dimension(&self) -> usize
    where
        Point<T, N>: Sub<Output = Vector<T, N>>,
    {
        let d = self.diagonal();
        let mut max_dim = 0;
        let mut max_span = d[0];

        for i in 1..N {
            if d[i] > max_span {
                max_span = d[i];
                max_dim = i;
            }
        }
        max_dim
    }

    pub fn offset(&self, p: &Point<T, N>) -> Vector<T, N>
    where
        Point<T, N>: Sub<Output = Vector<T, N>>,
        Vector<T, N>: DivAssign<T>
    {
        let mut o = *p - self.p_min;
        let d = self.diagonal();
        for i in 0..N { 
            if d[i] > T::zero() {
                o[i] = o[i] / d[i];
            }
        }
        o
    }

    pub fn contains(&self, p: Point<T, N>) -> bool {
        (0..N).all(|i| p[i] >= self.p_min[i] && p[i] <= self.p_max[i])
    }

    pub fn contains_exclusive(&self, p: Point<T, N>) -> bool {
        (0..N).all(|i| p[i] >= self.p_min[i] && p[i] < self.p_max[i])
    }

    pub fn is_empty(&self) -> bool {
        (0..N).any(|i| self.p_min[i] >= self.p_max[i])
    }

    pub fn is_degenerate(&self) -> bool {
        (0..N).any(|i| self.p_min[i] > self.p_max[i])
    }
}

impl<T, const N: usize> Default for Bounds<T, N>
where
    T: Bounded + Copy,
{
    fn default() -> Self {
        Self {
            p_min: Point([T::max_value(); N]),
            p_max: Point([T::min_value(); N]),
        }
    }
}

pub type Bounds2<T> = Bounds<T, 2>;
pub type Bounds3<T> = Bounds<T, 3>;
pub type Bounds3f = Bounds3<Float>;

impl<T> Bounds3<T> 
where 
    T: Num + PartialOrd + Copy + Default
{
    pub fn surface_area(&self) -> T {
        let d = self.diagonal();
        let two = T::one() + T::one();
        two * (d.x() * d.y() + d.x() * d.z() + d.y() * d.z())
    }

    pub fn bounding_sphere(&self) -> (Point3<T>, T) 
        where 
            <<Point3<T> as Sub>::Output as Normed>::Scalar: NumFloat,
    {
        let two = T::one() + T::one();
        let center = (self.p_min + self.diagonal()) / two;
        let radius = if self.contains(center) { center.distance(self.p_max) } else { T::zero() };
        (center, radius)
    }
}

#[derive(Copy, Clone, Default, PartialEq)]
pub struct Frame {
    pub x: Vector3f,
    pub y: Vector3f,
    pub z: Vector3f,
}

impl Frame {
    pub fn from_z(z: Vector3f) -> Self {
        let (x, y) = z.normalize().coordinate_system();
        Self { x, y, z: z.normalize() }
    }

    pub fn to_local(&self, v: Vector3f) -> Vector3f {
        Vector3f::new(v.dot(self.x), v.dot(self.y), v.dot(self.z))
    }

    pub fn from_local(&self, v: Vector3f) -> Vector3f {
        self.x * v.x() + self.y * v.y() + self.z * v.z()
    }
}

#[derive(Clone, Default)]
pub struct Ray {
    pub o: Point3f,
    pub d: Vector3f,
    pub medium: Option<Arc<Medium>>,
    pub time: Float,
    // We do this instead of creating a trait for Rayable or some gnarly thing like that
    pub differential: Option<RayDifferential>,
}

impl Ray {
    pub fn operator(&self, t: Float) -> Point3f {
        self.o + self.d * t
    }

    pub fn scale_differentials(&mut self, s: Float) {
        if let Some(differential) = &mut self.differential {
            differential.rx_origin = self.o + (differential.rx_origin - self.o) * s;
            differential.ry_origin = self.o + (differential.ry_origin - self.o) * s;
            differential.rx_direction = self.d + (differential.rx_direction - self.d) * s;
            differential.ry_direction = self.d + (differential.ry_direction - self.d) * s;
        }
    }
}

#[derive(Debug, Default, Copy, Clone)]
pub struct RayDifferential {
    pub rx_origin: Point3f,
    pub ry_origin: Point3f,
    pub rx_direction: Vector3f,
    pub ry_direction: Vector3f,
}