pbrt/shared/src/core/sampler.rs

use crate::core::filter::FilterTrait;
use crate::core::geometry::{Bounds2f, Point2f, Point2i, Vector2f};
use crate::core::pbrt::{Float, ONE_MINUS_EPSILON, PI, PI_OVER_2, PI_OVER_4};
use crate::utils::math::{
    clamp, encode_morton_2, inverse_radical_inverse, lerp, log2_int,
    owen_scrambled_radical_inverse, permutation_element, radical_inverse, round_up_pow2,
    scrambled_radical_inverse, sobol_interval_to_index, sobol_sample, BinaryPermuteScrambler,
    DigitPermutation, FastOwenScrambler, NoRandomizer, OwenScrambler, Scrambler, PRIME_TABLE_SIZE,
};
use crate::utils::rng::Rng;
use crate::utils::sobol::N_SOBOL_DIMENSIONS;
use crate::utils::{hash::*, sobol};
use crate::{gvec, GVec, Ptr};
use enum_dispatch::enum_dispatch;

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct CameraSample {
    pub p_film: Point2f,
    pub p_lens: Point2f,
    pub time: Float,
    pub filter_weight: Float,
}

pub fn get_camera_sample<S, F>(sampler: &mut S, p_pixel: Point2i, filter: &F) -> CameraSample
where
    S: SamplerTrait,
    F: FilterTrait,
{
    let fs = filter.sample(sampler.get_pixel2d());
    CameraSample {
        p_film: Point2f::from(p_pixel) + Vector2f::from(fs.p) + Vector2f::new(0.5, 0.5),
        time: sampler.get1d(),
        p_lens: sampler.get2d(),
        filter_weight: fs.weight,
    }
}

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct IndependentSampler {
    pub samples_per_pixel: i32,
    pub seed: u64,
    pub rng: Rng,
}

impl IndependentSampler {
    pub fn new(samples_per_pixel: i32, seed: u64) -> Self {
        Self {
            samples_per_pixel,
            seed,
            rng: Rng::default(),
        }
    }
}

impl SamplerTrait for IndependentSampler {
    fn samples_per_pixel(&self) -> i32 {
        self.samples_per_pixel
    }
    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        let hash_input = [p.x() as u64, p.y() as u64, self.seed];
        let sequence_index = hash_buffer(&hash_input, 0);
        self.rng.set_sequence(sequence_index);
        self.rng
            .advance((sample_index as u64) * 65536 + (dim.unwrap_or(0) as u64));
    }

    fn get1d(&mut self) -> Float {
        self.rng.uniform::<Float>()
    }
    fn get2d(&mut self) -> Point2f {
        Point2f::new(self.rng.uniform::<Float>(), self.rng.uniform::<Float>())
    }
    fn get_pixel2d(&mut self) -> Point2f {
        self.get2d()
    }
}

pub const MAX_HALTON_RESOLUTION: i32 = 128;

#[repr(C)]
#[derive(Debug, Default, Clone, PartialEq, Eq, Copy)]
pub enum RandomizeStrategy {
    #[default]
    None,
    PermuteDigits,
    FastOwen,
    Owen,
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct HaltonSampler {
    pub samples_per_pixel: i32,
    pub randomize: RandomizeStrategy,
    pub base_scales: [u64; 2],
    pub base_exponents: [u64; 2],
    pub mult_inverse: [u64; 2],
    pub halton_index: u64,
    pub dim: u32,
    pub digit_permutations: Ptr<DigitPermutation>,
}

#[allow(clippy::derivable_impls)]
impl Default for HaltonSampler {
    fn default() -> Self {
        Self {
            samples_per_pixel: 0,
            randomize: RandomizeStrategy::default(),
            base_scales: [0; 2],
            base_exponents: [0; 2],
            mult_inverse: [0; 2],
            halton_index: 0,
            dim: 0,
            digit_permutations: Ptr::default(),
        }
    }
}

impl HaltonSampler {
    pub fn sample_dimension(&self, dimension: u32) -> Float {
        if self.randomize == RandomizeStrategy::None {
            radical_inverse(dimension, self.halton_index)
        } else if self.randomize == RandomizeStrategy::PermuteDigits {
            let digit_perm = unsafe { &*self.digit_permutations.add(dimension as usize) };
            scrambled_radical_inverse(dimension, self.halton_index, digit_perm)
        } else {
            owen_scrambled_radical_inverse(
                dimension,
                self.halton_index,
                mix_bits(1 + ((dimension as u64) << 4)) as u32,
            )
        }
    }

