mod pipeline;

use shared::core::bsdf::BSDF;
use shared::core::bssrdf::{BSSRDFTrait, SubsurfaceInteraction};
use shared::core::bxdf::{BxDFFlags, BxDFTrait, FArgs, TransportMode};
use shared::core::camera::Camera;
use shared::core::film::VisibleSurface;
use shared::core::geometry::{Bounds2i, Point2f, Point2i, Point3fi, Ray, Vector3f, VectorLike};
use shared::core::interaction::{
    Interaction, InteractionTrait, MediumInteraction, SimpleInteraction, SurfaceInteraction,
};
use shared::core::light::{Light, LightSampleContext};
use shared::core::medium::{MediumTrait, PhaseFunctionTrait};
use shared::core::options::get_options;
use shared::core::pbrt::{Float, SHADOW_EPSILON};
use shared::core::primitive::{Primitive, PrimitiveTrait};
use shared::core::sampler::{Sampler, SamplerTrait};
use shared::core::shape::ShapeIntersection;
use shared::lights::sampler::LightSamplerTrait;
use shared::lights::sampler::{LightSampler, UniformLightSampler};
use shared::spectra::{SampledSpectrum, SampledWavelengths};
use shared::utils::hash::{hash_buffer, mix_bits};
use shared::utils::math::{float_to_bits, sample_discrete, square};
use shared::utils::rng::Rng;
use shared::utils::sampling::{
    WeightedReservoirSampler, power_heuristic, sample_uniform_hemisphere, sample_uniform_sphere,
    uniform_hemisphere_pdf, uniform_sphere_pdf,
};

use std::sync::Arc;

use crate::Arena;
use crate::spectra::get_spectra_context;

#[derive(Clone, Debug)]
pub struct IntegratorBase {
    pub aggregate: Arc<Primitive>,
    pub lights: Vec<Arc<Light>>,
    pub infinite_lights: Vec<Arc<Light>>,
}

impl IntegratorBase {
    pub fn new(aggregate: Arc<Primitive>, lights: Vec<Arc<Light>>) -> Self {
        let infinite_lights = lights
            .iter()
            .filter(|light| light.light_type().is_infinite())
            .cloned()
            .collect();

        Self {
            aggregate,
            lights,
            infinite_lights,
        }
    }

    pub fn intersect(&self, ray: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        self.aggregate.intersect(ray, t_max)
    }

    pub fn intersect_p(&self, ray: &Ray, t_max: Option<Float>) -> bool {
        self.aggregate.intersect_p(ray, t_max)
    }

    pub fn unoccluded(&self, p0: &Interaction, p1: &Interaction) -> bool {
        !self.intersect_p(&p0.spawn_ray_to_interaction(p1), Some(1. - SHADOW_EPSILON))
    }
}

pub trait IntegratorTrait {
    fn render(&self);
}

pub trait RayIntegratorTrait {
    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &mut Arena,
    );

    fn li(
        &self,
        ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        visible_surface: bool,
        arena: &mut Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>);
}

struct NextRaySample {
    wi: Vector3f,
    f: SampledSpectrum,
    pdf: Float,
    is_specular: bool,
}

#[derive(Clone, Debug)]
pub struct SimplePathIntegrator {
    base: IntegratorBase,
    camera: Arc<Camera>,
    light_sampler: UniformLightSampler,
    sample_lights: bool,
    sample_bsdf: bool,
    max_depth: usize,
}

unsafe impl Send for SimplePathIntegrator {}
unsafe impl Sync for SimplePathIntegrator {}

