forge-ibl-bake 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0
// forge-ibl-bake · deveraux.dev
//! The three IBL bake passes and their orchestrator.
//!
//! `bake_all` is the primary entry point: validates params, runs specular prefilter
//! + irradiance (or SH9) + BRDF LUT, and writes KTX2 to `out_dir`. Individual
//! passes are public for testing. All sampling uses Hammersley → fully deterministic.
//! DESIGN_TARGET: white-furnace and prefilter-monotonic tests gate correctness.

use std::f32::consts::PI;
use std::path::Path;

macro_rules! par_dispatch {
    ($range:expr, map $f:expr) => {{
        #[cfg(feature = "parallel")]
        { use rayon::prelude::*; $range.into_par_iter().map($f).collect() }
        #[cfg(not(feature = "parallel"))]
        { $range.map($f).collect() }
    }};
    ($range:expr, flat_map $f:expr) => {{
        #[cfg(feature = "parallel")]
        { use rayon::prelude::*; $range.into_par_iter().map($f).flatten().collect() }
        #[cfg(not(feature = "parallel"))]
        { $range.map($f).flatten().collect() }
    }};
}

use crate::{
    error::BakeError,
    hdr::HdrImage,
    ktx2::{
        write_ktx2_2d, write_ktx2_cubemap, VK_FORMAT_R16G16B16A16_SFLOAT,
        VK_FORMAT_R16G16_SFLOAT,
    },
    sample::{
        cube_face_dir, f32_to_f16_bits, hammersley, importance_sample_ggx,
        integrate_brdf, normalize3, orthonormal_basis, sample_equirect, tbn_to_world, dot3,
    },
    sh,
};

/// Parameters controlling all three bake passes.
#[derive(Clone, Debug)]
pub struct BakeParams {
    /// Cubemap face size (pixels) for prefiltered specular.
    pub specular_size: u32,
    /// Cubemap face size (pixels) for diffuse irradiance.
    pub irradiance_size: u32,
    /// BRDF LUT texture dimension (square).
    pub lut_size: u32,
    /// Monte-Carlo samples per texel for all passes.
    pub samples: u32,
    /// If true, emit SH9 JSON instead of irradiance cubemap.
    pub emit_sh: bool,
    /// Reserved; v0 bake is deterministic via Hammersley (no RNG).
    pub seed: u64,
}

impl Default for BakeParams {
    fn default() -> Self {
        Self {
            specular_size: 256,
            irradiance_size: 32,
            lut_size: 128,
            samples: 1024,
            emit_sh: false,
            seed: 0,
        }
    }
}

/// Full IBL bake: specular cubemap + irradiance (or SH9) + BRDF LUT → `out_dir`.
pub fn bake_all(env: &HdrImage, params: &BakeParams, out_dir: &Path) -> Result<(), BakeError> {
    if params.specular_size == 0 || params.irradiance_size == 0 || params.lut_size == 0 {
        return Err(BakeError::InvalidSize("specular/irradiance/lut sizes must be > 0".into()));
    }
    if params.samples == 0 {
        return Err(BakeError::InvalidSize("samples must be > 0".into()));
    }
    if !params.specular_size.is_power_of_two()
        || !params.irradiance_size.is_power_of_two()
        || !params.lut_size.is_power_of_two()
    {
        return Err(BakeError::InvalidSize(
            "specular_size, irradiance_size, and lut_size must be powers of two".into(),
        ));
    }

    std::fs::create_dir_all(out_dir)?;

    // 1. Prefiltered specular env (cubemap mip-chain).
    let spec_mips = bake_specular_env(env, params);
    write_ktx2_cubemap(
        &out_dir.join("ibl_specular.ktx2"),
        params.specular_size,
        VK_FORMAT_R16G16B16A16_SFLOAT,
        &spec_mips,
    )?;

    // 2. Diffuse irradiance or SH9.
    if params.emit_sh {
        let coeffs = sh::project_sh9(env, params.samples.max(64));
        sh::write_sh9_json(&out_dir.join("ibl_sh9.json"), &coeffs)?;
    } else {
        let irr = bake_irradiance(env, params);
        write_ktx2_cubemap(
            &out_dir.join("ibl_irradiance.ktx2"),
            params.irradiance_size,
            VK_FORMAT_R16G16B16A16_SFLOAT,
            &[irr],
        )?;
    }

    // 3. Split-sum BRDF LUT.
    let lut = bake_brdf_lut(params);
    write_ktx2_2d(
        &out_dir.join("ibl_brdf_lut.ktx2"),
        params.lut_size,
        params.lut_size,
        VK_FORMAT_R16G16_SFLOAT,
        &lut,
    )?;

    Ok(())
}

