mod bvh;
mod camera;
mod medium;
mod photons;
mod scene;
mod scenes;
mod vec3;

use std::f64::consts::PI;
use std::time::Instant;

use indicatif::{ProgressBar, ProgressStyle};
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use rayon::prelude::*;

use camera::Camera;
use medium::{hg_phase, sample_hg, MediumState};
use photons::CausticMaps;
use scene::{spot_emission, Light, Material, Ray, Scene, EPSILON};
use vec3::Vec3;

/// Shadow rays stop this fraction short of the light, so they don't hit the
/// light's own surface. Relative, because scenes range from ~1 to ~1000 units.
const SHADOW_RELATIVE_EPSILON: f64 = 1e-4;

/// Where light sampled by next event estimation arrives.
enum Receiver {
    /// A diffuse surface: Lambertian BRDF times the cosine at the surface.
    Surface { normal: Vec3, albedo: Vec3 },
    /// A scattering point in a medium: the phase function for turning from
    /// the path's direction `dir` towards the light.
    Medium { dir: Vec3, g: f64 },
}

/// Direct light at `point` from one uniformly sampled point on `light`'s
/// surface, attenuated by whatever media the shadow ray passes through.
fn sample_direct<R: Rng>(
    scene: &Scene,
    state: MediumState,
    light: &Light,
    point: Vec3,
    receiver: Receiver,
    rng: &mut R,
) -> Vec3 {
    let (lp, light_normal) = light.sample(rng);

    let to_light = lp - point;
    let dist_sq = to_light.dot(to_light);
    if dist_sq < 1e-12 {
        return Vec3::ZERO;
    }
    let dist = dist_sq.sqrt();
    let d_to_light = to_light / dist;
    let cos_at_light = -light_normal.dot(d_to_light);
    let (response, cos_at_rec, origin) = match receiver {
        Receiver::Surface { normal, albedo } => (albedo / PI, normal.dot(d_to_light), point + EPSILON * normal),
        Receiver::Medium { dir, g } => (Vec3::splat(hg_phase(dir.dot(d_to_light), g)), 1.0, point),
    };
    if cos_at_rec <= 0.0 || cos_at_light <= 0.0 {
        return Vec3::ZERO;
    }

    let tr = scene.transmittance(state, origin, d_to_light, dist * (1.0 - SHADOW_RELATIVE_EPSILON));
    if tr == Vec3::ZERO {
        return Vec3::ZERO;
    }

    // Uniform area sampling has pdf 1 / area.
    let geometry = cos_at_rec * cos_at_light / dist_sq;
    let emission = spot_emission(light.emission, light.spot, cos_at_light);
    (emission.mul_elem(response) * (geometry * light.area())).mul_elem(tr)
}

/// Bounces before Russian roulette may terminate a path.
const RR_MIN_DEPTH: u32 = 3;
/// Guard against paths that keep crossing invisible volume boundaries.
const MAX_VOLUME_CROSSINGS: u32 = 256;

/// Russian roulette: after RR_MIN_DEPTH bounces, end dim paths early,
/// boosting survivors so the estimate stays unbiased. Glass/mirror paths
/// (throughput ~1) survive. Returns false if the path ends.
fn russian_roulette<R: Rng>(throughput: &mut Vec3, bounces: u32, rng: &mut R) -> bool {
    if bounces >= RR_MIN_DEPTH {
        let survive = throughput.x.max(throughput.y).max(throughput.z).min(0.95);
        if rng.gen::<f64>() >= survive {
            return false;
        }
        *throughput = *throughput / survive;
    }
    true
}