    pub fn multiplicative_inverse(a: i64, n: i64) -> u64 {
        let (x, _) = Self::extended_gcd(a as u64, n as u64);
        x.rem_euclid(n) as u64
    }

    pub fn extended_gcd(a: u64, b: u64) -> (i64, i64) {
        if b == 0 {
            return (1, 0);
        }

        let (xp, yp) = Self::extended_gcd(b, a % b);

        let d = (a / b) as i64;

        (yp, xp - d * yp)
    }
}

impl SamplerTrait for HaltonSampler {
    fn samples_per_pixel(&self) -> i32 {
        self.samples_per_pixel
    }

    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        self.halton_index = 0;

        let sample_stride = self.base_scales[0] * self.base_scales[1];
        if sample_stride > 1 {
            let pm_x = p.x().rem_euclid(MAX_HALTON_RESOLUTION) as u64;
            let pm_y = p.y().rem_euclid(MAX_HALTON_RESOLUTION) as u64;

            let dim_offset_x = inverse_radical_inverse(pm_x, 2, self.base_exponents[0]);

            self.halton_index = self.halton_index.wrapping_add(
                dim_offset_x
                    .wrapping_mul(sample_stride / self.base_scales[0])
                    .wrapping_mul(self.mult_inverse[0]),
            );

            let dim_offset_y = inverse_radical_inverse(pm_y, 3, self.base_exponents[1]);

            self.halton_index = self.halton_index.wrapping_add(
                dim_offset_y
                    .wrapping_mul(sample_stride / self.base_scales[1])
                    .wrapping_mul(self.mult_inverse[1]),
            );

            self.halton_index %= sample_stride;
        }

        self.halton_index = self
            .halton_index
            .wrapping_add((sample_index as u64).wrapping_mul(sample_stride));

        self.dim = 2.max(dim.unwrap_or(0));
    }

    fn get1d(&mut self) -> Float {
        if self.dim + 1 >= PRIME_TABLE_SIZE as u32 {
            self.dim = 2;
        }
        self.dim += 1;
        self.sample_dimension(self.dim)
    }

    fn get2d(&mut self) -> Point2f {
        if self.dim > PRIME_TABLE_SIZE as u32 {
            self.dim = 2;
        }
        let dim = self.dim;
        self.dim += 2;
        Point2f::new(self.sample_dimension(dim), self.sample_dimension(dim + 1))
    }

    fn get_pixel2d(&mut self) -> Point2f {
        Point2f::new(
            radical_inverse(0, self.halton_index >> self.base_exponents[0]),
            radical_inverse(1, self.halton_index / self.base_scales[1]),
        )
    }
}

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct StratifiedSampler {
    x_pixel_samples: i32,
    y_pixel_samples: i32,
    jitter: bool,
    seed: u64,
    rng: Rng,
    pixel: Point2i,
    sample_index: i32,
    dim: u32,
}

impl StratifiedSampler {
    pub fn new(
        x_pixel_samples: i32,
        y_pixel_samples: i32,
        seed: Option<u64>,
        jitter: bool,
    ) -> Self {
        Self {
            x_pixel_samples,
            y_pixel_samples,
            jitter,
            seed: seed.unwrap_or(0),
            rng: Rng::default(),
            pixel: Point2i::default(),
            sample_index: 0,
            dim: 0,
        }
    }
}

impl SamplerTrait for StratifiedSampler {
    fn samples_per_pixel(&self) -> i32 {
        self.x_pixel_samples * self.y_pixel_samples
    }

    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        self.pixel = p;
        self.sample_index = sample_index;
        let hash_input = [p.x() as u64, p.y() as u64, self.seed];
        let sequence_index = hash_buffer(&hash_input, 0);
        self.rng.set_sequence(sequence_index);
        self.rng
            .advance((sample_index as u64) * 65536 + (dim.unwrap_or(0) as u64));
    }

    fn get1d(&mut self) -> Float {
        let hash_input = [
            self.pixel.x() as u64,
            self.pixel.y() as u64,
            self.dim as u64,
            self.seed,
        ];
        let hash = hash_buffer(&hash_input, 0);
        let stratum = permutation_element(
            self.sample_index as u32,
            self.samples_per_pixel() as u32,
            hash as u32,
        );
        self.dim += 1;

        let delta = if self.jitter {
            self.rng.uniform::<Float>()
        } else {
            0.5
        };
        (stratum as Float + delta) / (self.samples_per_pixel() as Float)
    }

    fn get2d(&mut self) -> Point2f {
        let hash_input = [
            self.pixel.x() as u64,
            self.pixel.y() as u64,
            self.dim as u64,
            self.seed,
        ];
        let hash = hash_buffer(&hash_input, 0);
        let stratum = permutation_element(
            self.sample_index as u32,
            self.samples_per_pixel() as u32,
            hash as u32,
        );
        self.dim += 2;
        let x = stratum % self.x_pixel_samples as u32;
        let y = stratum / self.y_pixel_samples as u32;
        let dx = if self.jitter {
            self.rng.uniform::<Float>()
        } else {
            0.5
        };
        let dy = if self.jitter {
            self.rng.uniform::<Float>()
        } else {
            0.5
        };
        Point2f::new(
            (x as Float + dx) / self.x_pixel_samples as Float,
            (y as Float + dy) / self.y_pixel_samples as Float,
        )
    }

    fn get_pixel2d(&mut self) -> Point2f {
        self.get2d()
    }
}

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct PaddedSobolSampler {
    samples_per_pixel: i32,
    seed: u64,
    randomize: RandomizeStrategy,
    pixel: Point2i,
    sample_index: i32,
    dim: u32,
}