// ── Pass 1: prefiltered specular cubemap mip-chain ───────────────────────────

/// Returns one `Vec<u8>` per mip level (mip[0] = finest, roughness≈0).
/// Each Vec contains 6 face images × face_size² × RGBA16F (8 bytes/pixel).
pub fn bake_specular_env(env: &HdrImage, params: &BakeParams) -> Vec<Vec<u8>> {
    let face_size = params.specular_size;
    let mip_count = mip_levels_for(face_size);

    let bake_mip = |mip: u32| -> Vec<u8> {
        let roughness = mip as f32 / (mip_count - 1).max(1) as f32;
        let dim = (face_size >> mip).max(1);
        let s = if roughness < 0.01 { 1u32 } else { params.samples };
        let mut data = Vec::with_capacity(6 * (dim * dim) as usize * 8);
        for face in 0..6u32 {
            for y in 0..dim {
                for x in 0..dim {
                    let u = (2.0 * x as f32 + 1.0) / dim as f32 - 1.0;
                    let v = (2.0 * y as f32 + 1.0) / dim as f32 - 1.0;
                    let n = cube_face_dir(face, u, v);
                    let color = prefilter_sample(env, n, roughness, s);
                    push_rgba16f(&mut data, color);
                }
            }
        }
        data
    };

    par_dispatch!(0..mip_count, map bake_mip)
}

/// GGX importance-sampled prefilter for one texel (split-sum, N=V=R).
fn prefilter_sample(env: &HdrImage, n: [f32; 3], roughness: f32, samples: u32) -> [f32; 3] {
    let v = n; // split-sum approximation: view direction = normal
    let (t, b, _) = orthonormal_basis(n);
    let mut acc = [0.0f32; 3];
    let mut total_w = 0.0f32;

    for i in 0..samples {
        let xi = hammersley(i, samples);
        let h_local = importance_sample_ggx(xi, roughness);
        let h = normalize3(tbn_to_world(h_local, t, b, n));
        let v_dot_h = dot3(v, h).max(0.0);
        let l = normalize3([
            2.0 * v_dot_h * h[0] - v[0],
            2.0 * v_dot_h * h[1] - v[1],
            2.0 * v_dot_h * h[2] - v[2],
        ]);
        let n_dot_l = dot3(n, l).max(0.0);
        if n_dot_l > 0.0 {
            let s = sample_equirect(&env.pixels, env.width, env.height, l);
            acc[0] += s[0] * n_dot_l;
            acc[1] += s[1] * n_dot_l;
            acc[2] += s[2] * n_dot_l;
            total_w += n_dot_l;
        }
    }

    if total_w < 1e-6 {
        sample_equirect(&env.pixels, env.width, env.height, n)
    } else {
        [acc[0] / total_w, acc[1] / total_w, acc[2] / total_w]
    }
}

// ── Pass 2: diffuse irradiance cubemap ───────────────────────────────────────