/// With `caustics`, caustic light (light that went through glass or off
/// mirrors before reaching fog or a diffuse surface) comes from the photon
/// maps, and path tracing skips the matching paths (from a diffuse or medium
/// interaction through mirrors/glass to a light):
/// - normally (all lights emit photons), only at the path's first such
///   interaction: the maps are gathered along fog segments and at the first
///   diffuse surface until then, and later caustics are path traced;
/// - with `direct_in_media` (only spotlights emit photons), for spotlights
///   at every interaction: the volume map also holds their direct light, so
///   points in fog don't sample spotlights, and every fog segment and
///   diffuse surface of the path gathers the maps. Other lights are path
///   traced as usual.
fn trace<R: Rng>(scene: &Scene, primary: &Ray, max_depth: u32, rng: &mut R, caustics: Option<&CausticMaps>) -> Vec3 {
    let mut radiance = Vec3::ZERO;
    // Product of the attenuations along the path so far (divided by the RR
    // survival probabilities), i.e. how much of any light found here reaches the camera.
    let mut throughput = Vec3::ONE;
    // False after a diffuse bounce or medium scattering: that light was
    // already gathered by next event estimation, so hitting it again would
    // double count.
    let mut count_emission = true;
    let mut ray = *primary;
    let mut state = MediumState::Outside;
    let mut bounces = 0;
    let mut crossings = 0;
    // Diffuse and medium interactions so far.
    let mut non_specular = 0;
    let direct_in_media = caustics.is_some_and(|m| m.direct_in_media);
    // Whether this interaction gathers from the photon maps.
    let maps_here = |non_specular: u32| caustics.filter(|_| direct_in_media || non_specular == 0);

    while bounces < max_depth {
        let hit = scene.hit(&ray, EPSILON, f64::INFINITY);

        // Travel through the medium (if any) up to the surface: either
        // scatter somewhere along the way, or arrive with the absorption
        // along the way applied.
        let t_hit = hit.as_ref().map_or(f64::INFINITY, |h| h.t);
        if let Some((medium, t0, t1)) = scene.medium_segment(state, &ray, t_hit) {
            if let Some(maps) = maps_here(non_specular) {
                radiance += throughput.mul_elem(maps.beam(&medium, &ray, t0, t1));
            }
            let t = t0 + medium.sample_distance(rng);
            if t < t1 {
                throughput = throughput.mul_elem(medium.absorption(t - t0));
                let p = ray.at(t);
                // Spotlights are left out when the photon map carries their
                // direct light into the fog.
                if let Some((light, pick_prob)) = scene.pick_light(p, direct_in_media, rng) {
                    let receiver = Receiver::Medium { dir: ray.direction, g: medium.g };
                    let direct = sample_direct(scene, state, light, p, receiver, rng) / pick_prob;
                    radiance += throughput.mul_elem(direct);
                }
                count_emission = false;
                non_specular += 1;
                ray = Ray::from_unit(p, sample_hg(ray.direction, medium.g, rng));
                bounces += 1;
                if !russian_roulette(&mut throughput, bounces, rng) {
                    break;
                }
                continue;
            }
            throughput = throughput.mul_elem(medium.absorption(t1 - t0));
        }

        let Some(rec) = hit else {
            radiance += throughput.mul_elem(scene.background.radiance(ray.direction));
            break;
        };

        match rec.material {
            Material::Light { emission, spot } => {
                // Lights emit from their front side only (matters for one-sided rects).
                // Reached from a diffuse/medium interaction via mirrors or
                // glass: a caustic, which the photon maps may already cover.
                let from_maps = match caustics {
                    Some(_) if non_specular == 0 => false,
                    Some(m) if m.direct_in_media => spot.is_some(),
                    Some(_) => non_specular == 1,
                    None => false,
                };
                if count_emission && rec.front_face && !from_maps {
                    let emission = spot_emission(emission, spot, -ray.direction.dot(rec.normal));
                    radiance += throughput.mul_elem(emission);
                }
                break;
            }
            Material::Volume { medium } => {
                // Invisible boundary: carry on in the same direction, now
                // inside (or back outside) the volume. Not a bounce.
                state = if rec.front_face { MediumState::Inside(Some(medium)) } else { MediumState::Outside };
                ray = Ray::from_unit(rec.point, ray.direction);
                crossings += 1;
                if crossings > MAX_VOLUME_CROSSINGS {
                    break;
                }
                continue;
            }
            _ => {}
        }

        let Some((attenuation, scattered)) = rec.material.scatter(&ray, &rec, rng) else {
            break;
        };

        match rec.material {
            Material::Diffuse { albedo } => {
                if let Some(maps) = maps_here(non_specular) {
                    radiance += throughput.mul_elem(maps.surface(rec.point, rec.normal, albedo));
                }
                if let Some((light, pick_prob)) = scene.pick_light(rec.point, false, rng) {
                    let receiver = Receiver::Surface { normal: rec.normal, albedo };
                    let direct = sample_direct(scene, state, light, rec.point, receiver, rng) / pick_prob;
                    radiance += throughput.mul_elem(direct);
                }
                count_emission = false;
                non_specular += 1;
            }
            Material::Dielectric { medium, .. } => {
                // Refracted through the surface (rather than reflected):
                // entering or leaving the object's interior medium.
                if scattered.direction.dot(rec.normal) < 0.0 {
                    state = if rec.front_face { MediumState::Inside(medium) } else { MediumState::Outside };
                }
                count_emission = true;
            }
            _ => count_emission = true,
        }

        throughput = throughput.mul_elem(attenuation);
        ray = scattered;
        bounces += 1;
        if !russian_roulette(&mut throughput, bounces, rng) {
            break;
        }
    }

    radiance
}

