pbrt/shared/src/lights/infinite.rs

use crate::{
    core::{
        geometry::{Frame, VectorLike},
        interaction::InteractionBase,
    },
    spectra::{RGBColorSpace, RGBIlluminantSpectrum},
    utils::{
        math::{clamp, equal_area_sphere_to_square, equal_area_square_to_sphere, square},
        sampling::{
            AliasTable, PiecewiseConstant2D, WindowedPiecewiseConstant2D, sample_uniform_sphere,
            uniform_sphere_pdf,
        },
    },
};

use crate::core::color::RGB;
use crate::core::geometry::{
    Bounds2f, Bounds3f, Normal3f, Point2f, Point2i, Point3f, Ray, Vector2f, Vector3f,
};
use crate::core::image::{Image, PixelFormat, WrapMode};
use crate::core::interaction::{Interaction, SimpleInteraction};
use crate::core::light::{
    LightBase, LightBounds, LightLiSample, LightSampleContext, LightTrait, LightType,
};
use crate::core::medium::{Medium, MediumInterface};
use crate::core::spectrum::{Spectrum, SpectrumTrait};
use crate::spectra::{DenselySampledSpectrum, SampledSpectrum, SampledWavelengths};
use crate::utils::Transform;
use crate::utils::ptr::Ptr;
use crate::{Float, PI};
use std::sync::Arc;

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct InfiniteUniformLight {
    pub base: LightBase,
    pub lemit: Ptr<DenselySampledSpectrum>,
    pub scale: Float,
    pub scene_center: Point3f,
    pub scene_radius: Float,
}

unsafe impl Send for InfiniteUniformLight {}
unsafe impl Sync for InfiniteUniformLight {}

impl LightTrait for InfiniteUniformLight {
    fn base(&self) -> &LightBase {
        &self.base
    }
    fn sample_li(
        &self,
        ctx: &LightSampleContext,
        u: Point2f,
        lambda: &SampledWavelengths,
        allow_incomplete_pdf: bool,
    ) -> Option<LightLiSample> {
        if allow_incomplete_pdf {
            return None;
        }
        let wi = sample_uniform_sphere(u);
        let pdf = uniform_sphere_pdf();
        let base = InteractionBase::new_boundary(
            ctx.p() + wi * (2. * self.scene_radius),
            0.,
            MediumInterface::default(),
        );
        let intr_simple = SimpleInteraction::new(base);

        let intr = Interaction::Simple(intr_simple);
        Some(LightLiSample::new(
            self.scale * self.lemit.sample(lambda),
            wi,
            pdf,
            intr,
        ))
    }

    fn pdf_li(
        &self,
        _ctx: &LightSampleContext,
        _wi: Vector3f,
        allow_incomplete_pdf: bool,
    ) -> Float {
        if allow_incomplete_pdf {
            return 0.;
        }
        uniform_sphere_pdf()
    }

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

    fn le(&self, _ray: &Ray, lambda: &SampledWavelengths) -> SampledSpectrum {
        self.scale * self.lemit.sample(lambda)
    }

    #[cfg(not(target_os = "cuda"))]
    fn preprocess(&mut self, _scene_bounds: &Bounds3f) {
        todo!()
    }

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

