codecraft 0.2.0

//! Microfacet multiple-scattering energy compensation (Turquin 2019 albedo scaling):
//! single-scattering directional albedo `E_ss(μ)` tables, computed on the CPU and uploaded.

use bytemuck::{Pod, Zeroable};

use super::OpenPbrSurface;

/// Number of μ = cosθ bins per table.
pub const BINS: usize = 32;

/// GPU layout: eight tables of `BINS` floats each (see `energy_tables` in WGSL).
/// New tables go on the end: the prebuilt rust-gpu SPIR-V carries its own copy of these offsets.
#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
pub struct EnergyTables {
    /// Conductor with F = 1.
    pub metal: [f32; BINS],
    /// Dielectric R + T from outside (μ > 0) / inside (μ < 0).
    pub dielectric_outside: [f32; BINS],
    pub dielectric_inside: [f32; BINS],
    pub dielectric_r_outside: [f32; BINS],
    pub dielectric_r_inside: [f32; BINS],
    pub dielectric_t_outside: [f32; BINS],
    pub dielectric_t_inside: [f32; BINS],
    /// Coat reflection albedo at `coat_ior`; tabulated so the shader need not estimate it per vertex (that flickered under DLSS).
    pub coat_r: [f32; BINS],
}

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> f32 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        ((self.0 >> 40) as f32) / (1u64 << 24) as f32
    }
}

fn ggx_lambda(w: [f32; 3], a: f32) -> f32 {
    if w[2].abs() < 1e-7 {
        return 0.0;
    }
    (-1.0 + (1.0 + (a * a * (w[0] * w[0] + w[1] * w[1])) / (w[2] * w[2])).sqrt()) / 2.0
}

fn ggx_sample(wi: [f32; 3], a: f32, rng: &mut Rng) -> [f32; 3] {
    let (u1, u2) = (rng.next(), rng.next());
    let mut v = [wi[0] * a, wi[1] * a, wi[2]];
    let l = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
    v = [v[0] / l, v[1] / l, v[2] / l];
    let phi = 2.0 * std::f32::consts::PI * u1;
    let z = (1.0 - u2) * (1.0 + v[2]) - v[2];
    let st = (1.0 - z * z).clamp(0.0, 1.0).sqrt();
    let mut h = [st * phi.cos() + v[0], st * phi.sin() + v[1], z + v[2]];
    h = [h[0] * a, h[1] * a, h[2]];
    let l = (h[0] * h[0] + h[1] * h[1] + h[2] * h[2]).sqrt();
    [h[0] / l, h[1] / l, h[2] / l]
}