struct Settings {
    width: usize,
    height: usize,
    /// Maximum samples per pixel (adaptive sampling may stop earlier).
    spp: u32,
    max_depth: u32,
    output: String,
    /// Target standard error of a pixel in display (gamma) space, 0..1.
    /// 0 disables adaptive sampling.
    noise_threshold: f64,
}

/// Adaptive sampling: every pixel first takes this many samples...
const ADAPTIVE_MIN_SPP: u32 = 32;
/// ...then later passes grow each still-noisy pixel's count by this fraction.
const ADAPTIVE_GROWTH: f64 = 0.5;

/// Running mean and variance of a pixel's luminance (Welford's algorithm).
#[derive(Default)]
struct PixelStats {
    n: u32,
    mean: f64,
    m2: f64,
}

impl PixelStats {
    fn add(&mut self, x: f64) {
        self.n += 1;
        let delta = x - self.mean;
        self.mean += delta / self.n as f64;
        self.m2 += delta * (x - self.mean);
    }

    /// Standard error of the pixel mean after gamma (sqrt) encoding, using
    /// d(sqrt(x))/dx = 1 / (2 sqrt(x)) to map it from linear space.
    fn display_error(&self) -> f64 {
        if self.n < 2 {
            return f64::INFINITY;
        }
        let variance = self.m2 / (self.n - 1) as f64;
        let std_error = (variance / self.n as f64).sqrt();
        std_error / (2.0 * self.mean.max(1e-4).sqrt())
    }
}

struct Pixel {
    sum: Vec3,
    stats: PixelStats,
    rng: SmallRng,
}

struct Renderer<'a> {
    scene: &'a Scene,
    cam: Camera,
    s: &'a Settings,
    caustics: Option<&'a CausticMaps>,
}

