pbrt/shared/src/lights/sampler.rs

use crate::core::geometry::primitives::OctahedralVector;
use crate::core::geometry::{Bounds3f, Normal3f, Point3f, Vector3f, VectorLike};
use crate::core::geometry::{DirectionCone, Normal};
use crate::core::light::Light;
use crate::core::light::{LightBounds, LightSampleContext};
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::math::{clamp, lerp, sample_discrete};
use crate::utils::math::{safe_sqrt, square};
use crate::utils::ptr::Ptr;
use crate::utils::sampling::AliasTable;
use crate::{Float, ONE_MINUS_EPSILON, PI};
use enum_dispatch::enum_dispatch;
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct CompactLightBounds {
    pub w: OctahedralVector,
    pub phi: Float,

    // [0..15]  = qCosTheta_o
    // [15..30] = qCosTheta_e
    // [30..31] = twoSided
    // [31..32] = Unused/Padding
    packed_info: u32,

    pub qb: [[u16; 3]; 2],
}

#[allow(clippy::derivable_impls)]
impl Default for CompactLightBounds {
    fn default() -> Self {
        Self {
            w: OctahedralVector::default(),
            phi: Float::default(),
            packed_info: u32::default(),
            qb: [[u16::default(); 3]; 2],
        }
    }
}

const _: () = assert!(core::mem::size_of::<CompactLightBounds>() == 24);

impl CompactLightBounds {
    pub fn new(lb: &LightBounds, all_b: &Bounds3f) -> Self {
        let q_cos_o = Self::quantize_cos(lb.cos_theta_o);
        let q_cos_e = Self::quantize_cos(lb.cos_theta_e);
        let two_sided = if lb.two_sided { 1 } else { 0 };

        // | - twoSided (1) - | - qCosTheta_e (15) - | - qCosTheta_o (15) - |
        let packed_info = (q_cos_o & 0x7FFF) | ((q_cos_e & 0x7FFF) << 15) | (two_sided << 30);

        let mut qb = [[0u16; 3]; 2];
        for i in 0..3 {
            qb[0][i] = Self::quantize_bounds(lb.bounds.p_min[i], all_b.p_min[i], all_b.p_max[i])
                .floor() as u16;
            qb[1][i] = Self::quantize_bounds(lb.bounds.p_max[i], all_b.p_min[i], all_b.p_max[i])
                .ceil() as u16;
        }

        Self {
            w: OctahedralVector::new(lb.w.normalize()),
            phi: lb.phi,
            packed_info,
            qb,
        }
    }

    #[inline(always)]
    pub fn two_sided(&self) -> bool {
        (self.packed_info >> 30) & 1 == 1
    }

    #[inline(always)]
    fn q_cos_theta_o(&self) -> u32 {
        self.packed_info & 0x7FFF
    }

    #[inline(always)]
    fn q_cos_theta_e(&self) -> u32 {
        (self.packed_info >> 15) & 0x7FFF
    }

    #[inline]
    pub fn cos_theta_o(&self) -> Float {
        2.0 * (self.q_cos_theta_o() as Float / 32767.0) - 1.0
    }

    #[inline]
    pub fn cos_theta_e(&self) -> Float {
        2.0 * (self.q_cos_theta_e() as Float / 32767.0) - 1.0
    }