fn dot(a: [f32; 3], b: [f32; 3]) -> f32 {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

fn fresnel_dielectric(mui: f32, eta_ti: f32) -> f32 {
    let c = mui;
    let mut2 = eta_ti * eta_ti + c * c - 1.0;
    if mut2 <= 0.0 {
        return 1.0;
    }
    let g = mut2.sqrt();
    0.5 * ((g - c) / (g + c)).powi(2) * (1.0 + (((g + c) * c - 1.0) / ((g - c) * c + 1.0)).powi(2))
}

fn metal_albedo(mu: f32, a: f32, rng: &mut Rng, n: usize) -> f32 {
    let wi = [(1.0 - mu * mu).max(0.0).sqrt(), 0.0, mu];
    let g1 = 1.0 / (1.0 + ggx_lambda(wi, a));
    let mut acc = 0.0;
    for _ in 0..n {
        let m = ggx_sample(wi, a, rng);
        let d = dot(wi, m);
        let wo = [
            -wi[0] + 2.0 * d * m[0],
            -wi[1] + 2.0 * d * m[1],
            -wi[2] + 2.0 * d * m[2],
        ];
        if wo[2] > 0.0 {
            let g2 = 1.0 / (1.0 + ggx_lambda(wi, a) + ggx_lambda(wo, a));
            acc += g2 / g1;
        }
    }
    acc / n as f32
}

/// Single-scatter (R, T) albedos at incidence cosine `mu` (`mu < 0`: from inside); `eta_ie` = interior / exterior IOR.
fn dielectric_albedo(mu: f32, a: f32, eta_ie: f32, rng: &mut Rng, n: usize) -> (f32, f32) {
    let inside = mu < 0.0;
    let mu_abs = mu.abs();
    let wi = [(1.0 - mu_abs * mu_abs).max(0.0).sqrt(), 0.0, mu];
    let wi_up = [wi[0], wi[1], mu_abs];
    let g1 = 1.0 / (1.0 + ggx_lambda(wi, a));
    let eta_ti_refl = if inside { 1.0 / eta_ie } else { eta_ie };
    let eta_ti_photon = 1.0 / eta_ti_refl;
    let mut acc_r = 0.0;
    let mut acc_t = 0.0;
    for _ in 0..n {
        let mut m = ggx_sample(wi_up, a, rng);
        if inside {
            m[2] = -m[2];
        }
        let d = dot(wi, m);
        let wo = [
            -wi[0] + 2.0 * d * m[0],
            -wi[1] + 2.0 * d * m[1],
            -wi[2] + 2.0 * d * m[2],
        ];
        if wi[2] * wo[2] > 0.0 {
            let g2 = 1.0 / (1.0 + ggx_lambda(wi, a) + ggx_lambda(wo, a));
            acc_r += fresnel_dielectric(d.abs(), eta_ti_refl) * g2 / g1;
        }
        let mut m = ggx_sample(wi_up, a, rng);
        if inside {
            m[2] = -m[2];
        }
        let wtn = dot(wi, m);
        let disc = 1.0 - eta_ti_photon * eta_ti_photon * (1.0 - wtn * wtn);
        if disc >= 0.0 {
            let s = wtn.signum() * (eta_ti_photon * wtn.abs() - disc.sqrt());
            let bi = [
                eta_ti_photon * wi[0] - m[0] * s,
                eta_ti_photon * wi[1] - m[1] * s,
                eta_ti_photon * wi[2] - m[2] * s,
            ];
            let l = (bi[0] * bi[0] + bi[1] * bi[1] + bi[2] * bi[2]).sqrt();
            let wo = [-bi[0] / l, -bi[1] / l, -bi[2] / l];
            if wo[2] * wi[2] < 0.0 {
                let g2 = 1.0 / (1.0 + ggx_lambda(wi, a) + ggx_lambda(wo, a));
                let t = (1.0 - fresnel_dielectric(wtn.abs(), eta_ti_refl)).clamp(0.0, 1.0);
                acc_t += t * g2 / g1;
            }
        }
    }
    (acc_r / n as f32, acc_t / n as f32)
}

/// The dielectric IOR ratio the shader's `eta_s()` uses, including coat blend and `specular_weight` modulation.
pub fn eta_s(s: &OpenPbrSurface) -> f32 {
    let n_s = s.specular_ior;
    let mut eta_sc = n_s / s.coat_ior;
    if eta_sc < 1.0 {
        eta_sc = 1.0 / eta_sc;
    }
    let eta = (1.0 - s.coat_weight) * n_s + s.coat_weight * eta_sc;
    if s.specular_weight == 1.0 || (eta - 1.0).abs() < 1e-7 {
        return eta;
    }
    let f0 = (eta - 1.0).powi(2) / (1.0 + eta).powi(2);
    let eps = (eta - 1.0).signum() * (s.specular_weight * f0).clamp(0.0, 0.99999).sqrt();
    (1.0 + eps) / (1.0 - eps).max(1e-10)
}

/// Isotropic-equivalent GGX alpha of the coat NDF (`coat_ndf_roughnesses` in `coat_brdf.wgsl`).
pub fn coat_alpha(s: &OpenPbrSurface) -> f32 {
    let an = s.coat_roughness_anisotropy;
    let at = s.coat_roughness * s.coat_roughness * (2.0 / (1.0 + (1.0 - an).powi(2))).sqrt();
    let ab = (1.0 - an) * at;
    ((at * at + ab * ab) / 2.0).sqrt().max(1e-4)
}

/// Isotropic-equivalent GGX alpha of the base specular NDF, including coat roughening.
pub fn base_alpha(s: &OpenPbrSurface) -> f32 {
    let r = s.specular_roughness;
    let rc = s.coat_roughness;
    let roughened = (r.powi(4) + 2.0 * rc.powi(4)).min(1.0).powf(0.25);
    let r_eff = (1.0 - s.coat_weight) * r + s.coat_weight * roughened;
    let an = s.specular_roughness_anisotropy;
    let at = r_eff * r_eff * (2.0 / (1.0 + (1.0 - an).powi(2))).sqrt();
    let ab = (1.0 - an) * at;
    ((at * at + ab * ab) / 2.0).sqrt().max(1e-4)
}

impl EnergyTables {
    pub fn compute(s: &OpenPbrSurface) -> Self {
        let a = base_alpha(s);
        let eta = eta_s(s);
        let mut rng = Rng(0x9E37_79B9_7F4A_7C15);
        let n = 2048;
        let mut t = Self {
            metal: [1.0; BINS],
            dielectric_outside: [1.0; BINS],
            dielectric_inside: [1.0; BINS],
            dielectric_r_outside: [0.0; BINS],
            dielectric_r_inside: [0.0; BINS],
            dielectric_t_outside: [0.0; BINS],
            dielectric_t_inside: [0.0; BINS],
            coat_r: [0.0; BINS],
        };
        for i in 0..BINS {
            let mu = (i as f32 + 0.5) / BINS as f32;
            t.metal[i] = metal_albedo(mu, a, &mut rng, n).clamp(0.05, 1.0);
            let (r_out, t_out) = dielectric_albedo(mu, a, eta, &mut rng, n);
            let (r_in, t_in) = dielectric_albedo(-mu, a, eta, &mut rng, n);
            t.dielectric_outside[i] = (r_out + t_out).clamp(0.05, 1.0);
            t.dielectric_inside[i] = (r_in + t_in).clamp(0.05, 1.0);
            t.dielectric_r_outside[i] = r_out;
            t.dielectric_r_inside[i] = r_in;
            t.dielectric_t_outside[i] = t_out;
            t.dielectric_t_inside[i] = t_in;
        }
        // Separate RNG so adding the coat left the tables above bit-identical.
        if s.coat_weight > 0.0 && (s.coat_ior - 1.0).abs() > 1e-5 {
            let ca = coat_alpha(s);
            let mut coat_rng = Rng(0xD1B5_4A32_D192_ED03);
            for i in 0..BINS {
                let mu = (i as f32 + 0.5) / BINS as f32;
                // Reflection only: transmission through the coat is handled by the layering weights.
                t.coat_r[i] = dielectric_albedo(mu, ca, s.coat_ior, &mut coat_rng, n)
                    .0
                    .clamp(0.0, 1.0);
            }
        }
        t
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn smooth_surfaces_conserve_energy() {
        let s = OpenPbrSurface {
            specular_roughness: 0.0,
            ..Default::default()
        };
        let t = EnergyTables::compute(&s);
        for i in 0..BINS {
            assert!(t.metal[i] > 0.99, "metal {i}: {}", t.metal[i]);
            assert!(
                t.dielectric_outside[i] > 0.98,
                "outside {i}: {}",
                t.dielectric_outside[i]
            );
        }
    }

    #[test]
    fn rough_surfaces_lose_energy() {
        let s = OpenPbrSurface {
            specular_roughness: 0.9,
            ..Default::default()
        };
        let t = EnergyTables::compute(&s);
        // Brute-force check: E_ss ≈ 0.42 at normal incidence for α = 0.81.
        let e = t.metal[BINS - 1];
        assert!((0.35..0.5).contains(&e), "metal normal incidence {e}");
        assert!(t.dielectric_inside[0] < t.dielectric_outside[BINS - 1]);
    }
}

#[cfg(test)]
mod print_tables {
    use super::*;
    #[test]
    #[ignore]
    fn print_rough_glass_tables() {
        let s = OpenPbrSurface {
            transmission_weight: 1.0,
            specular_roughness: 0.5,
            ..Default::default()
        };
        let t = EnergyTables::compute(&s);
        for i in 0..BINS {
            let mu = (i as f32 + 0.5) / BINS as f32;
            println!(
                "mu {mu:.3}  metal {:.3}  out {:.3}  in {:.3}",
                t.metal[i], t.dielectric_outside[i], t.dielectric_inside[i]
            );
        }
    }
}

#[cfg(test)]
mod uploaded_tables {
    use super::*;

    /// The shader divides by these; a table of zeroes would silently kill the specular lobe.
    #[test]
    fn every_table_is_filled_and_positive() {
        let surface = OpenPbrSurface {
            specular_roughness: 0.2,
            coat_weight: 1.0,
            coat_roughness: 0.05,
            ..OpenPbrSurface::default()
        };
        let tables = EnergyTables::compute(&surface);
        let named: [(&str, &[f32; BINS]); 8] = [
            ("metal", &tables.metal),
            ("dielectric_outside", &tables.dielectric_outside),
            ("dielectric_inside", &tables.dielectric_inside),
            ("dielectric_r_outside", &tables.dielectric_r_outside),
            ("dielectric_r_inside", &tables.dielectric_r_inside),
            ("dielectric_t_outside", &tables.dielectric_t_outside),
            ("dielectric_t_inside", &tables.dielectric_t_inside),
            ("coat_r", &tables.coat_r),
        ];
        for (name, table) in named {
            assert!(
                table.iter().all(|v| v.is_finite() && *v >= 0.0),
                "{name} has a value that is not a finite albedo: {table:?}",
            );
            assert!(
                table.iter().any(|v| *v > 0.0),
                "{name} is all zeroes, which would divide the lobe away",
            );
        }
    }
}