pbrt/shared/src/core/light.rs

use crate::core::color::RGB;
use crate::core::geometry::{
    Bounds2f, Bounds3f, DirectionCone, Normal3f, Point2f, Point2i, Point3f, Point3fi, Ray,
    Vector3f, VectorLike, cos_theta,
};
use crate::core::image::DeviceImage;
use crate::core::interaction::{
    Interaction, InteractionBase, InteractionTrait, MediumInteraction, SimpleInteraction,
    SurfaceInteraction,
};
use crate::core::medium::MediumInterface;
use crate::core::spectrum::{Spectrum, SpectrumTrait};
use crate::lights::*;
use crate::spectra::{
    DenselySampledSpectrum, LAMBDA_MAX, LAMBDA_MIN, RGBColorSpace, RGBIlluminantSpectrum,
    SampledSpectrum, SampledWavelengths,
};
use crate::utils::Transform;
use crate::utils::math::{equal_area_sphere_to_square, radians, safe_sqrt, smooth_step, square};
use crate::utils::ptr::{DevicePtr, Ptr};
use crate::utils::sampling::DevicePiecewiseConstant2D;
use crate::{Float, PI};
use bitflags::bitflags;

use enum_dispatch::enum_dispatch;

bitflags! {
    #[repr(C)]
    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
    pub struct LightType: u32 {
        const DeltaPosition  = 1;
        const DeltaDirection = 2;
        const Area           = 4;
        const Infinite       = 8;
    }
}

impl LightType {
    pub fn is_infinite(&self) -> bool {
        self.contains(LightType::Infinite)
    }

    pub fn is_delta_light(&self) -> bool {
        self.contains(LightType::DeltaPosition) || self.contains(LightType::DeltaDirection)
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct LightLeSample {
    pub l: SampledSpectrum,
    pub ray: Ray,
    pub pdf_pos: Float,
    pub pdf_dir: Float,
    pub intr: Interaction,
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct LightLiSample {
    pub l: SampledSpectrum,
    pub wi: Vector3f,
    pub pdf: Float,
    pub p_light: Interaction,
}

#[cfg(not(target_os = "cuda"))]
impl LightLiSample {
    pub fn new(l: SampledSpectrum, wi: Vector3f, pdf: Float, p_light: Interaction) -> Self {
        Self {
            l,
            wi,
            pdf,
            p_light,
        }
    }
}
#[repr(C)]
#[derive(Debug, Copy, Default, Clone)]
pub struct LightSampleContext {
    pub pi: Point3fi,
    pub n: Normal3f,
    pub ns: Normal3f,
}

impl LightSampleContext {
    pub fn new(pi: Point3fi, n: Normal3f, ns: Normal3f) -> Self {
        Self { pi, n, ns }
    }

    pub fn p(&self) -> Point3f {
        self.pi.into()
    }
}

impl From<&SurfaceInteraction> for LightSampleContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            pi: si.common.pi,
            n: si.common.n,
            ns: si.shading.n,
        }
    }
}

impl From<&MediumInteraction> for LightSampleContext {
    fn from(mi: &MediumInteraction) -> Self {
        Self {
            pi: mi.common.pi,
            n: Normal3f::default(),
            ns: Normal3f::default(),
        }
    }
}

