pbrt/src/geometry/primitives.rs

use super::traits::Tuple;
use super::{Float, NumFloat, PI};
use crate::utils::interval::Interval;
use crate::utils::math::safe_asin;
use num_traits::{Num, Signed, Zero, FloatConst};
use std::iter::Sum;
use std::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};

// 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_export]
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> $Struct<T, N>
        where
            T: Zero + Copy,
        {
            #[inline]
            pub fn zero() -> Self {
                Self([T::zero(); N])
            }
        }

        impl<const N: usize> $Struct<f32, N> {
            #[inline]
            pub fn floor(&self) -> $Struct<i32, N> {
                $Struct(self.0.map(|v| v.floor() as i32))
            }
        }

        impl<T, const N: usize> $Struct<T, N>
        where
            T: Copy,
        {
            #[inline]
            pub fn fill(value: T) -> Self {
                Self([value; 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 {
                Self(self.0.map(|c| -c))
            }
        }
    };
}

#[macro_export]
macro_rules! impl_scalar_ops {
    ($Struct:ident) => {
        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<const N: usize> Mul<$Struct<Float, N>> for Float {
            type Output = $Struct<Float, N>;
            fn mul(self, rhs: $Struct<Float, N>) -> Self::Output {
                rhs * self
            }
        }

        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;
                }
            }
        }
    };
}

#[macro_export]
macro_rules! impl_op {
    ($Op:ident, $op:ident, $Lhs:ident, $Rhs:ident, $Output:ident) => {
        impl<T, const N: usize> $Op<$Rhs<T, N>> for $Lhs<T, N>
        where
            T: $Op<Output = T> + Copy,
        {
            type Output = $Output<T, N>;
            fn $op(self, rhs: $Rhs<T, N>) -> Self::Output {
                let mut result = self.0;
                for i in 0..N {
                    result[i] = $Op::$op(self.0[i], rhs.0[i]);
                }
                $Output(result)
            }
        }
    };
}

#[macro_export]
macro_rules! impl_op_assign {
    ($OpAssign:ident, $op_assign:ident, $Lhs:ident, $Rhs:ident) => {
        impl<T, const N: usize> $OpAssign<$Rhs<T, N>> for $Lhs<T, N>
        where
            T: $OpAssign + Copy,
        {
            fn $op_assign(&mut self, rhs: $Rhs<T, N>) {
                for i in 0..N {
                    $OpAssign::$op_assign(&mut self.0[i], rhs.0[i]);
                }
            }
        }
    };
}

#[macro_export]
macro_rules! impl_float_vector_ops {
    ($Struct:ident) => {
        // This impl block is constrained to only apply when the scalar type `T` is a float.
        impl<T, const N: usize> $Struct<T, N>
        where
            T: NumFloat,
        {
            pub fn dot(self, rhs: Self) -> T {
                let mut sum = T::zero();
                for i in 0..N {
                    sum = sum + self[i] * rhs[i];
                }
                sum
            }

            pub fn norm_squared(&self) -> T {
                let mut sum = T::zero();
                for i in 0..N {
                    sum = sum + self[i] * self[i];
                }
                sum
            }

            pub fn norm(&self) -> T {
                self.norm_squared().sqrt()
            }

            pub fn normalize(self) -> Self {
                let n = self.norm();
                if n.is_zero() {
                    self
                } else {
                    self / n
                }
            }

            pub fn angle_between(self, rhs: Self) -> T {
                let dot_product = self.normalize().dot(rhs.normalize());
                let clamped_dot = dot_product.min(T::one()).max(-T::one());
                clamped_dot.acos()
            }

            pub fn project_on(self, rhs: Self) -> Self {
                let rhs_norm_sq = rhs.norm_squared();
                if rhs_norm_sq.is_zero() {
                    // This now calls the inherent `zero()` method from `impl_tuple_core!`
                    Self::zero()
                } else {
                    // Note: This requires Mul<T> to be implemented for the struct, 
                    // which your `impl_scalar_ops!` macro already does.
                    rhs * (self.dot(rhs) / rhs_norm_sq)
                }
            }
        }
    };
}

macro_rules! impl_abs {
    ($Struct:ident) => {
        impl<T, const N: usize> $Struct<T, N>
        where
            T: Signed + Copy,
        {
            pub fn abs(self) -> Self {
                let mut result = self.0;
                for i in 0..N {
                    result[i] = result[i].abs();
                }
                Self(result)
            }
        }
    };
}

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_tuple_core!(Vector);
impl_tuple_core!(Point);
impl_tuple_core!(Normal);

impl_scalar_ops!(Vector);
impl_scalar_ops!(Normal);

// Addition
impl_op!(Add, add, Vector, Vector, Vector);
impl_op!(Add, add, Point, Vector, Point);
impl_op!(Add, add, Vector, Point, Point);
impl_op!(Add, add, Normal, Normal, Normal);

// Subtraction
impl_op!(Sub, sub, Vector, Vector, Vector);
impl_op!(Sub, sub, Point, Vector, Point);
impl_op!(Sub, sub, Point, Point, Vector);
impl_op!(Sub, sub, Normal, Normal, Normal);

// AddAssign
impl_op_assign!(AddAssign, add_assign, Vector, Vector);
impl_op_assign!(AddAssign, add_assign, Point, Vector);
impl_op_assign!(AddAssign, add_assign, Normal, Normal);