    pub fn importance(&self, p: Point3f, n: Normal3f, all_b: &Bounds3f) -> Float {
        let bounds = self.bounds(all_b);
        let cos_o = self.cos_theta_o();
        let cos_e = self.cos_theta_e();

        let pc = bounds.centroid();
        let d2 = p.distance_squared(pc).max(bounds.diagonal().norm() * 0.5);

        let cos_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
            if cos_a > cos_b {
                1.0
            } else {
                cos_a * cos_b + sin_a * sin_b
            }
        };
        let sin_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
            if cos_a > cos_b {
                0.0
            } else {
                sin_a * cos_b - cos_a * sin_b
            }
        };

        let wi = (p - pc).normalize();
        let w_vec = self.w.to_vector();
        let mut cos_w = w_vec.dot(wi);
        if self.two_sided() {
            cos_w = cos_w.abs();
        }
        let sin_w = safe_sqrt(1.0 - square(cos_w));
        let cos_b = DirectionCone::bound_subtended_directions(&bounds, p).cos_theta;
        let sin_b = safe_sqrt(1. - square(cos_b));
        let sin_o = safe_sqrt(1. - square(cos_o));
        let cos_x = cos_sub_clamped(sin_w, cos_w, sin_o, cos_o);
        let sin_x = sin_sub_clamped(sin_w, cos_w, sin_o, cos_o);
        let cos_p = cos_sub_clamped(sin_x, cos_x, sin_b, cos_b);
        if cos_p <= cos_e {
            return 0.;
        }
        let mut importance = self.phi * cos_p / d2;
        if n != Normal3f::zero() {
            let cos_i = wi.abs_dot(n.into());
            let sin_i = safe_sqrt(1. - square(cos_i));
            let cos_pi = cos_sub_clamped(sin_i, cos_i, sin_b, cos_b);
            importance *= cos_pi;
        }
        importance
    }

    pub fn bounds(&self, all_b: &Bounds3f) -> Bounds3f {
        let mut p_min = Point3f::default();
        let mut p_max = Point3f::default();

        for i in 0..3 {
            let t_min = self.qb[0][i] as Float / 65535.0;
            let t_max = self.qb[1][i] as Float / 65535.0;
            p_min[i] = lerp(t_min, all_b.p_min[i], all_b.p_max[i]);
            p_max[i] = lerp(t_max, all_b.p_min[i], all_b.p_max[i]);
        }
        Bounds3f::from_points(p_min, p_max)
    }

    fn quantize_cos(c: Float) -> u32 {
        (32767.0 * ((c + 1.0) * 0.5)).floor() as u32
    }

    fn quantize_bounds(c: Float, min: Float, max: Float) -> Float {
        if min == max {
            return 0.0;
        }
        65535.0 * clamp((c - min) / (max - min), 0.0, 1.0)
    }
}

#[derive(Debug, Clone)]
pub struct SampledLight {
    pub light: Ptr<Light>,
    pub p: Float,
}

impl SampledLight {
    pub fn new(light: Light, p: Float) -> Self {
        Self {
            light: Ptr::from(&light),
            p,
        }
    }
}

#[enum_dispatch]
pub trait LightSamplerTrait {
    fn sample_with_context(&self, ctx: &LightSampleContext, u: Float) -> Option<SampledLight>;
    fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Light) -> Float;
    fn sample(&self, u: Float) -> Option<SampledLight>;
    fn pmf(&self, light: &Light) -> Float;
}

#[derive(Clone, Debug)]
#[enum_dispatch(LightSamplerTrait)]
pub enum LightSampler {
    Uniform(UniformLightSampler),
    Power(PowerLightSampler),
    BVH(BVHLightSampler),
}

#[derive(Clone, Debug)]
pub struct UniformLightSampler {
    lights: *const Light,
    lights_len: u32,
}

impl UniformLightSampler {
    pub fn new(lights: *const Light, lights_len: u32) -> Self {
        Self { lights, lights_len }
    }

    #[inline(always)]
    fn light(&self, idx: usize) -> Light {
        unsafe { *self.lights.add(idx) }
    }
}

impl LightSamplerTrait for UniformLightSampler {
    fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
        self.sample(u)
    }
    fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Light) -> Float {
        self.pmf(light)
    }
    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.lights_len == 0 {
            return None;
        }

        let light_index = (u as u32 * self.lights_len).min(self.lights_len - 1) as usize;
        Some(SampledLight {
            light: Ptr::from(&self.light(light_index)),
            p: 1. / self.lights_len as Float,
        })
    }
    fn pmf(&self, _light: &Light) -> Float {
        if self.lights_len == 0 {
            return 0.;
        }
        1. / self.lights_len as Float
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Alias {
    pub q: Float,
    pub alias: u32,
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct PowerLightSampler {
    pub lights: Ptr<Light>,
    pub lights_len: u32,
    pub alias_table: AliasTable,
}

unsafe impl Send for PowerLightSampler {}
unsafe impl Sync for PowerLightSampler {}

impl LightSamplerTrait for PowerLightSampler {
    fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
        self.sample(u)
    }

    fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Light) -> Float {
        self.pmf(light)
    }

    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.alias_table.size() == 0 {
            return None;
        }

        let (light_index, pmf, _) = self.alias_table.sample(u);

        let light_ref = unsafe { self.lights.add(light_index as usize) };
        Some(SampledLight {
            light: light_ref,
            p: pmf,
        })
    }

    fn pmf(&self, light: &Light) -> Float {
        let target = light as *const Light;
        for i in 0..self.lights_len as usize {
            if unsafe { self.lights.add(i) }.as_raw() == target {
                return self.alias_table.pmf(i as u32);
            }
        }
        0.0
    }
}