    #[cfg(not(target_os = "cuda"))]
    fn phi(&self, lambda: SampledWavelengths) -> SampledSpectrum {
        4. * PI * PI * square(self.scene_radius) * self.scale * self.lemit.sample(&lambda)
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct InfiniteImageLight {
    pub base: LightBase,
    pub image: Ptr<Image>,
    pub image_color_space: Ptr<RGBColorSpace>,
    pub distrib: Ptr<PiecewiseConstant2D>,
    pub compensated_distrib: Ptr<PiecewiseConstant2D>,
    pub scale: Float,
    pub scene_radius: Float,
    pub scene_center: Point3f,
}

unsafe impl Send for InfiniteImageLight {}
unsafe impl Sync for InfiniteImageLight {}

impl InfiniteImageLight {
    fn image_le(&self, uv: Point2f, lambda: &SampledWavelengths) -> SampledSpectrum {
        let mut rgb = RGB::default();
        for c in 0..3 {
            rgb[c] = self.image.lookup_nearest_channel_with_wrap(
                uv,
                c as i32,
                WrapMode::OctahedralSphere.into(),
            );
        }
        let spec = RGBIlluminantSpectrum::new(&*self.image_color_space, rgb.clamp_zero());
        self.scale * spec.sample(lambda)
    }
}

impl LightTrait for InfiniteImageLight {
    fn base(&self) -> &LightBase {
        &self.base
    }

    fn sample_li(
        &self,
        ctx: &LightSampleContext,
        u: Point2f,
        lambda: &SampledWavelengths,
        allow_incomplete_pdf: bool,
    ) -> Option<LightLiSample> {
        let (uv, map_pdf, _) = if allow_incomplete_pdf {
            self.compensated_distrib.sample(u)
        } else {
            self.distrib.sample(u)
        };

        if map_pdf == 0. {
            return None;
        }
        // Convert infinite light sample point to direction
        let w_light = equal_area_square_to_sphere(uv);
        let wi = self.base.render_from_light.apply_to_vector(w_light);
        let pdf = map_pdf / (4. * PI);

        // Return radiance value for infinite light direction
        let base = InteractionBase::new_boundary(
            ctx.p() + wi * (2. * self.scene_radius),
            0.,
            self.base.medium_interface,
        );
        let simple_intr = SimpleInteraction::new(base);
        let intr = Interaction::Simple(simple_intr);

        Some(LightLiSample::new(self.image_le(uv, lambda), wi, pdf, intr))
    }

    fn pdf_li(&self, _ctx: &LightSampleContext, wi: Vector3f, allow_incomplete_pdf: bool) -> Float {
        let w_light = self.base.render_from_light.apply_inverse_vector(wi);
        let uv = equal_area_sphere_to_square(w_light);
        let pdf = if allow_incomplete_pdf {
            self.compensated_distrib.pdf(uv)
        } else {
            self.distrib.pdf(uv)
        };
        pdf / (4. * PI)
    }

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

    fn le(&self, ray: &Ray, lambda: &SampledWavelengths) -> SampledSpectrum {
        let w_light = self
            .base
            .render_from_light
            .apply_inverse_vector(ray.d)
            .normalize();
        let uv = equal_area_sphere_to_square(w_light);
        self.image_le(uv, lambda)
    }

    #[cfg(not(target_os = "cuda"))]
    fn phi(&self, lambda: SampledWavelengths) -> SampledSpectrum {
        let mut sum_l = SampledSpectrum::new(0.);
        let width = self.image.resolution.x();
        let height = self.image.resolution.y();
        for v in 0..height {
            for u in 0..width {
                let mut rgb = RGB::default();
                for c in 0..3 {
                    rgb[c] = self.image.get_channel_with_wrap(
                        Point2i::new(u, v),
                        c as i32,
                        WrapMode::OctahedralSphere.into(),
                    );
                }
                sum_l += RGBIlluminantSpectrum::new(&*self.image_color_space, rgb.clamp_zero())
                    .sample(&lambda);
            }
        }
        4. * PI * PI * square(self.scene_radius) * self.scale * sum_l / (width * height) as Float
    }

    #[cfg(not(target_os = "cuda"))]
    fn preprocess(&mut self, scene_bounds: &Bounds3f) {
        let (scene_center, scene_radius) = scene_bounds.bounding_sphere();
        self.scene_center = scene_center;
        self.scene_radius = scene_radius;
    }

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

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct InfinitePortalLight {
    pub base: LightBase,
    pub image: Ptr<Image>,
    pub image_color_space: Ptr<RGBColorSpace>,
    pub scale: Float,
    pub portal: [Point3f; 4],
    pub portal_frame: Frame,
    pub distribution: WindowedPiecewiseConstant2D,
    pub scene_center: Point3f,
    pub scene_radius: Float,
}

impl InfinitePortalLight {
    pub fn image_lookup(&self, uv: Point2f, lambda: &SampledWavelengths) -> SampledSpectrum {
        let mut rgb = RGB::default();
        for c in 0..3 {
            rgb[c] = self.image.lookup_nearest_channel(uv, c as i32)
        }
        let spec = RGBIlluminantSpectrum::new(&*self.image_color_space, rgb.clamp_zero());
        self.scale * spec.sample(lambda)
    }

    pub fn image_from_render(&self, w_render: Vector3f) -> Option<(Point2f, Float)> {
        let w = self.portal_frame.to_local(w_render);

        if w.z() <= 0.0 {
            return None;
        }

        let alpha = w.x().atan2(w.z());
        let beta = w.y().atan2(w.z());

        let duv_dw = square(PI) * (1. - square(w.x())) * (1. - square(w.y())) / w.z();

        Some((
            Point2f::new(
                clamp((alpha + PI / 2.0) / PI, 0.0, 1.0),
                clamp((beta + PI / 2.0) / PI, 0.0, 1.0),
            ),
            duv_dw,
        ))
    }

    pub fn image_bounds(&self, p: Point3f) -> Option<Bounds2f> {
        let (p0, _) = self.image_from_render((self.portal[0] - p).normalize())?;
        let (p1, _) = self.image_from_render((self.portal[2] - p).normalize())?;

        Some(Bounds2f::from_points(p0, p1))
    }

    pub fn area(&self) -> Float {
        (self.portal[1] - self.portal[0]).norm() * (self.portal[3] - self.portal[0]).norm()
    }

    pub fn render_from_image(portal_frame: Frame, uv: Point2f) -> (Vector3f, Float) {
        let alpha = -PI / 2.0 + uv.x() * PI;
        let beta = -PI / 2.0 + uv.y() * PI;

        let x = alpha.tan();
        let y = beta.tan();

        let w = Vector3f::new(x, y, 1.0).normalize();

        let duv_dw = square(PI) * (1.0 - square(w.x())) * (1.0 - square(w.y())) / w.z();

        (portal_frame.from_local(w), duv_dw)
    }
}

impl LightTrait for InfinitePortalLight {
    fn base(&self) -> &LightBase {
        &self.base
    }

    fn sample_li(
        &self,
        ctx: &LightSampleContext,
        u: Point2f,
        lambda: &SampledWavelengths,
        _allow_incomplete_pdf: bool,
    ) -> Option<LightLiSample> {
        let b = self.image_bounds(ctx.p())?;
        let (uv, map_pdf) = self.distribution.sample(u, b)?;
        let (wi, duv_dw) = Self::render_from_image(self.portal_frame, uv);
        if duv_dw == 0. {
            return None;
        }
        let pdf = map_pdf / duv_dw;
        let l = self.image_lookup(uv, lambda);
        let pl = ctx.p() + 2. * self.scene_radius * wi;
        let base = InteractionBase::new_boundary(pl, 0., self.base.medium_interface);
        let sintr = SimpleInteraction::new(base);
        let intr = Interaction::Simple(sintr);
        Some(LightLiSample::new(l, wi, pdf, intr))
    }

    fn pdf_li(&self, ctx: &LightSampleContext, wi: Vector3f, _allow_incomplete_pdf: bool) -> Float {
        let Some((uv, duv_dw)) = self.image_from_render(wi) else {
            return 0.;
        };
        let Some(b) = self.image_bounds(ctx.p()) else {
            return 0.;
        };
        let pdf = self.distribution.pdf(uv, b);
        pdf / duv_dw
    }

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

    fn le(&self, ray: &Ray, lambda: &SampledWavelengths) -> SampledSpectrum {
        let uv = self.image_from_render(ray.d.normalize());
        let b = self.image_bounds(ray.o);
        match (uv, b) {
            (Some((p, duv_dw)), Some(bounds)) if bounds.contains(p) => self.image_lookup(p, lambda),
            _ => SampledSpectrum::new(0.0),
        }
    }

    #[cfg(not(target_os = "cuda"))]
    fn phi(&self, _lambda: SampledWavelengths) -> SampledSpectrum {
        todo!()
    }

    #[cfg(not(target_os = "cuda"))]
    fn preprocess(&mut self, scene_bounds: &Bounds3f) {
        (self.scene_center, self.scene_radius) = scene_bounds.bounding_sphere();
    }

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