impl From<&Interaction> for LightSampleContext {
    fn from(intr: &Interaction) -> Self {
        match intr {
            Interaction::Surface(si) => Self {
                pi: si.common.pi,
                n: si.common.n,
                ns: si.shading.n,
            },

            Interaction::Medium(mi) => Self {
                pi: mi.common.pi,
                n: mi.common.n,
                ns: mi.common.n,
            },

            Interaction::Simple(sim) => Self {
                pi: sim.common.pi,
                n: sim.common.n,
                ns: sim.common.n,
            },
        }
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct LightBase {
    pub render_from_light: Transform,
    pub light_type: LightType,
    pub medium_interface: MediumInterface,
}

impl LightBase {
    pub fn new(
        light_type: LightType,
        render_from_light: Transform,
        medium_interface: MediumInterface,
    ) -> Self {
        Self {
            light_type,
            render_from_light,
            medium_interface,
        }
    }

    fn l(
        &self,
        _p: Point3f,
        _n: Normal3f,
        _uv: Point2f,
        _w: Vector3f,
        _lambda: &SampledWavelengths,
    ) -> SampledSpectrum {
        SampledSpectrum::default()
    }

    pub fn light_type(&self) -> LightType {
        self.light_type
    }

    pub fn le(&self, _ray: &Ray, _lambda: &SampledWavelengths) -> SampledSpectrum {
        SampledSpectrum::default()
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct LightBounds {
    pub bounds: Bounds3f,
    pub phi: Float,
    pub w: Vector3f,
    pub cos_theta_o: Float,
    pub cos_theta_e: Float,
    pub two_sided: bool,
}

#[cfg(not(target_os = "cuda"))]
impl LightBounds {
    pub fn new(
        bounds: &Bounds3f,
        w: Vector3f,
        phi: Float,
        cos_theta_o: Float,
        cos_theta_e: Float,
        two_sided: bool,
    ) -> Self {
        Self {
            bounds: *bounds,
            phi,
            w,
            cos_theta_o,
            cos_theta_e,
            two_sided,
        }
    }
}

impl LightBounds {
    pub fn centroid(&self) -> Point3f {
        self.bounds.p_min + Vector3f::from(self.bounds.p_max) / 2.
    }

    pub fn importance(&self, p: Point3f, n: Normal3f) -> Float {
        // Compute clamped squared distance to reference point
        let pc = self.centroid();
        let d2_raw = p.distance_squared(pc);
        let d2 = d2_raw.max(self.bounds.diagonal().norm()) / 2.;
        let cos_sub_clamped = |sin_theta_a: Float,
                               cos_theta_a: Float,
                               sin_theta_b: Float,
                               cos_theta_b: Float|
         -> Float {
            if cos_theta_a > cos_theta_b {
                return 1.;
            }
            cos_theta_a * cos_theta_b + sin_theta_a * sin_theta_b
        };

        let sin_sub_clamped = |sin_theta_a: Float,
                               cos_theta_a: Float,
                               sin_theta_b: Float,
                               cos_theta_b: Float|
         -> Float {
            if cos_theta_a > cos_theta_b {
                return 1.;
            }
            sin_theta_a * cos_theta_b - cos_theta_a * sin_theta_b
        };

        let wi = (p - pc).normalize();
        let mut cos_theta_w = self.w.dot(wi);
        if self.two_sided {
            cos_theta_w = cos_theta_w.abs();
        }
        let sin_theta_w = safe_sqrt(1. - square(cos_theta_w));
        let cos_theta_b = DirectionCone::bound_subtended_directions(&self.bounds, p).cos_theta;
        let sin_theta_b = safe_sqrt(1. - square(cos_theta_b));
        let sin_theta_o = safe_sqrt(1. - square(self.cos_theta_o));
        let cos_theta_x = cos_sub_clamped(sin_theta_w, cos_theta_w, sin_theta_o, self.cos_theta_o);
        let sin_theta_x = sin_sub_clamped(sin_theta_w, cos_theta_w, sin_theta_o, self.cos_theta_o);
        let cos_theta_p = cos_sub_clamped(sin_theta_x, cos_theta_x, sin_theta_b, cos_theta_b);

        if cos_theta_p <= self.cos_theta_e {
            return 0.;
        }

        let mut importance = self.phi * cos_theta_p / d2;
        if n != Normal3f::new(0., 0., 0.) {
            let cos_theta_i = wi.abs_dot(n.into());
            let sin_theta_i = safe_sqrt(1. - square(cos_theta_i));
            let cos_thetap_i = cos_sub_clamped(sin_theta_i, cos_theta_i, sin_theta_b, cos_theta_b);
            importance *= cos_thetap_i;
        }
        importance
    }

    pub fn union(a: &Self, b: &Self) -> Self {
        if a.phi == 0. {
            return a.clone();
        }
        if b.phi == 0. {
            return b.clone();
        }

        let a_cone = DirectionCone::new(a.w, a.cos_theta_o);
        let b_cone = DirectionCone::new(b.w, b.cos_theta_o);
        let cone = DirectionCone::union(&a_cone, &b_cone);
        let cos_theta_o = cone.cos_theta;
        let cos_theta_e = a.cos_theta_e.min(b.cos_theta_e);
        LightBounds::new(
            &a.bounds.union(b.bounds),
            cone.w,
            a.phi + b.phi,
            cos_theta_o,
            cos_theta_e,
            a.two_sided || b.two_sided,
        )
    }
}

#[enum_dispatch]
pub trait LightTrait {
    fn base(&self) -> &LightBase;

    fn sample_li(
        &self,
        ctx: &LightSampleContext,
        u: Point2f,
        lambda: &SampledWavelengths,
        allow_incomplete_pdf: bool,
    ) -> Option<LightLiSample>;

    fn pdf_li(&self, ctx: &LightSampleContext, wi: Vector3f, allow_incomplete_pdf: bool) -> Float;

    fn l(
        &self,
        p: Point3f,
        n: Normal3f,
        uv: Point2f,
        w: Vector3f,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum;

    fn le(&self, ray: &Ray, lambda: &SampledWavelengths) -> SampledSpectrum;

    fn light_type(&self) -> LightType {
        self.base().light_type
    }

    #[cfg(not(target_os = "cuda"))]
    fn bounds(&self) -> Option<LightBounds>;

    #[cfg(not(target_os = "cuda"))]
    fn preprocess(&mut self, scene_bounds: &Bounds3f);

    #[cfg(not(target_os = "cuda"))]
    fn phi(&self, lambda: SampledWavelengths) -> SampledSpectrum;
}

#[enum_dispatch(LightTrait)]
#[repr(C)]
#[derive(Debug, Clone, Copy)]
#[allow(clippy::large_enum_variant)]
pub enum Light {
    DiffuseArea(DiffuseAreaLight),
    Distant(DistantLight),
    Goniometric(GoniometricLight),
    InfiniteUniform(InfiniteUniformLight),
    InfiniteImage(InfiniteImageLight),
    InfinitePortal(InfinitePortalLight),
    Point(PointLight),
    Projection(ProjectionLight),
    Spot(SpotLight),
}