/// Returns 6 face images × irradiance_size² × RGBA16F in one flat Vec.
pub fn bake_irradiance(env: &HdrImage, params: &BakeParams) -> Vec<u8> {
    let dim = params.irradiance_size;
    let s = params.samples;
    let total_texels = 6u32 * dim * dim;

    let bake_texel = |idx: u32| -> [u8; 8] {
        let face = idx / (dim * dim);
        let local = idx % (dim * dim);
        let y = local / dim;
        let x = local % dim;
        let u = (2.0 * x as f32 + 1.0) / dim as f32 - 1.0;
        let v = (2.0 * y as f32 + 1.0) / dim as f32 - 1.0;
        let n = cube_face_dir(face, u, v);
        rgba16f_bytes(irradiance_sample(env, n, s))
    };

    par_dispatch!(0..total_texels, flat_map bake_texel)
}

/// Cosine-weighted hemisphere integration for one surface normal `n`.
/// pdf(ω) = cos(θ)/π → estimator = π/N × Σ L(ωᵢ) for cosine-weighted samples.
fn irradiance_sample(env: &HdrImage, n: [f32; 3], samples: u32) -> [f32; 3] {
    if samples == 0 {
        return [0.0, 0.0, 0.0];
    }
    let (t, b, _) = orthonormal_basis(n);
    let mut acc = [0.0f32; 3];

    for i in 0..samples {
        let xi = hammersley(i, samples);
        // Cosine-weighted hemisphere: r=√(xi[1]), φ=2π·xi[0], z=√(1−xi[1]).
        let r = xi[1].max(0.0).sqrt();
        let phi = 2.0 * PI * xi[0];
        let local = [r * phi.cos(), r * phi.sin(), (1.0 - xi[1]).max(0.0).sqrt()];
        let dir = normalize3(tbn_to_world(local, t, b, n));
        let s = sample_equirect(&env.pixels, env.width, env.height, dir);
        acc[0] += s[0];
        acc[1] += s[1];
        acc[2] += s[2];
    }

    // Estimator: irradiance E ≈ π/N × Σ L, PRE-DIVIDED by π so consumers can do
    // `diffuse = albedo * sample` (glTF/Khronos/Bevy/LearnOpenGL convention).
    // White-furnace ⇒ ≈1.0, not π. Locked by the magnitude assertion in proof t2.
    let w = 1.0 / samples as f32;
    [acc[0] * w, acc[1] * w, acc[2] * w]
}

// ── Pass 3: split-sum BRDF LUT ───────────────────────────────────────────────

/// Returns lut_size² × R16G16 pixel data (scale, bias) in row-major order.
/// Axes: x = NdotV (left=0, right=1), y = roughness (bottom=0, top=1).
pub fn bake_brdf_lut(params: &BakeParams) -> Vec<u8> {
    let dim = params.lut_size;
    let mut data = Vec::with_capacity((dim * dim * 4) as usize);
    for y in 0..dim {
        let roughness = (y as f32 + 0.5) / dim as f32;
        for x in 0..dim {
            let n_dot_v = (x as f32 + 0.5) / dim as f32;
            let (scale, bias) = integrate_brdf(n_dot_v, roughness, params.samples);
            data.extend_from_slice(&f32_to_f16_bits(scale).to_le_bytes());
            data.extend_from_slice(&f32_to_f16_bits(bias).to_le_bytes());
        }
    }
    data
}

// ── Helpers ───────────────────────────────────────────────────────────────────

/// Number of mip levels for a complete chain down to 1×1.
pub fn mip_levels_for(size: u32) -> u32 {
    if size == 0 {
        return 0;
    }
    32 - size.leading_zeros()
}

fn push_rgba16f(buf: &mut Vec<u8>, rgb: [f32; 3]) {
    buf.extend_from_slice(&rgba16f_bytes(rgb));
}

fn rgba16f_bytes(rgb: [f32; 3]) -> [u8; 8] {
    let r = f32_to_f16_bits(rgb[0]).to_le_bytes();
    let g = f32_to_f16_bits(rgb[1]).to_le_bytes();
    let b = f32_to_f16_bits(rgb[2]).to_le_bytes();
    let a = f32_to_f16_bits(1.0).to_le_bytes();
    [r[0], r[1], g[0], g[1], b[0], b[1], a[0], a[1]]
}