impl PaddedSobolSampler {
    pub fn new(samples_per_pixel: i32, randomize: RandomizeStrategy, seed: Option<u64>) -> Self {
        Self {
            samples_per_pixel,
            seed: seed.unwrap_or(0),
            randomize,
            pixel: Point2i::default(),
            sample_index: 0,
            dim: 0,
        }
    }

    fn sample_dimension(&self, dimension: u32, a: u32, hash: u32) -> Float {
        if self.randomize == RandomizeStrategy::None {
            return sobol_sample(a as u64, dimension, NoRandomizer);
        }
        match self.randomize {
            RandomizeStrategy::PermuteDigits => {
                sobol_sample(a as u64, dimension, BinaryPermuteScrambler::new(hash))
            }
            RandomizeStrategy::FastOwen => {
                sobol_sample(a as u64, dimension, FastOwenScrambler::new(hash))
            }
            RandomizeStrategy::Owen => sobol_sample(a as u64, dimension, OwenScrambler::new(hash)),
            RandomizeStrategy::None => unreachable!(),
        }
    }
}

impl SamplerTrait for PaddedSobolSampler {
    fn samples_per_pixel(&self) -> i32 {
        self.samples_per_pixel
    }
    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        self.pixel = p;
        self.sample_index = sample_index;
        self.dim = dim.unwrap_or(0);
    }

    fn get1d(&mut self) -> Float {
        let hash_input = [
            self.pixel.x() as u64,
            self.pixel.y() as u64,
            self.dim as u64,
            self.seed,
        ];
        let hash = hash_buffer(&hash_input, 0);
        let index = permutation_element(
            self.sample_index as u32,
            self.samples_per_pixel as u32,
            hash as u32,
        );
        self.sample_dimension(0, index, (hash >> 32) as u32)
    }
    fn get2d(&mut self) -> Point2f {
        let hash_input = [
            self.pixel.x() as u64,
            self.pixel.y() as u64,
            self.dim as u64,
            self.seed,
        ];
        let hash = hash_buffer(&hash_input, 0);
        let index = permutation_element(
            self.sample_index as u32,
            self.samples_per_pixel as u32,
            hash as u32,
        );
        self.dim += 2;
        Point2f::new(
            self.sample_dimension(0, index, hash as u32),
            self.sample_dimension(1, index, (hash >> 32) as u32),
        )
    }

    fn get_pixel2d(&mut self) -> Point2f {
        self.get2d()
    }
}

#[derive(Default, Debug, Clone)]
pub struct SobolSampler {
    samples_per_pixel: i32,
    scale: i32,
    seed: u64,
    randomize: RandomizeStrategy,
    pixel: Point2i,
    dim: u32,
    sobol_index: u64,
}

impl SobolSampler {
    pub fn new(
        samples_per_pixel: i32,
        full_resolution: Point2i,
        randomize: RandomizeStrategy,
        seed: Option<u64>,
    ) -> Self {
        let scale = round_up_pow2(full_resolution.x().max(full_resolution.y()));
        Self {
            samples_per_pixel,
            scale,
            seed: seed.unwrap_or(0),
            randomize,
            pixel: Point2i::default(),
            dim: 0,
            sobol_index: 0,
        }
    }

