pbrt/shared/src/spectra/simple.rs

use super::cie::*;
use super::sampled::{LAMBDA_MAX, LAMBDA_MIN};
use crate::core::spectrum::{Spectrum, SpectrumTrait};
use crate::spectra::{SampledSpectrum, SampledWavelengths, N_SPECTRUM_SAMPLES};
use crate::utils::find_interval;
use crate::utils::math::square;
use crate::{gvec, gvec_from_slice, gvec_with_capacity, Float, GVec, Ptr};
use core::hash::{Hash, Hasher};
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ConstantSpectrum {
    pub c: Float,
}

impl ConstantSpectrum {
    pub fn new(c: Float) -> Self {
        Self { c }
    }
}

impl SpectrumTrait for ConstantSpectrum {
    fn evaluate(&self, _lambda: Float) -> Float {
        self.c
    }

    fn max_value(&self) -> Float {
        self.c
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct DenselySampledSpectrum {
    pub lambda_min: i32,
    pub lambda_max: i32,
    pub values: GVec<Float>,
}

unsafe impl Send for DenselySampledSpectrum {}
unsafe impl Sync for DenselySampledSpectrum {}

impl DenselySampledSpectrum {
    pub fn new(lambda_min: i32, lambda_max: i32, values: GVec<Float>) -> Self {
        let func_integral = 0.0;
        Self {
            lambda_min,
            lambda_max,
            values,
        }
    }

    pub fn new_zero(lambda_min: i32, lambda_max: i32) -> Self {
        let n = (lambda_max - lambda_min + 1).max(0) as usize;
        let mut values = gvec_with_capacity(n);
        values.resize(n, 0.0);
        Self {
            lambda_min,
            lambda_max,
            values,
        }
    }

    pub fn from_spectrum(spec: &Spectrum) -> Self {
        let mut values = gvec_with_capacity((LAMBDA_MAX - LAMBDA_MIN + 1) as usize);
        for lambda in LAMBDA_MIN..=LAMBDA_MAX {
            values.push(spec.evaluate(lambda as Float));
        }
        Self {
            lambda_min: LAMBDA_MIN,
            lambda_max: LAMBDA_MAX,
            values,
        }
    }

    pub fn from_function<F>(f: F, lambda_min: i32, lambda_max: i32) -> Self
    where
        F: Fn(Float) -> Float,
    {
        let mut values = gvec_with_capacity((lambda_max - lambda_min + 1) as usize);
        for lambda in lambda_min..=lambda_max {
            values.push(f(lambda as Float));
        }
        Self {
            lambda_min,
            lambda_max,
            values,
        }
    }

    pub fn generate_cie_d(temperature: Float) -> Self {
        let cct = temperature * 1.4388 / 1.4380;

        if cct < 4000.0 {
            return Self::from_function(
                |lambda| BlackbodySpectrum::new(cct).evaluate(lambda),
                LAMBDA_MIN,
                LAMBDA_MAX,
            );
        }

        let x = if cct < 7000. {
            -4.607 * 1e9 / cct.powi(3) + 2.9678 * 1e6 / square(cct) + 0.09911 * 1e3 / cct + 0.244063
        } else {
            -2.0064 * 1e9 / cct.powi(3) + 1.9018 * 1e6 / square(cct) + 0.24748 * 1e3 / cct + 0.23704
        };
        let y = -3. * x + 2.87 * x - 0.275;
        let m = 0.0241 + 0.2562 * x - 0.7341 * y;
        let m1 = (-1.3515 - 1.7703 * x + 5.9114 * y) / m;
        let m2 = (0.0300 - 31.4424 * x + 30.0717 * y) / m;

        let mut coarse_values = gvec_with_capacity(N_CIES);
        for i in 0..N_CIES {
            coarse_values.push((CIE_S0[i] + CIE_S1[i] * m1 + CIE_S2[i] * m2) * 0.01);
        }

        let temp_pls = PiecewiseLinearSpectrum {
            lambdas: gvec_from_slice(&CIE_S_LAMBDA),
            values: gvec_from_slice(&coarse_values),
            count: N_CIES as u32,
        };

        Self::from_function(|lambda| temp_pls.evaluate(lambda), LAMBDA_MIN, LAMBDA_MAX)
    }

    pub fn scale(&mut self, s: Float) {
        for v in &mut self.values {
            *v *= s;
        }
    }

    #[inline(always)]
    pub fn count(&self) -> usize {
        if self.values.is_empty() {
            0
        } else {
            (self.lambda_max - self.lambda_min + 1) as usize
        }
    }

    #[inline(always)]
    pub fn value(&self, idx: u32) -> Float {
        unsafe { *self.values.as_ptr().add(idx as usize) }
    }
}

impl PartialEq for DenselySampledSpectrum {
    fn eq(&self, other: &Self) -> bool {
        self.lambda_min == other.lambda_min
            && self.lambda_max == other.lambda_max
            && self.values == other.values
    }
}

impl Eq for DenselySampledSpectrum {}

impl Hash for DenselySampledSpectrum {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.lambda_min.hash(state);
        self.lambda_max.hash(state);
        for &val in self.values.iter() {
            val.to_bits().hash(state);
        }
    }
}

impl SpectrumTrait for DenselySampledSpectrum {
    fn max_value(&self) -> Float {
        if self.values.is_empty() {
            return 0.0;
        }
        let mut max_val = Float::NEG_INFINITY;
        for i in 0..self.count() {
            let val = self.value(i as u32);
            if val > max_val {
                max_val = val;
            }
        }
        max_val
    }