impl Renderer<'_> {
    fn sample(&self, px: &mut Pixel, i: usize, j: usize, count: u32) {
        let (width, height) = (self.s.width, self.s.height);
        for _ in 0..count {
            let u = (i as f64 + px.rng.gen::<f64>()) / width as f64;
            let v = ((height - 1 - j) as f64 + px.rng.gen::<f64>()) / height as f64;
            let sample = trace(self.scene, &self.cam.get_ray(u, v), self.s.max_depth, &mut px.rng, self.caustics);
            px.sum += sample;
            px.stats.add(sample.luminance());
        }
    }

    /// Per-pixel noise, taken as the maximum over each 3x3 neighbourhood. A
    /// pixel's own estimate is unreliable while it has only seen a few samples
    /// (a pixel that has not yet hit a rare bright path looks converged), so
    /// a neighbour that did keeps it sampling.
    fn neighbourhood_error(&self, pixels: &[Pixel]) -> Vec<f64> {
        let (width, height) = (self.s.width, self.s.height);
        let own: Vec<f64> = pixels.iter().map(|p| p.stats.display_error()).collect();
        (0..width * height)
            .into_par_iter()
            .map(|idx| {
                let (i, j) = (idx % width, idx / width);
                let mut worst = 0.0_f64;
                for nj in j.saturating_sub(1)..(j + 2).min(height) {
                    for ni in i.saturating_sub(1)..(i + 2).min(width) {
                        worst = worst.max(own[nj * width + ni]);
                    }
                }
                worst
            })
            .collect()
    }
}

fn bar(len: u64, template: &str) -> ProgressBar {
    let pb = ProgressBar::new(len);
    pb.set_style(
        ProgressStyle::with_template(template)
            .unwrap()
            .progress_chars("=> "),
    );
    pb
}

/// Returns the gamma-encoded image and the total number of samples taken.
fn render(scene: &Scene, s: &Settings, caustics: Option<&CausticMaps>) -> (Vec<Vec3>, u64) {
    let (width, height) = (s.width, s.height);
    let view = scene.view;
    let r = Renderer {
        scene,
        cam: Camera::new(
            view.lookfrom,
            view.lookat,
            Vec3::new(0.0, 1.0, 0.0),
            view.vfov,
            width as f64 / height as f64,
            (view.lookfrom - view.lookat).length(),
        ),
        s,
        caustics,
    };

    let mut pixels: Vec<Pixel> = (0..width * height)
        .map(|idx| Pixel {
            sum: Vec3::ZERO,
            stats: PixelStats::default(),
            rng: SmallRng::seed_from_u64(idx as u64 ^ 0x9E37_79B9_7F4A_7C15),
        })
        .collect();

    // Base pass: every pixel gets the full budget, or the adaptive minimum.
    let adaptive = s.noise_threshold > 0.0;
    let base_spp = if adaptive { s.spp.min(ADAPTIVE_MIN_SPP) } else { s.spp };
    let pb = bar(
        height as u64,
        "{spinner:.green} Rendering [{bar:50.bright_green/blue}] {percent:>3}% {pos}/{len} rows  {elapsed_precise}  ETA {eta_precise}",
    );
    pixels.par_chunks_mut(width).enumerate().for_each(|(j, row)| {
        for (i, px) in row.iter_mut().enumerate() {
            r.sample(px, i, j, base_spp);
        }
        pb.inc(1);
    });
    pb.finish();

    // Adaptive passes: keep sampling pixels whose neighbourhood is still noisy.
    if adaptive {
        let pb = bar(
            (width * height) as u64,
            "{spinner:.green} Refining  [{bar:50.bright_green/blue}] {percent:>3}% {pos}/{len} pixels converged  {msg}",
        );
        for pass in 1.. {
            let error = r.neighbourhood_error(&pixels);
            let active: Vec<bool> = pixels
                .iter()
                .zip(&error)
                .map(|(p, &e)| p.stats.n < s.spp && e >= s.noise_threshold)
                .collect();
            let remaining = active.iter().filter(|&&a| a).count();
            pb.set_position((width * height - remaining) as u64);
            pb.set_message(format!("pass {pass}"));
            if remaining == 0 {
                break;
            }
            pixels
                .par_chunks_mut(width)
                .zip(active.par_chunks(width))
                .enumerate()
                .for_each(|(j, (row, row_active))| {
                    for (i, (px, &a)) in row.iter_mut().zip(row_active).enumerate() {
                        if a {
                            let n = px.stats.n;
                            let extra = ((n as f64 * ADAPTIVE_GROWTH).ceil() as u32).clamp(1, s.spp - n);
                            r.sample(px, i, j, extra);
                        }
                    }
                });
        }
        pb.finish();
    }

    let total_samples = pixels.iter().map(|p| p.stats.n as u64).sum();
    let image = pixels
        .iter()
        .map(|p| (p.sum / p.stats.n as f64).map(|c| c.max(0.0).sqrt().min(1.0)))
        .collect();
    (image, total_samples)
}