#[derive(Clone, Copy, Debug, Default)]
#[repr(C, align(32))]
pub struct LightBVHNode {
    pub light_bounds: CompactLightBounds,
    // Bit 31 (MSB) : isLeaf (1 bit)
    // Bits 0..31   : childOrLightIndex (31 bits)
    packed_data: u32,
}

const _: () = assert!(core::mem::size_of::<LightBVHNode>() == 32);

impl LightBVHNode {
    /// Mask to isolate the Leaf Flag (Bit 31)
    const LEAF_MASK: u32 = 0x8000_0000;
    /// Mask to isolate the Index (Bits 0-30)
    const INDEX_MASK: u32 = 0x7FFF_FFFF;

    pub fn make_leaf(light_index: u32, cb: CompactLightBounds) -> Self {
        debug_assert!(
            (light_index & Self::LEAF_MASK) == 0,
            "Light index too large"
        );

        Self {
            light_bounds: cb,
            // Set index and flip the MSB to 1
            packed_data: light_index | Self::LEAF_MASK,
        }
    }

    pub fn make_interior(child_index: u32, cb: CompactLightBounds) -> Self {
        debug_assert!(
            (child_index & Self::LEAF_MASK) == 0,
            "Child index too large"
        );

        Self {
            light_bounds: cb,
            // Set index, MSB remains 0
            packed_data: child_index,
        }
    }

    #[inline(always)]
    pub fn is_leaf(&self) -> bool {
        (self.packed_data & Self::LEAF_MASK) != 0
    }

    #[inline(always)]
    pub fn light_index(&self) -> u32 {
        debug_assert!(self.is_leaf());
        self.packed_data & Self::INDEX_MASK
    }

    #[inline(always)]
    pub fn child_index(&self) -> u32 {
        debug_assert!(!self.is_leaf());
        self.packed_data & Self::INDEX_MASK
    }

    #[inline(always)]
    pub fn child_or_light_index(&self) -> u32 {
        self.packed_data & Self::INDEX_MASK
    }

    pub fn sample(&self, _ctx: &LightSampleContext, _u: Float) -> Option<SampledLight> {
        todo!("Implement LightBVHNode::Sample logic")
    }
}

#[derive(Clone, Debug, Copy)]
pub struct BVHLightSampler {
    pub nodes: Ptr<LightBVHNode>,
    pub lights: Ptr<Light>,
    pub infinite_lights: Ptr<Light>,
    pub bit_trails: Ptr<u64>,
    pub nodes_len: u32,
    pub lights_len: u32,
    pub infinite_lights_len: u32,
    pub all_light_bounds: Bounds3f,
}

unsafe impl Send for BVHLightSampler {}
unsafe impl Sync for BVHLightSampler {}

impl BVHLightSampler {
    #[inline(always)]
    fn node(&self, idx: usize) -> &LightBVHNode {
        unsafe { self.nodes.at(idx) }
    }

    #[inline(always)]
    fn light(&self, idx: usize) -> Light {
        unsafe { *self.lights.at(idx) }
    }

    #[inline(always)]
    fn infinite_light(&self, idx: usize) -> Light {
        unsafe { *self.infinite_lights.at(idx) }
    }

    #[inline(always)]
    fn bit_trail(&self, idx: usize) -> u64 {
        unsafe { *self.bit_trails.at(idx) }
    }

    #[inline(always)]
    fn light_index_in(&self, base: Ptr<Light>, len: u32, light: &Light) -> Option<usize> {
        let target = light as *const Light;
        for i in 0..len as usize {
            if unsafe { base.add(i) }.as_raw() == target {
                return Some(i);
            }
        }
        None
    }

    fn evaluate_cost(&self, b: &LightBounds, bounds: &Bounds3f, dim: usize) -> Float {
        let theta_o = b.cos_theta_o.acos();
        let theta_e = b.cos_theta_e.acos();
        let theta_w = (theta_o + theta_e).min(PI);
        let sin_o = safe_sqrt(1. - square(b.cos_theta_o));
        let m_omega = 2. * PI * (1. - b.cos_theta_o)
            + PI / 2.
                * (2. * theta_w - (theta_o - 2. * theta_w).cos() - 2. * theta_o * sin_o
                    + b.cos_theta_o);
        let kr = bounds.diagonal().max_component_value() / bounds.diagonal()[dim];
        b.phi * m_omega * kr * b.bounds.surface_area()
    }
}