impl SimplePathIntegrator {
    fn sample_direct_lighting(
        &self,
        isect: &SurfaceInteraction,
        bsdf: &BSDF,
        sampler: &mut Sampler,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum {
        let Some(sampled_light) = self.light_sampler.sample(sampler.get1d()) else {
            return SampledSpectrum::new(0.0);
        };
        let u_light = sampler.get2d();
        let p0 = LightSampleContext::from(isect);

        let Some(ls) = sampled_light.light.sample_li(&p0, u_light, lambda, false) else {
            return SampledSpectrum::new(0.0);
        };

        if ls.pdf == 0.0 || ls.l.is_black() {
            return SampledSpectrum::new(0.0);
        }

        let wi = ls.wi;
        let wo = -isect.wo();

        let Some(mut f) = bsdf.f(wo, wi, TransportMode::Radiance) else {
            return SampledSpectrum::new(0.);
        };

        f *= wi.abs_dot(isect.shading.n.into());

        if f.is_black() {
            return SampledSpectrum::new(0.0);
        }

        // Visibility check
        let p0_surf = Interaction::Surface(isect.clone());
        if !self.base.unoccluded(&p0_surf, ls.p_light.as_ref()) {
            return SampledSpectrum::new(0.0);
        }

        f * ls.l / (sampled_light.p * ls.pdf)
    }

    fn sample_bsdf_direction(
        &self,
        wo: Vector3f,
        bsdf: &BSDF,
        sampler: &mut Sampler,
        _isect: &SurfaceInteraction,
    ) -> Option<NextRaySample> {
        let u = sampler.get1d();
        let bs = bsdf.sample_f(wo, u, sampler.get2d(), FArgs::default())?;

        Some(NextRaySample {
            wi: bs.wi,
            f: bs.f,
            pdf: bs.pdf,
            is_specular: bs.is_specular(),
        })
    }

    fn sample_uniform_direction(
        &self,
        wo: Vector3f,
        bsdf: &BSDF,
        sampler: &mut Sampler,
        isect: &SurfaceInteraction,
    ) -> Option<NextRaySample> {
        let flags = bsdf.flags();
        let mut wi;
        let pdf;

        if BxDFFlags::is_reflective(&flags) && BxDFFlags::is_transmissive(&flags) {
            wi = sample_uniform_sphere(sampler.get2d());
            pdf = uniform_sphere_pdf();
        } else {
            wi = sample_uniform_hemisphere(sampler.get2d());
            pdf = uniform_hemisphere_pdf();

            let same_hemi = wo.dot(isect.n().into()) * wi.dot(isect.n().into()) > 0.0;
            if BxDFFlags::is_reflective(&flags) && !same_hemi
                || BxDFFlags::is_transmissive(&flags) && same_hemi
            {
                wi = -wi;
            }
        }

        let f = bsdf.f(wo, wi, TransportMode::Radiance)?;

        Some(NextRaySample {
            wi,
            f,
            pdf,
            is_specular: false,
        })
    }
}

impl RayIntegratorTrait for SimplePathIntegrator {
    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &mut Arena,
    ) {
        pipeline::evaluate_pixel_sample(
            self,
            self.camera.as_ref(),
            sampler,
            p_pixel,
            sample_ind,
            arena,
        );
    }

    fn li(
        &self,
        mut ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        _visible_surface: bool,
        arena: &mut Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>) {
        let mut l = SampledSpectrum::new(0.0);
        let mut beta = SampledSpectrum::new(1.0);
        let mut specular_bounce = true;
        let mut depth = 0;
        while !beta.is_black() {
            let Some(mut si) = self.base.intersect(&ray, None) else {
                if !self.sample_lights || specular_bounce {
                    for light in &self.base.infinite_lights {
                        l += light.le(&ray, lambda);
                    }
                }
                break;
            };

            let t_hit = si.t_hit();
            let isect = &mut si.intr;
            if !self.sample_lights || specular_bounce {
                l += beta * isect.le(-ray.d, lambda);
            }
            if depth == self.max_depth {
                break;
            }
            depth += 1;

            let Some(bsdf) = isect.get_bsdf(&ray, lambda, self.camera.as_ref(), sampler) else {
                specular_bounce = false;
                isect.skip_intersection(&mut ray, t_hit);
                continue;
            };

            let wo = -ray.d;
            if self.sample_lights {
                l += beta * self.sample_direct_lighting(isect, &bsdf, sampler, lambda);
            }

            let sample = if self.sample_bsdf {
                self.sample_bsdf_direction(wo, &bsdf, sampler, isect)
            } else {
                self.sample_uniform_direction(wo, &bsdf, sampler, isect)
            };

            let Some(bs) = sample else {
                break;
            };

            beta *= bs.f * bs.wi.abs_dot(isect.shading.n.into()) / bs.pdf;
            specular_bounce = bs.is_specular;
            ray = isect.spawn_ray(bs.wi);
        }
        let stdspec = get_spectra_context();
        assert!(beta.y(lambda, stdspec) > 0.);
        debug_assert!(beta.y(lambda, stdspec).is_finite());
        (l, None)
    }
}

#[derive(Clone, Debug)]
pub struct PathIntegrator {
    base: IntegratorBase,
    camera: Arc<Camera>,
    sampler_prototype: Sampler,
    light_sampler: LightSampler,
    regularize: bool,
    max_depth: usize,
}

unsafe impl Send for PathIntegrator {}
unsafe impl Sync for PathIntegrator {}