// SubAssign
impl_op_assign!(SubAssign, sub_assign, Vector, Vector);
impl_op_assign!(SubAssign, sub_assign, Point, Vector);
impl_op_assign!(SubAssign, sub_assign, Normal, Normal);

impl_float_vector_ops!(Vector);
impl_float_vector_ops!(Normal);
impl_abs!(Vector);
impl_abs!(Normal);

impl_accessors!(Vector);
impl_accessors!(Point);
impl_accessors!(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: Copy, const N: usize> From<Vector<T, N>> for Point<T, N> {
    fn from(v: Vector<T, N>) -> Self {
        Self(v.0)
    }
}

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

impl<T, const N: usize> Point<T, N>
where
    T: NumFloat,
{
    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 Point2f = Point2<Float>;
pub type Point2i = Point2<i32>;
pub type Point2fi = Point2<Interval>;
pub type Point3<T> = Point<T, 3>;
pub type Point3f = Point3<Float>;
pub type Point3i = Point3<i32>;
pub type Point3fi = Point3<Interval>;
pub type Vector2<T> = Vector<T, 2>;
pub type Vector2f = Vector2<Float>;
pub type Vector2i = Vector2<i32>;
pub type Vector2fi = Vector2<Interval>;
pub type Vector3<T> = Vector<T, 3>;
pub type Vector3f = Vector3<Float>;
pub type Vector3i = Vector3<i32>;
pub type Vector3fi = Vector3<Interval>;
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<const N: usize> Point<Interval, N> {
    pub fn new_from_point(p: Point<Float, N>) -> Self {
        let mut arr = [Interval::default(); N];
        for i in 0..N {
            arr[i] = Interval::new(p[i]);
        }
        Self(arr)
    }

    pub fn new_with_error(p: Point<Float, N>, e: Vector<Float, N>) -> Self {
        let mut arr = [Interval::default(); N];
        for i in 0..N {
            arr[i] = Interval::new_from_value_and_error(p[i], e[i]);
        }
        Self(arr)
    }

    pub fn error(&self) -> Vector<Float, N> {
        let mut arr = [0.0; N];
        for i in 0..N {
            arr[i] = self[i].width() / 2.0;
        }
        Vector(arr)
    }

    pub fn midpoint(&self) -> Point<Float, N> {
        let mut arr = [0.0; N];
        for i in 0..N {
            arr[i] = self[i].midpoint();
        }
        Point(arr)
    }

    pub fn is_exact(&self) -> bool {
        self.0.iter().all(|interval| interval.width() == 0.0)
    }
}

impl<const N: usize> Vector<Interval, N> {
    pub fn new_from_vector(v: Vector<Float, N>) -> Self {
        let mut arr = [Interval::default(); N];
        for i in 0..N {
            arr[i] = Interval::new(v[i]);
        }
        Self(arr)
    }

    pub fn new_with_error(v: Vector<Float, N>, e: Vector<Float, N>) -> Self {
        let mut arr = [Interval::default(); N];
        for i in 0..N {
            arr[i] = Interval::new_from_value_and_error(v[i], e[i]);
        }
        Self(arr)
    }

    pub fn error(&self) -> Vector<Float, N> {
        let mut arr = [0.0; N];
        for i in 0..N {
            arr[i] = self[i].width() / 2.0;
        }
        Vector(arr)
    }

    pub fn midpoint(&self) -> Vector<Float, N> {
        let mut arr = [0.0; N];
        for i in 0..N {
            arr[i] = self[i].midpoint();
        }
        Vector(arr)
    }

    pub fn is_exact(&self) -> bool {
        self.0.iter().all(|interval| interval.width() == 0.0)
    }
}

impl<const N: usize> From<Point<Interval, N>> for Point<Float, N> {
    fn from(pi: Point<Interval, N>) -> Self {
        let mut arr = [0.0; N];
        for i in 0..N {
            arr[i] = pi[i].midpoint();
        }
        Point(arr)
    }
}

impl<const N: usize> Mul<Vector<Interval, N>> for Interval {
    type Output = Vector<Interval, N>;
    fn mul(self, rhs: Vector<Interval, N>) -> Self::Output {
        rhs * self
    }
}

impl<const N: usize> Div<Vector<Interval, N>> for Interval {
    type Output = Vector<Interval, N>;
    fn div(self, rhs: Vector<Interval, N>) -> Self::Output {
        let mut result = rhs.0;
        for i in 0..N {
            result[i] = self / rhs[i];
        }
        Vector(result)
    }
}

impl<const N: usize> From<Vector<i32, N>> for Vector<f32, N> {
    fn from(v: Vector<i32, N>) -> Self {
        Self(v.0.map(|c| c as f32))
    }
}

impl<const N: usize> From<Point<i32, N>> for Point<Float, N> {
    fn from(p: Point<i32, N>) -> Self {
        Point(p.0.map(|c| c as Float))
    }
}

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

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

impl Frame {
    pub fn new(x: Vector3f, z: Vector3f) -> Self {
        Self {
            x,
            y: z.cross(x),
            z,
        }
    }

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

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

    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 to_local_normal(&self, n: Normal3f) -> Normal3f {
        let n: Vector3f = n.into();
        Normal3f::new(n.dot(self.x), n.dot(self.y), n.dot(self.z))
    }

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

    pub fn from_local_normal(&self, v: Normal3f) -> Normal3f {
        Normal3f::from(self.x * v.x() + self.y * v.y() + self.z * v.z())
    }
}