fn save_png(image: &[Vec3], width: usize, height: usize, path: &str) -> image::ImageResult<()> {
    let bytes: Vec<u8> = image
        .iter()
        .flat_map(|c| [c.x, c.y, c.z])
        .map(|c| (c.clamp(0.0, 1.0) * 255.0) as u8)
        .collect();
    image::save_buffer(path, &bytes, width as u32, height as u32, image::ColorType::Rgb8)
}

fn region_mean(image: &[Vec3], width: usize, rows: std::ops::Range<usize>, cols: std::ops::Range<usize>) -> f64 {
    let n = rows.len() * cols.len();
    if n == 0 {
        return 0.0;
    }
    let sum: f64 = rows
        .flat_map(|j| cols.clone().map(move |i| image[j * width + i].mean()))
        .sum();
    sum / n as f64
}

fn print_region_table(image: &[Vec3], width: usize, height: usize) {
    let (h_start, h_end) = (height / 3, 2 * height / 3);
    let (w_start, w_end) = (width / 3, 2 * width / 3);
    let regions = [
        ("Left Wall", 0..height, 0..width / 4),
        ("Right Wall", 0..height, 3 * width / 4..width),
        ("Floor", h_end..height, w_start..w_end),
        ("Ceiling", 0..h_start, w_start..w_end),
        ("Center", h_start..h_end, w_start..w_end),
    ];

    println!("Region Brightness Analysis");
    println!("{:<12} {:>10}  Visual", "Region", "Brightness");
    for (name, rows, cols) in regions {
        let brightness = region_mean(image, width, rows, cols) * 255.0;
        let filled = ((brightness / 255.0 * 20.0) as usize).min(20);
        let bar = format!("{}{}", "#".repeat(filled), "-".repeat(20 - filled));
        let color = if brightness > 80.0 { 32 } else if brightness > 40.0 { 33 } else { 31 };
        println!("{name:<12} {brightness:>10.1}  \x1b[{color}m{bar}\x1b[0m");
    }
}

const USAGE: &str = "usage: pathtracer [--scene NAME] [--list-scenes] [--photons N | --no-photons] \
[width] [height] [samples_per_pixel] [max_bounces] [output] [noise_threshold]";

fn list_scenes() {
    println!("Scenes:");
    for s in scenes::SCENES {
        println!("  {:<10} {}", s.name, s.description);
    }
}