    fn sample_dimension(&self, dimension: u32) -> Float {
        if self.randomize == RandomizeStrategy::None {
            return sobol_sample(self.sobol_index, dimension, NoRandomizer);
        }
        let hash_input = [self.pixel.x() as u64, self.pixel.y() as u64, self.seed];
        let hash = hash_buffer(&hash_input, 0) as u32;
        match self.randomize {
            RandomizeStrategy::PermuteDigits => sobol_sample(
                self.sobol_index,
                dimension,
                BinaryPermuteScrambler::new(hash),
            ),
            RandomizeStrategy::FastOwen => {
                sobol_sample(self.sobol_index, dimension, FastOwenScrambler::new(hash))
            }
            RandomizeStrategy::Owen => {
                sobol_sample(self.sobol_index, dimension, OwenScrambler::new(hash))
            }
            RandomizeStrategy::None => unreachable!(),
        }
    }
}

impl SamplerTrait for SobolSampler {
    fn samples_per_pixel(&self) -> i32 {
        self.samples_per_pixel
    }
    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        self.pixel = p;
        self.dim = 2.max(dim.unwrap_or(0));
        self.sobol_index =
            sobol_interval_to_index(log2_int(self.scale as Float) as u32, sample_index as u64, p)
    }

    fn get1d(&mut self) -> Float {
        if self.dim >= N_SOBOL_DIMENSIONS as u32 {
            self.dim = 2;
        }

        let dim = self.dim;
        self.dim += 1;
        self.sample_dimension(dim)
    }

    fn get2d(&mut self) -> Point2f {
        if self.dim >= N_SOBOL_DIMENSIONS as u32 {
            self.dim = 2;
        }
        let u = Point2f::new(
            self.sample_dimension(self.dim),
            self.sample_dimension(self.dim + 1),
        );
        self.dim += 2;
        u
    }

    fn get_pixel2d(&mut self) -> Point2f {
        let mut u = Point2f::new(
            sobol_sample(self.sobol_index, 0, NoRandomizer),
            sobol_sample(self.sobol_index, 1, NoRandomizer),
        );
        u[0] = clamp(
            u[0] * self.scale as Float - self.pixel[0] as Float,
            0.,
            ONE_MINUS_EPSILON,
        ) as Float;
        u[1] = clamp(
            u[1] * self.scale as Float - self.pixel[1] as Float,
            1.,
            ONE_MINUS_EPSILON,
        ) as Float;
        u
    }
}

#[repr(C)]
#[derive(Default, Copy, Debug, Clone)]
pub struct ZSobolSampler {
    randomize: RandomizeStrategy,
    seed: u64,
    log2_samples_per_pixel: i32,
    n_base4_digits: u32,
    morton_index: u64,
    dim: u32,
}

impl ZSobolSampler {
    pub fn new(
        samples_per_pixel: i32,
        full_resolution: Point2i,
        randomize: RandomizeStrategy,
        seed: Option<u64>,
    ) -> Self {
        let log2_samples_per_pixel = log2_int(samples_per_pixel as Float) as u32;
        let res = round_up_pow2(full_resolution.x().max(full_resolution.y()));
        let log4_samples_per_pixel = log2_samples_per_pixel.div_ceil(2);
        let n_base4_digits = log2_int(res as Float) as u32 + log4_samples_per_pixel as u32;
        Self {
            randomize,
            seed: seed.unwrap_or(0),
            log2_samples_per_pixel: log2_samples_per_pixel as i32,
            n_base4_digits,
            morton_index: 0,
            dim: 0,
        }
    }

    fn get_sample_index(&self) -> u64 {
        const PERMUTATIONS: [[u8; 4]; 24] = [
            [0, 1, 2, 3],
            [0, 1, 3, 2],
            [0, 2, 1, 3],
            [0, 2, 3, 1],
            [0, 3, 2, 1],
            [0, 3, 1, 2],
            [1, 0, 2, 3],
            [1, 0, 3, 2],
            [1, 2, 0, 3],
            [1, 2, 3, 0],
            [1, 3, 2, 0],
            [1, 3, 0, 2],
            [2, 1, 0, 3],
            [2, 1, 3, 0],
            [2, 0, 1, 3],
            [2, 0, 3, 1],
            [2, 3, 0, 1],
            [2, 3, 1, 0],
            [3, 1, 2, 0],
            [3, 1, 0, 2],
            [3, 2, 1, 0],
            [3, 2, 0, 1],
            [3, 0, 2, 1],
            [3, 0, 1, 2],
        ];

        let mut sample_index = 0;
        let pow2_samples = (self.log2_samples_per_pixel & 1) != 0;
        let last_digit = if pow2_samples { 1 } else { 0 };

        for i in (last_digit..self.n_base4_digits).rev() {
            let digit_shift = (2 * i) - if pow2_samples { 1 } else { 0 };

            let mut digit = (self.morton_index >> digit_shift) & 3;

            let higher_digits = self.morton_index >> (digit_shift + 2);

            let mix_input = higher_digits ^ (0x55555555 * self.dim as u64);
            let p = (mix_bits(mix_input) >> 24) % 24;

            digit = PERMUTATIONS[p as usize][digit as usize] as u64;

            sample_index |= digit << digit_shift;
        }

        if pow2_samples {
            let lsb = self.morton_index & 1;
            sample_index |= lsb;
        }

        sample_index
    }
}

