pbrt/src/utils/sampling.rs

use crate::core::image::Image;
use crate::utils::Arena;
use crate::utils::containers::Array2D;
use shared::Float;
use shared::core::geometry::{Point2i, Vector2f, Vector2i};
use shared::utils::Ptr;
use shared::utils::sampling::{
    AliasTable, Bin, DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceSummedAreaTable,
    DeviceWindowedPiecewiseConstant2D, PiecewiseLinear2D,
};
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct PiecewiseConstant1D {
    func: Box<[Float]>,
    cdf: Box<[Float]>,
    pub device: DevicePiecewiseConstant1D,
}

impl PiecewiseConstant1D {
    // Constructors
    pub fn new(f: &[Float]) -> Self {
        Self::new_with_bounds(f.to_vec(), 0.0, 1.0)
    }

    pub fn to_shared(&self, arena: &mut Arena) -> DevicePiecewiseConstant1D {
        let (func_ptr, _) = arena.alloc_slice(&self.func);
        let (cdf_ptr, _) = arena.alloc_slice(&self.cdf);

        DevicePiecewiseConstant1D {
            func: func_ptr,
            cdf: cdf_ptr,
            func_integral: self.func_integral,
            n: self.func.len() as u32,
            min: self.min,
            max: self.max,
        }
    }

    pub fn from_func<F>(f: F, min: Float, max: Float, n: usize) -> Self
    where
        F: Fn(Float) -> Float,
    {
        let delta = (max - min) / n as Float;
        let values: Vec<Float> = (0..n)
            .map(|i| {
                let x = min + (i as Float + 0.5) * delta;
                f(x)
            })
            .collect();

        Self::new_with_bounds(values, min, max)
    }

    pub fn new_with_bounds(f: Vec<Float>, min: Float, max: Float) -> Self {
        let n = f.len();
        let mut cdf = Vec::with_capacity(n + 1);
        cdf.push(0.0);

        let delta = (max - min) / n as Float;
        for i in 0..n {
            cdf.push(cdf[i] + f[i] * delta);
        }

        let func_integral = cdf[n];

        if func_integral > 0.0 {
            for c in &mut cdf {
                *c /= func_integral;
            }
        }

        // Convert to boxed slices (no more reallocation possible)
        let func: Box<[Float]> = f.into_boxed_slice();
        let cdf: Box<[Float]> = cdf.into_boxed_slice();

        let device = DevicePiecewiseConstant1D {
            func: func.as_ptr().into(),
            cdf: cdf.as_ptr().into(),
            min,
            max,
            n: n as u32,
            func_integral,
        };

        Self { func, cdf, device }
    }

    // Accessors
    pub fn min(&self) -> Float {
        self.device.min
    }
    pub fn max(&self) -> Float {
        self.device.max
    }
    pub fn n(&self) -> usize {
        self.device.n as usize
    }
    pub fn integral(&self) -> Float {
        self.device.func_integral
    }

    pub fn func(&self) -> &[Float] {
        &self.func
    }
    pub fn cdf(&self) -> &[Float] {
        &self.cdf
    }
}

impl std::ops::Deref for PiecewiseConstant1D {
    type Target = DevicePiecewiseConstant1D;

    fn deref(&self) -> &Self::Target {
        &self.device
    }
}

#[derive(Debug, Clone)]
pub struct PiecewiseConstant2D {
    pub conditionals: Vec<PiecewiseConstant1D>,
    pub marginal: PiecewiseConstant1D,
    pub conditional_devices: Box<[DevicePiecewiseConstant1D]>,
    pub device: DevicePiecewiseConstant2D,
}

impl std::ops::Deref for PiecewiseConstant2D {
    type Target = DevicePiecewiseConstant2D;

    fn deref(&self) -> &Self::Target {
        &self.device
    }
}

impl PiecewiseConstant2D {
    pub fn new(data: &[Float], n_u: usize, n_v: usize) -> Self {
        assert_eq!(data.len(), n_u * n_v);

        // Build conditional distributions p(u|v) for each row
        let mut conditionals = Vec::with_capacity(n_v);
        let mut marginal_func = Vec::with_capacity(n_v);

        for v in 0..n_v {
            let row_start = v * n_u;
            let row: Vec<Float> = data[row_start..row_start + n_u].to_vec();

            let conditional = PiecewiseConstant1D::new_with_bounds(row, 0.0, 1.0);
            marginal_func.push(conditional.integral());
            conditionals.push(conditional);
        }

        // Build marginal distribution p(v)
        let marginal = PiecewiseConstant1D::new_with_bounds(marginal_func, 0.0, 1.0);

        // Create array of device structs
        let conditional_devices: Box<[DevicePiecewiseConstant1D]> = conditionals
            .iter()
            .map(|c| c.device)
            .collect::<Vec<_>>()
            .into_boxed_slice();

        let device = DevicePiecewiseConstant2D {
            conditionals: conditional_devices.as_ptr().into(),
            marginal: marginal.device,
            n_u: n_u as u32,
            n_v: n_v as u32,
        };

        Self {
            conditionals,
            marginal,
            conditional_devices,
            device,
        }
    }

