pbrt/src/utils/quaternion.rs

use std::f32::consts::PI;
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
use std::ops::{Index, IndexMut};

use crate::core::pbrt::Float;
use crate::geometry::{Vector3f, VectorLike};
use crate::utils::math::{safe_asin, sinx_over_x};

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct Quaternion {
    pub v: Vector3f,
    pub w: Float,
}

impl Default for Quaternion {
    fn default() -> Self {
        Self {
            v: Vector3f::default(),
            w: 1.0,
        }
    }
}

impl Add for Quaternion {
    type Output = Self;
    fn add(self, rhs: Quaternion) -> Self {
        Self {
            v: self.v + rhs.v,
            w: self.w + rhs.w,
        }
    }
}

impl AddAssign for Quaternion {
    fn add_assign(&mut self, rhs: Self) {
        self.v += rhs.v;
        self.w += rhs.w;
    }
}

impl Sub for Quaternion {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            v: self.v - rhs.v,
            w: self.w - rhs.w,
        }
    }
}

impl SubAssign for Quaternion {
    fn sub_assign(&mut self, rhs: Self) {
        self.v -= rhs.v;
        self.w -= rhs.w;
    }
}

impl Mul<Float> for Quaternion {
    type Output = Self;
    fn mul(self, rhs: Float) -> Self {
        Self {
            v: self.v * rhs,
            w: self.w * rhs,
        }
    }
}

impl MulAssign<Float> for Quaternion {
    fn mul_assign(&mut self, rhs: Float) {
        self.v *= rhs;
        self.w *= rhs;
    }
}

impl Div<Float> for Quaternion {
    type Output = Self;
    fn div(self, rhs: Float) -> Self {
        Self {
            v: self.v / rhs,
            w: self.w / rhs,
        }
    }
}

impl DivAssign<Float> for Quaternion {
    fn div_assign(&mut self, rhs: Float) {
        self.v /= rhs;
        self.w /= rhs;
    }
}

impl Neg for Quaternion {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            v: -self.v,
            w: -self.w,
        }
    }
}

impl Index<usize> for Quaternion {
    type Output = Float;
    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        match index {
            0 => &self.v[0],
            1 => &self.v[1],
            2 => &self.v[2],
            3 => &self.w,
            _ => panic!("Quaternion index out of bounds: {}", index),
        }
    }
}

impl IndexMut<usize> for Quaternion {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        match index {
            0 => &mut self.v[0],
            1 => &mut self.v[1],
            2 => &mut self.v[2],
            3 => &mut self.w,
            _ => panic!("Quaternion index out of bounds: {}", index),
        }
    }
}

impl Quaternion {
    pub fn dot(&self, rhs: Quaternion) -> Float {
        self.v.dot(rhs.v) + self.w * rhs.w
    }

    #[inline]
    pub fn angle_between(&self, rhs: Quaternion) -> Float {
        if self.dot(rhs) < 0.0 {
            PI - 2. * safe_asin((self.v + rhs.v).norm() / 2.)
        } else {
            2. * safe_asin((rhs.v - self.v).norm() / 2.)
        }
    }

    pub fn slerp(t: Float, q1: Quaternion, q2: Quaternion) -> Quaternion {
        let theta = q1.angle_between(q2);
        let sin_theta_over_theta = sinx_over_x(theta);
        q1 * (1. - t) * sinx_over_x((1. - t) * theta) / sin_theta_over_theta
            + q2 * t * sinx_over_x(t * theta) / sin_theta_over_theta
    }

    pub fn length(&self) -> Float {
        self.v.norm()
    }

    pub fn normalize(&self) -> Self {
        Quaternion {
            v: self.v.normalize(),
            w: self.w,
        }
    }
}