    fn sample(&self, lambda: &SampledWavelengths) -> SampledSpectrum {
        let mut s = SampledSpectrum::default();
        let n = self.count() as i32;
        for i in 0..N_SPECTRUM_SAMPLES {
            let offset = lambda[i].round() as i32 - self.lambda_min;
            s[i] = if offset < 0 || offset >= n {
                0.0
            } else {
                self.value(offset as u32)
            };
        }
        s
    }

    fn evaluate(&self, lambda: Float) -> Float {
        let offset = (lambda.round() as i32) - self.lambda_min;
        let n = self.count() as i32;
        if offset < 0 || offset >= n {
            0.0
        } else {
            self.value(offset as u32)
        }
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct PiecewiseLinearSpectrum {
    pub lambdas: GVec<Float>,
    pub values: GVec<Float>,
    pub count: u32,
}

impl PiecewiseLinearSpectrum {
    #[inline(always)]
    pub fn count(&self) -> usize {
        self.count.try_into().unwrap()
    }

    #[inline(always)]
    pub fn lambda(&self, idx: u32) -> Float {
        unsafe { *self.lambdas.as_ptr().add(idx as usize) }
    }

    #[inline(always)]
    pub fn value(&self, idx: u32) -> Float {
        unsafe { *self.values.as_ptr().add(idx as usize) }
    }

    pub fn new(lambdas: GVec<Float>, values: GVec<Float>) -> Self {
        assert_eq!(lambdas.len(), values.len());
        let count = lambdas.len() as u32;
        Self {
            lambdas,
            values,
            count,
        }
    }

    pub fn from_interleaved(data: &[Float], _normalize: bool) -> Self {
        assert!(
            data.len() % 2 == 0,
            "Interleaved data must have even length"
        );

        let n = data.len() / 2;
        let mut pairs: GVec<(Float, Float)> = gvec_with_capacity(n);
        for chunk in data.chunks(2) {
            pairs.push((chunk[0], chunk[1]));
        }
        pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(core::cmp::Ordering::Equal));

        let mut lambdas = gvec_with_capacity(n);
        let mut values = gvec_with_capacity(n);
        for (l, v) in pairs.iter() {
            lambdas.push(*l);
            values.push(*v);
        }
        Self::new(lambdas, values)
    }
}

unsafe impl Send for PiecewiseLinearSpectrum {}
unsafe impl Sync for PiecewiseLinearSpectrum {}

impl SpectrumTrait for PiecewiseLinearSpectrum {
    fn evaluate(&self, lambda: Float) -> Float {
        if self.lambdas.is_empty() {
            return 0.0;
        }

        if lambda <= self.lambda(0) {
            return self.value(0);
        }

        if lambda >= self.lambda(self.count - 1) {
            return self.value(self.count - 1);
        }

        let i = find_interval(self.count, |idx| self.lambda(idx) <= lambda);

        let l0 = self.lambda(i);
        let l1 = self.lambda(i + 1);
        let v0 = self.value(i);
        let v1 = self.value(i + 1);

        let t = (lambda - l0) / (l1 - l0);
        v0 + t * (v1 - v0)
    }

    fn max_value(&self) -> Float {
        if self.values.is_empty() {
            return 0.;
        }

        let n = self.count;
        let mut max_val = Float::NEG_INFINITY;

        for i in 0..n {
            unsafe {
                let val = *self.values.as_ptr().add(i as usize);
                if val > max_val {
                    max_val = val;
                }
            }
        }
        max_val
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct BlackbodySpectrum {
    pub temperature: Float,
    pub normalization_factor: Float,
}

// Planck's Law
impl BlackbodySpectrum {
    const C: Float = 299792458.0;
    const H: Float = 6.62606957e-34;
    const KB: Float = 1.3806488e-23;

    pub fn new(temperature: Float) -> Self {
        // Physical constants
        let lambda_max = 2.8977721e-3 / temperature * 1e9;
        let max_val = Self::planck_law(lambda_max, temperature);
        Self {
            temperature,
            normalization_factor: if max_val > 0.0 { 1.0 / max_val } else { 0.0 },
        }
    }

    fn planck_law(lambda_nm: Float, temp: Float) -> Float {
        if temp <= 0.0 {
            return 0.0;
        }
        let lambda_m = lambda_nm * 1e-9;
        let c1 = 2.0 * Self::H * Self::C * Self::C;
        let c2 = (Self::H * Self::C) / Self::KB;

        let numerator = c1 / lambda_m.powi(5);
        let denominator = (c2 / (lambda_m * temp)).exp() - 1.0;

        if denominator.is_infinite() {
            0.0
        } else {
            numerator / denominator
        }
    }

    pub fn sample(&self, lambda: &SampledWavelengths) -> SampledSpectrum {
        SampledSpectrum::from_fn(|i| {
            Self::planck_law(lambda[i], self.temperature) * self.normalization_factor
        })
    }
}

impl SpectrumTrait for BlackbodySpectrum {
    fn evaluate(&self, lambda: Float) -> Float {
        Self::planck_law(lambda, self.temperature) * self.normalization_factor
    }

    fn max_value(&self) -> Float {
        let lambda_max = 2.8977721e-3 / self.temperature * 1e9;
        Self::planck_law(lambda_max, self.temperature)
    }
}