    pub fn from_image(image: &Image) -> Self {
        let res = image.resolution();
        let n_u = res.x() as usize;
        let n_v = res.y() as usize;

        let mut data = Vec::with_capacity(n_u * n_v);

        for v in 0..n_v {
            for u in 0..n_u {
                let p = Point2i::new(u as i32, v as i32);
                let luminance = image.get_channels(p).average();
                data.push(luminance);
            }
        }

        Self::new(&data, n_u, n_v)
    }

    pub fn integral(&self) -> Float {
        self.marginal.integral()
    }
}

struct PiecewiseLinear2DStorage<const N: usize> {
    data: Vec<Float>,
    marginal_cdf: Vec<Float>,
    conditional_cdf: Vec<Float>,
    param_values: [Vec<Float>; N],
}

pub struct PiecewiseLinear2DHost<const N: usize> {
    pub view: PiecewiseLinear2D<N>,
    _storage: Arc<PiecewiseLinear2DStorage<N>>,
}

impl<const N: usize> PiecewiseLinear2DHost<N> {
    pub fn new(
        data: &[Float],
        x_size: i32,
        y_size: i32,
        param_res: [usize; N],
        param_values: [&[Float]; N],
        normalize: bool,
        build_cdf: bool,
    ) -> Self {
        if build_cdf && !normalize {
            panic!("PiecewiseLinear2D::new: build_cdf implies normalize=true");
        }

        let size = Vector2i::new(x_size, y_size);
        let inv_patch_size = Vector2f::new(1. / (x_size - 1) as Float, 1. / (y_size - 1) as Float);

        let mut param_size = [0u32; N];
        let mut param_strides = [0u32; N];
        let param_values = std::array::from_fn(|i| param_values[i].to_vec());

        let mut slices: u32 = 1;
        for i in (0..N).rev() {
            if param_res[i] < 1 {
                panic!("PiecewiseLinear2D::new: parameter resolution must be >= 1!");
            }
            param_size[i] = param_res[i] as u32;
            param_strides[i] = if param_res[i] > 1 { slices } else { 0 };
            slices *= param_size[i];
        }

        let n_values = (x_size * y_size) as usize;
        let mut new_data = vec![0.0; slices as usize * n_values];
        let mut marginal_cdf = if build_cdf {
            vec![0.0; slices as usize * y_size as usize]
        } else {
            Vec::new()
        };
        let mut conditional_cdf = if build_cdf {
            vec![0.0; slices as usize * n_values]
        } else {
            Vec::new()
        };

        let mut data_offset = 0;
        for slice in 0..slices as usize {
            let slice_offset = slice * n_values;
            let current_data = &data[data_offset..data_offset + n_values];
            let mut sum = 0.;

            // Construct conditional CDF
            if normalize {
                for y in 0..(y_size - 1) {
                    for x in 0..(x_size - 1) {
                        let i = (y * x_size + x) as usize;
                        let v00 = current_data[i] as f64;
                        let v10 = current_data[i + 1] as f64;
                        let v01 = current_data[i + x_size as usize] as f64;
                        let v11 = current_data[i + 1 + x_size as usize] as f64;
                        sum += 0.25 * (v00 + v10 + v01 + v11);
                    }
                }
            }

            let normalization = if normalize && sum > 0.0 {
                1.0 / sum as Float
            } else {
                1.0
            };
            for k in 0..n_values {
                new_data[slice_offset + k] = current_data[k] * normalization;
            }

            if build_cdf {
                let marginal_slice_offset = slice * y_size as usize;
                // Construct marginal CDF
                for y in 0..y_size as usize {
                    let mut cdf_sum = 0.0;
                    let i_base = y * x_size as usize;
                    conditional_cdf[slice_offset + i_base] = 0.0;
                    for x in 0..(x_size - 1) as usize {
                        let i = i_base + x;
                        cdf_sum +=
                            0.5 * (new_data[slice_offset + i] + new_data[slice_offset + i + 1]);
                        conditional_cdf[slice_offset + i + 1] = cdf_sum;
                    }
                }
                // Construct marginal CDF
                marginal_cdf[marginal_slice_offset] = 0.0;
                let mut marginal_sum = 0.0;
                for y in 0..(y_size - 1) as usize {
                    let cdf1 = conditional_cdf[slice_offset + (y + 1) * x_size as usize - 1];
                    let cdf2 = conditional_cdf[slice_offset + (y + 2) * x_size as usize - 1];
                    marginal_sum += 0.5 * (cdf1 + cdf2);
                    marginal_cdf[marginal_slice_offset + y + 1] = marginal_sum;
                }
            }
            data_offset += n_values;
        }

        let view = PiecewiseLinear2D {
            size,
            inv_patch_size,
            param_size,
            param_strides,
            param_values: param_values.each_ref().map(|x| x.as_ptr().into()),
            data: Ptr::null(),
            marginal_cdf: marginal_cdf.as_ptr().into(),
            conditional_cdf: conditional_cdf.as_ptr().into(),
        };

        let storage = Arc::new(PiecewiseLinear2DStorage {
            data: new_data,
            marginal_cdf,
            conditional_cdf,
            param_values,
        });

        let mut final_view = view;
        final_view.data = storage.data.as_ptr().into();
        final_view.marginal_cdf = storage.marginal_cdf.as_ptr().into();
        final_view.conditional_cdf = storage.conditional_cdf.as_ptr().into();
        for i in 0..N {
            final_view.param_values[i] = storage.param_values[i].as_ptr().into();
        }

        Self {
            view: final_view,
            _storage: storage,
        }
    }
}