impl LightSamplerTrait for BVHLightSampler {
    fn sample_with_context(&self, ctx: &LightSampleContext, mut u: Float) -> Option<SampledLight> {
        let empty_nodes = if self.nodes_len == 0 { 0. } else { 1. };
        let inf_size = self.infinite_lights_len as Float;
        let light_size = self.lights_len as Float;

        let p_inf = inf_size / (inf_size + empty_nodes);

        if u < p_inf {
            u /= p_inf;
            let ind = (u * light_size).min(light_size - 1.) as usize;
            let pmf = p_inf / inf_size;
            return Some(SampledLight::new(self.infinite_light(ind), pmf));
        }

        if self.nodes_len == 0 {
            return None;
        }
        let p = ctx.p();
        let n = ctx.ns;
        u = ((u - p_inf) / (1. - p_inf)).min(ONE_MINUS_EPSILON);
        let mut node_ind = 0;
        let mut pmf = 1. - p_inf;

        loop {
            let node = self.node(node_ind);
            if !node.is_leaf() {
                let child0_idx = node_ind + 1;
                let child1_idx = node.child_or_light_index() as usize;
                let child0 = self.node(child0_idx);
                let child1 = self.node(child1_idx);

                let ci: [Float; 2] = [
                    child0.light_bounds.importance(p, n, &self.all_light_bounds),
                    child1.light_bounds.importance(p, n, &self.all_light_bounds),
                ];

                if ci[0] == 0. && ci[1] == 0. {
                    return None;
                }

                let mut node_pmf: Float = 0.;
                let child = sample_discrete(&ci, u, Some(&mut node_pmf), Some(&mut u));
                pmf *= node_pmf;
                node_ind = if child == 0 { child0_idx } else { child1_idx };
            } else {
                if node_ind > 0 || node.light_bounds.importance(p, n, &self.all_light_bounds) > 0. {
                    let light_idx = node.child_or_light_index() as usize;
                    return Some(SampledLight::new(self.light(light_idx), pmf));
                }
                return None;
            }
        }
    }

    fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Light) -> Float {
        let empty_nodes = if self.nodes_len == 0 { 0. } else { 1. };
        let n_infinite = self.infinite_lights_len as Float;
        if self
            .light_index_in(self.infinite_lights, self.infinite_lights_len, light)
            .is_some()
        {
            return 1.0 / (n_infinite + empty_nodes);
        }

        let Some(light_index) = self.light_index_in(self.lights, self.lights_len, light) else {
            return 0.0;
        };

        let mut bit_trail = self.bit_trail(light_index);
        let p_inf = n_infinite / (n_infinite + empty_nodes);
        let mut pmf = 1.0 - p_inf;
        let mut node_ind = 0;
        let p = ctx.p();
        let n = ctx.ns;

        loop {
            let node = self.node(node_ind);
            if node.is_leaf() {
                return pmf;
            }
            let child0 = self.node(node_ind + 1);
            let child1 = self.node(node.child_or_light_index() as usize);
            let ci = [
                child0.light_bounds.importance(p, n, &self.all_light_bounds),
                child1.light_bounds.importance(p, n, &self.all_light_bounds),
            ];

            let sum_importance = ci[0] + ci[1];
            if sum_importance == 0.0 {
                return 0.0;
            }

            let which_child = (bit_trail & 1) as usize;

            // Update probability: prob of picking the correct child
            pmf *= ci[which_child] / sum_importance;

            // Advance
            node_ind = if which_child == 1 {
                node.child_or_light_index() as usize
            } else {
                node_ind + 1
            };

            bit_trail >>= 1;
        }
    }

    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.lights_len == 0 {
            return None;
        }

        let light_ind = (u * self.lights_len as Float).min(self.lights_len as Float - 1.) as usize;

        Some(SampledLight::new(
            self.light(light_ind),
            1. / self.lights_len as Float,
        ))
    }

    fn pmf(&self, _light: &Light) -> Float {
        if self.lights_len == 0 {
            return 0.;
        }

        1. / self.lights_len as Float
    }
}