pbrt/shared/src/core/scattering.rs

use crate::core::geometry::{
    Normal3f, Point2f, Vector2f, Vector3f, VectorLike, abs_cos_theta, cos_phi, cos2_theta, sin_phi,
    tan2_theta,
};
use crate::core::pbrt::{Float, PI};
use crate::spectra::{N_SPECTRUM_SAMPLES, SampledSpectrum};
use crate::utils::math::{clamp, lerp, safe_sqrt, square};
use crate::utils::sampling::sample_uniform_disk_polar;

use num::complex::Complex;

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct TrowbridgeReitzDistribution {
    alpha_x: Float,
    alpha_y: Float,
}

impl TrowbridgeReitzDistribution {
    pub fn new(alpha_x: Float, alpha_y: Float) -> Self {
        Self { alpha_x, alpha_y }
    }

    pub fn d(&self, wm: Vector3f) -> Float {
        let tan2_theta = tan2_theta(wm);
        if tan2_theta.is_infinite() {
            return 0.;
        }
        let cos4_theta = square(cos2_theta(wm));
        let e =
            tan2_theta * (square(cos_phi(wm) / self.alpha_x) + square(sin_phi(wm) / self.alpha_y));
        1.0 / (PI * self.alpha_x * self.alpha_y * cos4_theta * square(1. + e))
    }

    pub fn effectively_smooth(&self) -> bool {
        self.alpha_x.max(self.alpha_y) < 1e-3
    }

    pub fn lambda(&self, w: Vector3f) -> Float {
        let tan2_theta = tan2_theta(w);
        if tan2_theta.is_infinite() {
            return 0.;
        }
        let alpha2 = square(cos_phi(w) * self.alpha_x) + square(sin_phi(w) * self.alpha_y);
        ((1. + alpha2 * tan2_theta).sqrt() - 1.) / 2.
    }

    pub fn g(&self, wo: Vector3f, wi: Vector3f) -> Float {
        1. / (1. + self.lambda(wo) + self.lambda(wi))
    }

    pub fn g1(&self, w: Vector3f) -> Float {
        1. / (1. / self.lambda(w))
    }

    pub fn d_from_w(&self, w: Vector3f, wm: Vector3f) -> Float {
        self.g1(w) / abs_cos_theta(w) * self.d(wm) * w.dot(wm).abs()
    }

    pub fn pdf(&self, w: Vector3f, wm: Vector3f) -> Float {
        self.d_from_w(w, wm)
    }

    pub fn sample_wm(&self, w: Vector3f, u: Point2f) -> Vector3f {
        let mut wh = Vector3f::new(self.alpha_x * w.x(), self.alpha_y * w.y(), w.z()).normalize();
        if wh.z() < 0. {
            wh = -wh;
        }
        let t1 = if wh.z() < 0.99999 {
            Vector3f::new(0., 0., 1.).cross(wh).normalize()
        } else {
            Vector3f::new(1., 0., 0.)
        };
        let t2 = wh.cross(t1);
        let mut p = sample_uniform_disk_polar(u);
        let h = (1. - square(p.x())).sqrt();
        p[1] = lerp((1. + wh.z()) / 2., h, p.y());
        let pz = 0_f32.max(1. - Vector2f::from(p).norm_squared());
        let nh = p.x() * t1 + p.y() * t2 + pz * wh;
        Vector3f::new(
            self.alpha_x * nh.x(),
            self.alpha_y * nh.y(),
            nh.z().max(1e-6),
        )
        .normalize()
    }

    pub fn roughness_to_alpha(roughness: Float) -> Float {
        roughness.sqrt()
    }

    pub fn regularize(&mut self) {
        if self.alpha_x < 0.3 {
            self.alpha_x = clamp(2. * self.alpha_x, 0.1, 0.3);
        }
        if self.alpha_y < 0.3 {
            self.alpha_y = clamp(2. * self.alpha_y, 0.1, 0.3);
        }
    }
}

pub fn refract(wi: Vector3f, n: Normal3f, eta_ratio: Float) -> Option<(Vector3f, Float)> {
    let mut n_interface = n;
    let mut eta = eta_ratio;

    let mut cos_theta_i = Vector3f::from(n_interface).dot(wi);

    if cos_theta_i < 0.0 {
        eta = 1.0 / eta;
        cos_theta_i = -cos_theta_i;
        n_interface = -n_interface;
    }

    let sin2_theta_i = (1.0 - square(cos_theta_i)).max(0.0_f32);
    let sin2_theta_t = sin2_theta_i / square(eta);

    // Handle total internal reflection
    if sin2_theta_t >= 1.0 {
        return None;
    }

    let cos_theta_t = (1.0 - sin2_theta_t).sqrt();

    let wt = -wi / eta + (cos_theta_i / eta - cos_theta_t) * Vector3f::from(n_interface);
    Some((wt, eta))
}

pub fn reflect(wo: Vector3f, n: Normal3f) -> Vector3f {
    -wo + Vector3f::from(2. * wo.dot(n.into()) * n)
}

pub fn fr_dielectric(cos_theta_i: Float, eta: Float) -> Float {
    let mut cos_safe = clamp(cos_theta_i, -1., 1.);
    let mut eta_corr = eta;
    if cos_theta_i < 0. {
        eta_corr = 1. / eta_corr;
        cos_safe = -cos_safe;
    }
    let sin2_theta_i = 1. - square(cos_safe);
    let sin2_theta_t = sin2_theta_i / square(eta_corr);
    if sin2_theta_t >= 1. {
        return 1.;
    }
    let cos_theta_t = safe_sqrt(1. - sin2_theta_t);
    let r_parl = (eta_corr * cos_safe - cos_theta_t) / (eta_corr * cos_safe + cos_theta_t);
    let r_perp = (cos_safe - eta_corr * cos_theta_t) / (cos_safe + eta_corr * cos_theta_t);

    (square(r_parl) + square(r_perp)) / 2.
}

pub fn fr_complex(cos_theta_i: Float, eta: Complex<Float>) -> Float {
    let cos_corr = clamp(cos_theta_i, 0., 1.);
    let sin2_theta_i = 1. - square(cos_corr);
    let sin2_theta_t: Complex<Float> = sin2_theta_i / square(eta);
    let cos2_theta_t: Complex<Float> = (1. - sin2_theta_t).sqrt();

    let r_parl = (eta * cos_corr - cos2_theta_t) / (eta * cos_corr + cos2_theta_t);
    let r_perp = (cos_corr - eta * cos2_theta_t) / (cos_corr + eta * cos2_theta_t);

    (r_parl.norm() + r_perp.norm()) / 2.
}

pub fn fr_complex_from_spectrum(
    cos_theta_i: Float,
    eta: SampledSpectrum,
    k: SampledSpectrum,
) -> SampledSpectrum {
    let mut result = SampledSpectrum::default();
    for i in 0..N_SPECTRUM_SAMPLES {
        result[i] = fr_complex(cos_theta_i, Complex::new(eta[i], k[i]));
    }
    result
}