#[derive(Debug, Clone)]
pub struct AliasTableHost {
    pub view: AliasTable,
    _storage: Vec<Bin>,
}

impl AliasTableHost {
    pub fn new(weights: &[Float]) -> Self {
        let n = weights.len();
        if n == 0 {
            return Self {
                view: AliasTable {
                    bins: Ptr::null(),
                    size: 0,
                },
                _storage: Vec::new(),
            };
        }

        let sum: f64 = weights.iter().map(|&w| w as f64).sum();
        assert!(sum > 0.0, "Sum of weights must be positive");
        let mut bins = Vec::with_capacity(n);
        for &w in weights {
            bins.push(Bin {
                p: (w as f64 / sum) as Float,
                q: 0.0,
                alias: 0,
            });
        }

        struct Outcome {
            p_hat: f64,
            index: usize,
        }

        let mut under = Vec::with_capacity(n);
        let mut over = Vec::with_capacity(n);

        for (i, bin) in bins.iter().enumerate() {
            let p_hat = (bin.p as f64) * (n as f64);
            if p_hat < 1.0 {
                under.push(Outcome { p_hat, index: i });
            } else {
                over.push(Outcome { p_hat, index: i });
            }
        }

        while !under.is_empty() && !over.is_empty() {
            let un = under.pop().unwrap();
            let ov = over.pop().unwrap();

            bins[un.index].q = un.p_hat as Float;
            bins[un.index].alias = ov.index as u32;

            let p_excess = un.p_hat + ov.p_hat - 1.0;

            if p_excess < 1.0 {
                under.push(Outcome {
                    p_hat: p_excess,
                    index: ov.index,
                });
            } else {
                over.push(Outcome {
                    p_hat: p_excess,
                    index: ov.index,
                });
            }
        }

        while let Some(ov) = over.pop() {
            bins[ov.index].q = 1.0;
            bins[ov.index].alias = ov.index as u32;
        }

        while let Some(un) = under.pop() {
            bins[un.index].q = 1.0;
            bins[un.index].alias = un.index as u32;
        }

        let view = AliasTable {
            bins: bins.as_ptr().into(),
            size: bins.len() as u32,
        };

        Self {
            view,
            _storage: bins,
        }
    }
}

#[derive(Clone, Debug)]
pub struct SummedAreaTable {
    pub device: DeviceSummedAreaTable,
    sum: Array2D<f64>,
}

impl std::ops::Deref for SummedAreaTable {
    type Target = DeviceSummedAreaTable;
    fn deref(&self) -> &Self::Target {
        &self.device
    }
}

impl SummedAreaTable {
    pub fn new(values: &Array2D<Float>) -> Self {
        let width = values.x_size() as i32;
        let height = values.y_size() as i32;

        let mut sum = Array2D::<f64>::new_dims(width, height);
        sum[(0, 0)] = values[(0, 0)] as f64;

        for x in 1..width {
            sum[(x, 0)] = values[(x, 0)] as f64 + sum[(x - 1, 0)];
        }

        for y in 1..height {
            sum[(0, y)] = values[(0, y)] as f64 + sum[(0, y - 1)];
        }

        for y in 1..height {
            for x in 1..width {
                let term = values[(x, y)] as f64;
                let left = sum[(x - 1, y)];
                let up = sum[(x, y - 1)];
                let diag = sum[(x - 1, y - 1)];

                sum[(x, y)] = term + left + up - diag;
            }
        }

        let device = DeviceSummedAreaTable { sum: *sum };
        Self { device, sum }
    }
}

#[derive(Clone, Debug)]
pub struct WindowedPiecewiseConstant2D {
    pub device: DeviceWindowedPiecewiseConstant2D,
    sat: DeviceSummedAreaTable,
    func: Array2D<Float>,
}

impl std::ops::Deref for WindowedPiecewiseConstant2D {
    type Target = DeviceWindowedPiecewiseConstant2D;
    fn deref(&self) -> &Self::Target {
        &self.device
    }
}

impl WindowedPiecewiseConstant2D {
    pub fn new(func: Array2D<Float>) -> Self {
        let sat = *SummedAreaTable::new(&func);
        let device = DeviceWindowedPiecewiseConstant2D { sat, func: *func };

        Self { sat, func, device }
    }
}