impl PathIntegrator {
    fn sample_ld(
        &self,
        intr: &SurfaceInteraction,
        bsdf: &BSDF,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
    ) -> Option<SampledSpectrum> {
        let mut ctx: LightSampleContext = intr.into();
        let flags = bsdf.flags();
        if BxDFFlags::is_reflective(&flags) ^ !BxDFFlags::is_transmissive(&flags) {
            ctx.pi = Point3fi::new_from_point(intr.offset_ray_vector(-intr.wo()));
        }

        let u = sampler.get1d();
        let sampled_light = self.light_sampler.sample_with_context(&ctx, u)?;

        let u_light = sampler.get2d();
        let light = sampled_light.light;
        let ls = light.sample_li(&ctx, u_light, lambda, true)?;

        if ls.l.is_black() || ls.pdf == 0.0 {
            return None;
        }

        let wo = intr.wo();
        let wi = ls.wi;
        let f = bsdf.f(wo, wi, TransportMode::Radiance)? * wi.abs_dot(intr.shading.n.into());

        if !self
            .base
            .unoccluded(&Interaction::Surface(intr.clone()), &ls.p_light)
        {
            return None;
        }

        let pl = sampled_light.p * ls.pdf;
        if light.light_type().is_delta_light() {
            Some(ls.l * f / pl)
        } else {
            let pb = bsdf.pdf(wo, wi, FArgs::default());
            let wl = power_heuristic(1, pl, 1, pb);
            Some(wl * ls.l * f / pl)
        }
    }
}

impl RayIntegratorTrait for PathIntegrator {
    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &mut Arena,
    ) {
        pipeline::evaluate_pixel_sample(
            self,
            self.camera.as_ref(),
            sampler,
            p_pixel,
            sample_ind,
            arena,
        );
    }

    fn li(
        &self,
        mut ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        visible_surface: bool,
        _arena: &mut Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>) {
        let mut l = SampledSpectrum::new(0.0);
        let mut beta = SampledSpectrum::new(1.0);
        let mut visible: Option<VisibleSurface> = None;
        let mut depth = 0;
        let mut pb = 1.;
        let mut eta_scale = 1.;
        let mut specular_bounce = false;
        let mut any_non_specular_bounces = false;

        let mut prev_int_ctx = LightSampleContext::default();
        loop {
            let Some(mut si) = self.base.intersect(&ray, None) else {
                for light in &self.base.infinite_lights {
                    let le = light.le(&ray, lambda);
                    if depth == 0 || specular_bounce {
                        l += beta * le;
                    } else {
                        let pl = self.light_sampler.pmf_with_context(&prev_int_ctx, light)
                            * light.pdf_li(&prev_int_ctx, ray.d, true);
                        let w_b = power_heuristic(1, pb, 1, pl);
                        l += beta * w_b * le;
                    }
                }
                break;
            };

            let t_hit = si.t_hit();
            let isect = &mut si.intr;
            let le = isect.le(-ray.d, lambda);
            if depth == 0 || specular_bounce {
                l += beta * le;
            } else if let Some(area_light) = &isect.area_light {
                let pl = self
                    .light_sampler
                    .pmf_with_context(&prev_int_ctx, area_light)
                    + area_light.pdf_li(&prev_int_ctx, ray.d, true);
                let wl = power_heuristic(1, pb, 1, pl);
                l += beta * wl * le;
            }

            let Some(mut bsdf) = isect.get_bsdf(&ray, lambda, self.camera.as_ref(), sampler) else {
                specular_bounce = true;
                isect.skip_intersection(&mut ray, t_hit);
                continue;
            };

            if depth == 0 && visible_surface {
                const N_RHO_SAMPLES: usize = 16;
                const UC_RHO: [Float; N_RHO_SAMPLES] = [
                    0.75741637,
                    0.37870818,
                    0.7083487,
                    0.18935409,
                    0.9149363,
                    0.35417435,
                    0.5990858,
                    0.09467703,
                    0.8578725,
                    0.45746812,
                    0.686759,
                    0.17708716,
                    0.9674518,
                    0.2995429,
                    0.5083201,
                    0.047338516,
                ];
                let u_rho: [Point2f; N_RHO_SAMPLES] = [
                    Point2f::new(0.855985, 0.570367),
                    Point2f::new(0.381823, 0.851844),
                    Point2f::new(0.285328, 0.764262),
                    Point2f::new(0.733380, 0.114073),
                    Point2f::new(0.542663, 0.344465),
                    Point2f::new(0.127274, 0.414848),
                    Point2f::new(0.964700, 0.947162),
                    Point2f::new(0.594089, 0.643463),
                    Point2f::new(0.095109, 0.170369),
                    Point2f::new(0.825444, 0.263359),
                    Point2f::new(0.429467, 0.454469),
                    Point2f::new(0.244460, 0.816459),
                    Point2f::new(0.756135, 0.731258),
                    Point2f::new(0.516165, 0.152852),
                    Point2f::new(0.180888, 0.214174),
                    Point2f::new(0.898579, 0.503897),
                ];

                let albedo = bsdf.rho_wo(isect.wo(), &UC_RHO, &u_rho);
                visible = Some(VisibleSurface::new(isect, &albedo, lambda));
            }

            if self.regularize && any_non_specular_bounces {
                bsdf.regularize();
            }

            if depth == self.max_depth {
                break;
            }

            depth += 1;

            if bsdf.flags().is_non_specular()
                && let Some(ld) = self.sample_ld(isect, &bsdf, lambda, sampler)
            {
                l += beta * ld;
            }

            let wo = -ray.d;
            let u = sampler.get1d();
            let Some(bs) = bsdf.sample_f(wo, u, sampler.get2d(), FArgs::default()) else {
                break;
            };

            beta *= bs.f * bs.wi.abs_dot(isect.shading.n.into()) / bs.pdf;
            pb = if bs.pdf_is_proportional {
                bsdf.pdf(wo, bs.wi, FArgs::default())
            } else {
                bs.pdf
            };
            specular_bounce = bs.is_specular();
            any_non_specular_bounces |= !bs.is_specular();
            if bs.is_transmissive() {
                eta_scale *= square(bs.eta);
            }
            prev_int_ctx = LightSampleContext::from(&si.intr);
            ray = si
                .intr
                .spawn_ray_with_differentials(&ray, bs.wi, bs.flags, bs.eta);

            let rr_beta = beta * eta_scale;
            if rr_beta.max_component_value() < 1. && depth > 1 {
                let q = (1. - rr_beta.max_component_value()).max(0.);
                if sampler.get1d() < q {
                    break;
                }
                beta /= 1. - q;
            }
        }

        (l, visible)
    }
}