impl SamplerTrait for ZSobolSampler {
    fn samples_per_pixel(&self) -> i32 {
        todo!()
    }
    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
        self.dim = dim.unwrap_or(0);
        self.morton_index = (encode_morton_2(p.x() as u32, p.y() as u32)
            << self.log2_samples_per_pixel)
            | sample_index as u64
    }

    fn get1d(&mut self) -> Float {
        let sample_index = self.get_sample_index();
        let hash_input = [self.dim as u64, self.seed];
        let hash = hash_buffer(&hash_input, 0) as u32;
        self.dim += 1;
        if self.randomize == RandomizeStrategy::None {
            return sobol_sample(sample_index, self.dim, NoRandomizer);
        }
        match self.randomize {
            RandomizeStrategy::PermuteDigits => {
                sobol_sample(sample_index, self.dim, BinaryPermuteScrambler::new(hash))
            }
            RandomizeStrategy::FastOwen => {
                sobol_sample(sample_index, self.dim, FastOwenScrambler::new(hash))
            }
            RandomizeStrategy::Owen => {
                sobol_sample(sample_index, self.dim, OwenScrambler::new(hash))
            }
            RandomizeStrategy::None => unreachable!(),
        }
    }

    fn get2d(&mut self) -> Point2f {
        let sample_index = self.get_sample_index();
        self.dim += 2;
        let hash_input = [self.dim as u64, self.seed];
        let hash = hash_buffer(&hash_input, 0);
        let sample_hash = [hash as u32, (hash >> 32) as u32];
        if self.randomize == RandomizeStrategy::None {
            return Point2f::new(
                sobol_sample(sample_index, 0, NoRandomizer),
                sobol_sample(sample_index, 1, NoRandomizer),
            );
        }
        match self.randomize {
            RandomizeStrategy::PermuteDigits => Point2f::new(
                sobol_sample(sample_index, 0, BinaryPermuteScrambler::new(sample_hash[0])),
                sobol_sample(sample_index, 1, BinaryPermuteScrambler::new(sample_hash[1])),
            ),
            RandomizeStrategy::FastOwen => Point2f::new(
                sobol_sample(sample_index, 0, FastOwenScrambler::new(sample_hash[0])),
                sobol_sample(sample_index, 1, FastOwenScrambler::new(sample_hash[1])),
            ),
            RandomizeStrategy::Owen => Point2f::new(
                sobol_sample(sample_index, 0, OwenScrambler::new(sample_hash[0])),
                sobol_sample(sample_index, 1, OwenScrambler::new(sample_hash[1])),
            ),
            RandomizeStrategy::None => unreachable!(),
        }
    }

    fn get_pixel2d(&mut self) -> Point2f {
        self.get2d()
    }
}

#[derive(Default, Debug, Clone)]
pub struct MLTSampler;
impl SamplerTrait for MLTSampler {
    fn samples_per_pixel(&self) -> i32 {
        todo!()
    }
    fn start_pixel_sample(&mut self, _p: Point2i, _sample_index: i32, _dim: Option<u32>) {
        todo!()
    }
    fn get1d(&mut self) -> Float {
        todo!()
    }
    fn get2d(&mut self) -> Point2f {
        todo!()
    }
    fn get_pixel2d(&mut self) -> Point2f {
        todo!()
    }
}

#[enum_dispatch]
pub trait SamplerTrait {
    fn samples_per_pixel(&self) -> i32;
    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>);
    fn get1d(&mut self) -> Float;
    fn get2d(&mut self) -> Point2f;
    fn get_pixel2d(&mut self) -> Point2f;
}

#[enum_dispatch(SamplerTrait)]
#[derive(Debug, Clone)]
pub enum Sampler {
    Independent(IndependentSampler),
    Stratified(StratifiedSampler),
    Halton(HaltonSampler),
    PaddedSobol(PaddedSobolSampler),
    Sobol(SobolSampler),
    ZSobol(ZSobolSampler),
    MLT(MLTSampler),
}