fn main() {
    // Pull out the --scene / --list-scenes options; the rest are positional.
    let mut scene_name = String::from("cornell");
    // Override for the number of caustic photons (0: path trace only).
    let mut photons_override: Option<usize> = None;
    let mut args: Vec<String> = Vec::new();
    let mut raw = std::env::args().skip(1);
    while let Some(a) = raw.next() {
        if a == "--list-scenes" {
            list_scenes();
            return;
        } else if a == "--scene" {
            scene_name = raw.next().unwrap_or_else(|| {
                eprintln!("{USAGE}");
                std::process::exit(2);
            });
        } else if let Some(name) = a.strip_prefix("--scene=") {
            scene_name = name.to_string();
        } else if a == "--photons" {
            photons_override = raw.next().and_then(|n| n.parse().ok());
            if photons_override.is_none() {
                eprintln!("{USAGE}");
                std::process::exit(2);
            }
        } else if a == "--no-photons" {
            photons_override = Some(0);
        } else if a == "--help" || a == "-h" {
            println!("{USAGE}");
            list_scenes();
            return;
        } else {
            args.push(a);
        }
    }
    let Some(scene_info) = scenes::find(&scene_name) else {
        eprintln!("unknown scene '{scene_name}'");
        list_scenes();
        std::process::exit(2);
    };

    let arg = |i: usize, default: &str| args.get(i).cloned().unwrap_or_else(|| default.to_string());
    fn parse<T: std::str::FromStr>(s: String) -> T {
        s.parse().unwrap_or_else(|_| {
            eprintln!("{USAGE}");
            std::process::exit(2);
        })
    }

    let settings = Settings {
        width: parse(arg(0, "400")),
        height: parse(arg(1, "300")),
        spp: parse(arg(2, "50")),
        max_depth: parse(arg(3, "15")),
        output: arg(4, &format!("{scene_name}.png")),
        noise_threshold: parse(arg(5, "0")),
    };
    if settings.width == 0 || settings.height == 0 || settings.spp == 0 || settings.noise_threshold < 0.0 {
        eprintln!("width, height and samples_per_pixel must be >= 1, noise_threshold >= 0");
        std::process::exit(2);
    }

    let scene = (scene_info.build)();

    println!("=== Path Tracer (Rust) ===");
    println!("Features: reflections, refraction/caustics, multiple bounces, NEE, BVH");
    println!();
    println!("Scene:             {} ({} objects, {} lights)", scene_name, scene.object_count(), scene.light_count());
    println!("Resolution:        {}x{}", settings.width, settings.height);
    println!("Samples per pixel: {} (max)", settings.spp);
    println!("Max bounces:       {} (Russian roulette after {})", settings.max_depth, RR_MIN_DEPTH);
    if settings.noise_threshold > 0.0 {
        println!("Adaptive sampling: stop at noise < {}", settings.noise_threshold);
    } else {
        println!("Adaptive sampling: off");
    }
    println!("Threads:           {}", rayon::current_num_threads());
    println!(
        "BVH:               8-wide, depth {}, {}",
        scene.bvh_depth(),
        if scene.uses_avx512() { "AVX-512" } else { "scalar" }
    );
    println!("Output:            {}", settings.output);
    println!();

    let start = Instant::now();
    let photon_settings = scene.caustic_photons.map(|mut p| {
        if let Some(n) = photons_override {
            p.emitted = n;
        }
        p
    });
    let caustics = match photon_settings {
        Some(p) if p.emitted > 0 => {
            let t = Instant::now();
            let maps = CausticMaps::build(&scene, p);
            let (volume, surface) = maps.counts();
            println!(
                "Caustic photons:   {} emitted, {} stored in media, {} on surfaces ({:.2}s)",
                p.emitted,
                volume,
                surface,
                t.elapsed().as_secs_f64()
            );
            Some(maps)
        }
        _ => None,
    };
    let (image, total_samples) = render(&scene, &settings, caustics.as_ref());
    let elapsed = start.elapsed().as_secs_f64();
    let pixels = (settings.width * settings.height) as f64;

    if let Err(e) = save_png(&image, settings.width, settings.height, &settings.output) {
        eprintln!("failed to write {}: {e}", settings.output);
        std::process::exit(1);
    }

    println!();
    println!("Rendered {}", settings.output);
    println!(
        "Total time: {:.2}s | Speed: {:.2} MPix*samp/s | Avg samples/pixel: {:.1}",
        elapsed,
        total_samples as f64 / elapsed / 1e6,
        total_samples as f64 / pixels
    );
    // The regions (left wall, floor, ...) only mean something for the Cornell box.
    if scene_name == "cornell" {
        println!();
        print_region_table(&image, settings.width, settings.height);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use medium::Medium;
    use scene::{Background, SceneBuilder, View};

    fn view() -> View {
        View { lookfrom: Vec3::new(0.0, 0.0, -10.0), lookat: Vec3::ZERO, vfov: 40.0 }
    }

    /// Average of `n` path samples along one ray.
    fn mean_radiance(scene: &Scene, ray: Ray, n: usize) -> Vec3 {
        let mut rng = SmallRng::seed_from_u64(1);
        let mut sum = Vec3::ZERO;
        for _ in 0..n {
            sum += trace(scene, &ray, 10_000, &mut rng, None);
        }
        sum / n as f64
    }

    fn assert_close(v: Vec3, expected: Vec3, tol: f64) {
        for a in 0..3 {
            assert!((v[a] - expected[a]).abs() < tol, "got {v:?}, expected {expected:?}");
        }
    }

    const INTO_CENTER: fn() -> Ray = || Ray::from_unit(Vec3::new(0.0, 0.0, -10.0), Vec3::new(0.0, 0.0, 1.0));

    /// White furnace: with no absorption and a uniform background of 1,
    /// every path eventually leaves and sees 1, so the mean must be 1 however
    /// much the medium scatters.
    #[test]
    fn furnace_fog() {
        let mut b = SceneBuilder::new(view(), Background::Color(Vec3::ONE));
        b.fog(Vec3::splat(-1.0), Vec3::splat(1.0), Medium::new(Vec3::ZERO, 3.0, 0.6));
        assert_close(mean_radiance(&b.build(), INTO_CENTER(), 200_000), Vec3::ONE, 0.01);
    }

    #[test]
    fn furnace_volume_sphere() {
        let mut b = SceneBuilder::new(view(), Background::Color(Vec3::ONE));
        let smoke = Medium::new(Vec3::ZERO, 4.0, -0.3);
        b.sphere(Vec3::ZERO, 1.0, Material::Volume { medium: smoke });
        assert_close(mean_radiance(&b.build(), INTO_CENTER(), 200_000), Vec3::ONE, 0.01);
    }

    #[test]
    fn furnace_glass_with_scattering_interior() {
        let mut b = SceneBuilder::new(view(), Background::Color(Vec3::ONE));
        let milk = Medium::new(Vec3::ZERO, 5.0, 0.2);
        b.sphere(Vec3::ZERO, 1.0, Material::Dielectric { ref_idx: 1.5, albedo: Vec3::ONE, medium: Some(milk) });
        assert_close(mean_radiance(&b.build(), INTO_CENTER(), 200_000), Vec3::ONE, 0.01);
    }

    /// Henyey-Greenstein: sampled directions have mean cosine g, and the
    /// phase function (used by next event estimation) integrates to 1.
    #[test]
    fn henyey_greenstein() {
        let mut rng = SmallRng::seed_from_u64(2);
        let dir = Vec3::new(0.3, -0.5, 0.8).normalize();
        for g in [-0.7, 0.0, 0.4, 0.9] {
            let n = 400_000;
            let mean_cos: f64 = (0..n).map(|_| sample_hg(dir, g, &mut rng).dot(dir)).sum::<f64>() / n as f64;
            assert!((mean_cos - g).abs() < 0.005, "g = {g}: mean cos {mean_cos}");

            // Midpoint rule over cos_theta (the phase function doesn't depend on phi).
            let steps = 200_000;
            let integral: f64 = (0..steps)
                .map(|i| {
                    let c = -1.0 + (i as f64 + 0.5) * 2.0 / steps as f64;
                    hg_phase(c, g) * 2.0 * PI * (2.0 / steps as f64)
                })
                .sum();
            assert!((integral - 1.0).abs() < 1e-3, "g = {g}: integral {integral}");
        }
    }

    /// A purely absorbing slab of thickness 2: Beer's law, exactly, per channel.
    #[test]
    fn beer_lambert_absorption() {
        let mut b = SceneBuilder::new(view(), Background::Color(Vec3::ONE));
        let sigma_a = Vec3::new(0.1, 0.5, 1.5);
        b.fog(Vec3::splat(-1.0), Vec3::splat(1.0), Medium::new(sigma_a, 0.0, 0.0));
        let expected = (sigma_a * -2.0).map(f64::exp);
        assert_close(mean_radiance(&b.build(), INTO_CENTER(), 10), expected, 1e-12);
    }
}