#[derive(Clone, Debug)]
pub struct SimpleVolPathIntegrator {
    base: IntegratorBase,
    max_depth: usize,
    camera: Arc<Camera>,
    sampler_prototype: Sampler,
}

impl RayIntegratorTrait for SimpleVolPathIntegrator {
    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &mut Arena,
    ) {
        pipeline::evaluate_pixel_sample(
            self,
            self.camera.as_ref(),
            sampler,
            p_pixel,
            sample_ind,
            arena,
        );
    }

    fn li(
        &self,
        mut ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        _visible_surface: bool,
        _arena: &mut Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>) {
        let mut l = SampledSpectrum::new(0.);
        let mut beta = 1.;
        let mut depth = 0;

        let lambda = lambda.terminate_secondary();

        loop {
            let si = self.base.intersect(&ray, None);
            let mut scattered = false;
            let mut terminated = false;
            let current_medium = ray.medium.take();
            if let Some(medium) = &current_medium {
                let hash0 = hash_buffer(&sampler.get1d(), 0);
                let hash1 = hash_buffer(&sampler.get1d(), 0);
                let mut rng = Rng::new_with_offset(hash0, hash1);
                let t_max = if let Some(hit) = &si {
                    hit.t_hit
                } else {
                    Float::INFINITY
                };
                let u = sampler.get1d();
                let mut u_mode = sampler.get1d();
                medium.as_ref().sample_t_maj(
                    ray.clone(),
                    t_max,
                    u,
                    &mut rng,
                    &lambda,
                    |p, mp, sigma_maj, _t_maj, rng| -> bool {
                        let sig_maj = sigma_maj[0];
                        // Compute medium event probabilities for interaction
                        let p_absorb: Float = mp.sigma_a[0] / sig_maj;
                        let p_scatter: Float = mp.sigma_s[0] / sig_maj;
                        let p_null = (1. - p_absorb - p_scatter).max(0_f32);
                        let mode =
                            sample_discrete(&[p_absorb, p_scatter, p_null], u_mode, None, None);
                        if mode == 0 {
                            // Handle absorption event for medium sample
                            l += beta * mp.le;
                            terminated = true;
                            false
                        } else if mode == 1 {
                            // Handle regular scattering event for medium sample
                            if depth >= self.max_depth {
                                return false;
                            }
                            depth += 1;

                            let u_phase = Point2f::new(rng.uniform(), rng.uniform());

                            let Some(ps) = mp.phase.sample_p(-ray.d, u_phase) else {
                                terminated = true;
                                return false;
                            };

                            beta *= ps.p / ps.pdf;
                            ray.o = p;
                            ray.d = ps.wi;
                            scattered = true;
                            false
                        } else {
                            // Handle null-scattering event for medium sample
                            u_mode = rng.uniform();
                            true
                        }
                    },
                );
            }
            // Handle terminated and unscattered rays after medium sampling
            if terminated {
                return (l, None);
            }

            if scattered {
                return (l, None);
            }

            let Some(mut hit) = si else {
                for light in &self.base.infinite_lights {
                    l += beta * light.le(&ray, &lambda);
                }
                return (l, None);
            };

            // Handle surface intersection along ray path
            if let Some(bsdf) = &hit.intr.get_bsdf(&ray, &lambda, &self.camera, sampler) {
                let uc = sampler.get1d();
                let u = sampler.get2d();
                if bsdf.sample_f(-ray.d, uc, u, FArgs::default()).is_some() {
                    panic!("SimpleVolPathIntegrator does not support surface scattering");
                }
            } else {
                hit.intr.skip_intersection(&mut ray, hit.t_hit);
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct VolPathIntegrator {
    base: IntegratorBase,
    sampler_prototype: Sampler,
    camera: Arc<Camera>,
    light_sampler: LightSampler,
    regularize: bool,
    max_depth: usize,
}

impl VolPathIntegrator {
    fn sample_ld(
        &self,
        intr: &Interaction,
        bsdf: Option<&BSDF>,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        beta: SampledSpectrum,
        r_p: SampledSpectrum,
    ) -> SampledSpectrum {
        let mut ctx = LightSampleContext::from(intr);
        if let Some(bsdf) = bsdf {
            let flags = bsdf.flags();
            let is_refl = flags.is_reflective();
            let is_trans = flags.is_transmissive();
            if is_refl ^ is_trans {
                let vector = if is_refl { intr.wo() } else { -intr.wo() };
                ctx.pi = Point3fi::new_from_point(intr.offset_ray_vector(vector));
            }
        }
        let u = sampler.get1d();
        let Some(sampled_light) = self.light_sampler.sample_with_context(&ctx, u) else {
            return SampledSpectrum::new(0.);
        };
        let u_light = sampler.get2d();
        let light = sampled_light.light;
        debug_assert!(sampled_light.p != 0.);
        let Some(ls) = light
            .sample_li(&ctx, u_light, lambda, true)
            .filter(|ls| !ls.l.is_black() && ls.pdf > 0.0)
        else {
            return SampledSpectrum::new(0.0);
        };
        let p_l = sampled_light.p * ls.pdf;
        let wo = intr.wo();
        let wi = ls.wi;
        let (f_hat, scatter_pdf) = match (bsdf, intr) {
            (Some(bsdf), Interaction::Surface(si)) => {
                let f = bsdf
                    .f(wo, wi, TransportMode::Radiance)
                    .map(|f_val| f_val * wi.abs_dot(si.shading.n.into()))
                    .unwrap_or(SampledSpectrum::new(0.0));

                let pdf = bsdf.pdf(wo, wi, FArgs::default());
                (f, pdf)
            }

            (None, Interaction::Medium(mi)) => {
                let phase = &mi.phase;
                let p = phase.p(wo, wi);
                let pdf = phase.pdf(wo, wi);

                (SampledSpectrum::new(p), pdf)
            }

            (Some(_), _) => panic!("BSDF provided for a non-Surface interaction"),
            (None, _) => panic!("No BSDF provided for a non-Medium interaction"),
        };

        if f_hat.is_black() {
            return SampledSpectrum::new(0.);
        }

        let mut light_ray = intr.spawn_ray_to_interaction(ls.p_light.as_ref());
        let mut t_ray = SampledSpectrum::new(1.);
        let mut r_l = SampledSpectrum::new(1.);
        let mut r_u = SampledSpectrum::new(1.);
        let hash0 = hash_buffer(&[light_ray.o.x(), light_ray.o.y(), light_ray.o.z()], 0);
        let hash1 = hash_buffer(&[light_ray.d.x(), light_ray.d.y(), light_ray.d.z()], 0);
        let mut rng = Rng::new_with_offset(hash0, hash1);

        while light_ray.d != Vector3f::zero() {
            // Trace ray through media to estimate transmittance
            let si = &self.base.intersect(&light_ray, Some(1. - SHADOW_EPSILON));
            if si.as_ref().filter(|s| s.intr.material.is_some()).is_some() {
                return SampledSpectrum::new(0.);
            }

            let current_medium = light_ray.medium.take();
            if let Some(medium) = &current_medium {
                let t_max = if let Some(s) = si {
                    s.t_hit()
                } else {
                    1. - SHADOW_EPSILON
                };

                let u: Float = rng.uniform();
                let t_maj = medium.as_ref().sample_t_maj(
                    light_ray.clone(),
                    t_max,
                    u,
                    &mut rng,
                    lambda,
                    |_p, mp, sigma_maj, t_maj, rng| {
                        let sigma_n =
                            SampledSpectrum::clamp_zero(&(sigma_maj - mp.sigma_a - mp.sigma_s));
                        let pdf = t_maj[0] * sigma_maj[0];
                        t_ray *= t_maj * sigma_n / pdf;
                        r_l *= t_maj * sigma_maj / pdf;
                        r_u *= t_maj * sigma_n / pdf;

                        let tr = t_ray / (r_l + r_u).average();
                        if tr.max_component_value() < 0.05 {
                            let q = 0.75;
                            if rng.uniform::<Float>() < 1. {
                                t_ray = SampledSpectrum::new(0.);
                            } else {
                                t_ray /= 1. - q;
                            }
                        }

                        if t_ray.is_black() {
                            return false;
                        }
                        true
                    },
                );
                t_ray *= t_maj / t_maj[0];
                r_l *= t_maj / t_maj[0];
                r_u *= t_maj / t_maj[0];
            }

            if t_ray.is_black() {
                return SampledSpectrum::new(0.);
            }

            if let Some(s) = si {
                light_ray = s.intr.spawn_ray_to_interaction(ls.p_light.as_ref());
            } else {
                break;
            }
        }

        r_l *= r_p * p_l;
        r_u *= r_p * scatter_pdf;

        if light.light_type().is_delta_light() {
            beta * f_hat * t_ray * ls.l / r_l.average()
        } else {
            beta * f_hat * t_ray * ls.l / (r_l + r_u).average()
        }
    }
}

impl RayIntegratorTrait for VolPathIntegrator {
    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &mut Arena,
    ) {
        pipeline::evaluate_pixel_sample(
            self,
            self.camera.as_ref(),
            sampler,
            p_pixel,
            sample_ind,
            arena,
        );
    }

    fn li(
        &self,
        mut ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        visible_surface: bool,
        arena: &mut Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>) {
        let mut l = SampledSpectrum::new(0.);
        let mut beta = SampledSpectrum::new(1.);
        let mut r_u = SampledSpectrum::new(1.);
        let mut r_l = SampledSpectrum::new(1.);
        let mut specular_bounce = false;
        let mut any_non_specular_bounces = false;
        let mut depth = 0;
        let mut eta_scale = 1.;
        let mut prev_int_ctx = LightSampleContext::default();
        let mut visible: Option<VisibleSurface> = None;

        loop {
            let si = self.base.intersect(&ray, None);
            let current_medium = ray.medium.take();
            if let Some(medium) = &current_medium {
                let mut scattered = false;
                let mut terminated = false;
                let hash0 = hash_buffer(&sampler.get1d(), 0);
                let hash1 = hash_buffer(&sampler.get1d(), 0);
                let mut rng = Rng::new_with_offset(hash0, hash1);

                let t_max = if let Some(hit) = &si {
                    hit.t_hit
                } else {
                    Float::INFINITY
                };
                let t_maj = medium.as_ref().sample_t_maj(
                    ray.clone(),
                    t_max,
                    sampler.get1d(),
                    &mut rng,
                    lambda,
                    |p, mp, sigma_maj, t_maj, rng| -> bool {
                        // Handle medium scattering event for ray
                        if beta.is_black() {
                            terminated = true;
                            return false;
                        }

                        if depth < self.max_depth && !mp.le.is_black() {
                            let pdf = sigma_maj[0] * t_maj[0];
                            let betap = beta * t_maj / pdf;
                            let r_e = r_u * sigma_maj * t_maj / pdf;
                            if !r_e.is_black() {
                                l += betap * mp.sigma_a * mp.le / r_e.average();
                            }
                        }
                        let p_absorb = mp.sigma_a[0] / sigma_maj[0];
                        let p_scatter = mp.sigma_s[0] / sigma_maj[0];
                        let p_null = (1. - p_absorb - p_scatter).max(0.);
                        assert!(1. - p_absorb - p_scatter > -1e-6);
                        // Sample medium scattering event type and update path
                        let um = rng.uniform();
                        let mode = sample_discrete(&[p_absorb, p_scatter, p_null], um, None, None);
                        if mode == 0 {
                            terminated = true;
                            false
                        } else if mode == 1 {
                            // Handle scattering along ray path
                            // Stop path sampling if maximum depth has been reached
                            if depth >= self.max_depth {
                                terminated = true;
                                return false;
                            }
                            depth += 1;

                            let pdf = t_maj[0] * mp.sigma_s[0];
                            beta *= t_maj * mp.sigma_s / pdf;
                            r_u *= t_maj * mp.sigma_s / pdf;

                            if !beta.is_black() && !r_u.is_black() {
                                let ray_medium = ray.medium.expect("Void scattering").clone();
                                let intr = MediumInteraction::new(
                                    p, -ray.d, ray.time, ray_medium, mp.phase,
                                );
                                l += self.sample_ld(
                                    &Interaction::Medium(intr.clone()),
                                    None,
                                    lambda,
                                    sampler,
                                    beta,
                                    r_u,
                                );
                                let u = sampler.get2d();
                                let Some(ps) =
                                    intr.phase.sample_p(-ray.d, u).filter(|ps| ps.pdf > 0.0)
                                else {
                                    terminated = true;
                                    // Stop traversal (Path died)
                                    return false;
                                };

                                beta *= ps.p / ps.pdf;
                                r_l = r_u / ps.pdf;
                                prev_int_ctx = LightSampleContext::from(&intr);
                                scattered = true;
                                ray.o = p;
                                ray.d = ps.wi;
                                specular_bounce = false;
                                any_non_specular_bounces = false;
                            }
                            false
                        } else {
                            let sigma_n =
                                SampledSpectrum::clamp_zero(&(sigma_maj - mp.sigma_a - mp.sigma_s));
                            let pdf = t_maj[0] * sigma_n[0];
                            beta *= t_maj * sigma_n / pdf;
                            if pdf == 0. {
                                beta = SampledSpectrum::new(0.);
                            }
                            r_u *= t_maj * sigma_n / pdf;
                            r_l *= t_maj * sigma_maj / pdf;
                            !beta.is_black() && !r_u.is_black()
                        }
                    },
                );

                if terminated || beta.is_black() || r_u.is_black() {
                    return (l, None);
                }
                if scattered {
                    continue;
                }
                beta *= t_maj / t_maj[0];
                r_u *= t_maj / t_maj[0];
                r_l *= t_maj / t_maj[0];
            }

            // Handle surviving unscattered rays
            // Add emitted light at volume path vertex or from the environment

            let stdspec = get_spectra_context();
            let Some(mut hit) = si else {
                for light in &self.base.infinite_lights {
                    let le = light.le(&ray, lambda);
                    if !le.is_black() {
                        if depth == 0 || specular_bounce {
                            l += beta * light.le(&ray, lambda);
                        }
                    } else {
                        // Add infinite light contribution using both PDFs with MIS
                        let p_l = self.light_sampler.pmf_with_context(&prev_int_ctx, light)
                            * light.pdf_li(&prev_int_ctx, ray.d, true);
                        r_l *= p_l;
                        l += beta * le / (r_u + r_l).average();
                    }
                }
                break;
            };

            let le = &hit.intr.le(-ray.d, lambda);
            if !le.is_black() {
                if depth == 0 || specular_bounce {
                    l += beta * *le / r_u.average();
                }
            } else {
                // Add surface light contribution using both PDFs with MIS
                let area_light = hit.intr.area_light.as_ref().expect("No surface light");
                let p_l = self
                    .light_sampler
                    .pmf_with_context(&prev_int_ctx, area_light);
                r_l *= p_l;
                l += beta * *le / (r_u + r_l).average();
            }

            let isect = &mut hit.intr;
            let Some(mut bsdf) = isect.get_bsdf(&ray, lambda, &self.camera, sampler) else {
                hit.intr.skip_intersection(&mut ray, hit.t_hit());
                continue;
            };

            if depth == 0 && visible_surface {
                const N_RHO_SAMPLES: usize = 16;
                const UC_RHO: [Float; N_RHO_SAMPLES] = [
                    0.75741637,
                    0.37870818,
                    0.7083487,
                    0.18935409,
                    0.9149363,
                    0.35417435,
                    0.5990858,
                    0.09467703,
                    0.8578725,
                    0.45746812,
                    0.686759,
                    0.17708716,
                    0.9674518,
                    0.2995429,
                    0.5083201,
                    0.047338516,
                ];
                let u_rho: [Point2f; N_RHO_SAMPLES] = [
                    Point2f::new(0.855985, 0.570367),
                    Point2f::new(0.381823, 0.851844),
                    Point2f::new(0.285328, 0.764262),
                    Point2f::new(0.733380, 0.114073),
                    Point2f::new(0.542663, 0.344465),
                    Point2f::new(0.127274, 0.414848),
                    Point2f::new(0.964700, 0.947162),
                    Point2f::new(0.594089, 0.643463),
                    Point2f::new(0.095109, 0.170369),
                    Point2f::new(0.825444, 0.263359),
                    Point2f::new(0.429467, 0.454469),
                    Point2f::new(0.244460, 0.816459),
                    Point2f::new(0.756135, 0.731258),
                    Point2f::new(0.516165, 0.152852),
                    Point2f::new(0.180888, 0.214174),
                    Point2f::new(0.898579, 0.503897),
                ];
                let albedo = bsdf.rho_wo(isect.wo(), &UC_RHO, &u_rho);
                visible = Some(VisibleSurface::new(isect, &albedo, lambda));
            }

            if depth >= self.max_depth {
                return (l, visible);
            }

            depth += 1;

            if self.regularize && any_non_specular_bounces {
                bsdf.regularize();
            }

            if BxDFFlags::is_non_specular(&bsdf.flags()) {
                l += self.sample_ld(
                    &Interaction::Surface(isect.clone()),
                    Some(&bsdf),
                    lambda,
                    sampler,
                    beta,
                    r_u,
                );
                debug_assert!(l.y(lambda, stdspec).is_finite());
            }
            let wo = isect.wo();
            let n = isect.shading.n;
            prev_int_ctx = LightSampleContext::from(&*isect);
            let u = sampler.get1d();
            let Some(bs) = bsdf.sample_f(wo, u, sampler.get2d(), FArgs::default()) else {
                break;
            };
            beta *= bs.f * bs.wi.abs_dot(n.into()) / bs.pdf;
            if bs.pdf_is_proportional {
                r_l = r_u / bsdf.pdf(wo, bs.wi, FArgs::default());
            } else {
                r_l = r_u / bs.pdf;
            }
            debug_assert!(beta.y(lambda, stdspec).is_finite());
            // Update volumetric integrator path state after surface scattering
            specular_bounce = bs.is_specular();
            if bs.is_transmissive() {
                eta_scale *= square(bs.eta);
            }
            ray = isect.spawn_ray_with_differentials(&ray, bs.wi, bs.flags, bs.eta);

            if let Some(bssrdf) = (*isect).get_bssrdf(&ray, lambda, self.camera.as_ref())
                && bs.is_transmissive()
            {
                let uc = sampler.get1d();
                let up = sampler.get2d();
                let Some(probe_seg) = bssrdf.sample_sp(uc, up) else {
                    break;
                };
                let seed = mix_bits(float_to_bits(sampler.get1d()).into());
                let mut interaction_sampler =
                    WeightedReservoirSampler::<SubsurfaceInteraction>::new(seed);
                let base = SimpleInteraction::new(Point3fi::new_from_point(probe_seg.p0), ray.time);
                loop {
                    let r = base.spawn_ray_to_point(probe_seg.p1);
                    if r.d == Vector3f::zero() {
                        break;
                    }
                    let Some(si) = self.base.intersect(&r, Some(1.)) else {
                        break;
                    };

                    let materials_match = match (&si.intr.material, &isect.material) {
                        (Some(m1), Some(m2)) => Arc::ptr_eq(m1, m2),
                        _ => false,
                    };

                    if materials_match {
                        interaction_sampler.add(SubsurfaceInteraction::new(&si.intr), 1.);
                    }
                }

                if !interaction_sampler.has_sample() {
                    break;
                }

                if let Some(ssi) = interaction_sampler.get_sample() {
                    let bssrdf_sample = bssrdf.probe_intersection_to_sample(ssi);
                    if bssrdf_sample.sp.is_black() || bssrdf_sample.pdf.is_black() {
                        break;
                    }
                    let pdf = interaction_sampler.sample_probability() * bssrdf_sample.pdf[0];
                    beta *= bssrdf_sample.sp / pdf;
                    r_u *= bssrdf_sample.pdf / bssrdf_sample.pdf[0];
                    let mut pi = SurfaceInteraction::from(ssi);
                    pi.common.wo = bssrdf_sample.wo;
                    prev_int_ctx = LightSampleContext::from(&pi);
                    let mut sw = bssrdf_sample.sw;
                    if self.regularize {
                        sw.regularize();
                    }

                    l += self.sample_ld(
                        &Interaction::Surface(pi.clone()),
                        Some(&sw),
                        lambda,
                        sampler,
                        beta,
                        r_u,
                    );

                    let u = sampler.get1d();
                    let Some(bs) = sw.sample_f(pi.wo(), u, sampler.get2d(), FArgs::default())
                    else {
                        break;
                    };
                    beta *= bs.f * bs.wi.abs_dot(pi.shading.n.into()) / bs.pdf;
                    r_l = r_u / bs.pdf;
                    debug_assert!(beta.y(lambda).is_finite());
                    specular_bounce = bs.is_specular();
                    ray = pi.spawn_ray(bs.wi);
                }
            }

            if beta.is_black() {
                break;
            }

            let rr_beta = beta * eta_scale / r_u.average();
            let u_rr = sampler.get1d();
            if rr_beta.max_component_value() < 1. && depth > 1 {
                let q = (1. - rr_beta.max_component_value()).max(0.);
                if u_rr < 1. {
                    break;
                }
                beta /= 1. - q;
            }
        }

        (l, visible)
    }
}

pub enum Integrator {
    BPDT(BPDTIntegrator),
    MLT(MLTIntegrator),
    SPPM(SPPMIntegrator),
    Sampler(SamplerIntegrator),
}

pub struct BPDTIntegrator;
pub struct MLTIntegrator;
pub struct SPPMIntegrator;
pub struct SamplerIntegrator;