pbrt/src/core/material.rs

use bumpalo::Bump;
use enum_dispatch::enum_dispatch;
use std::ops::Deref;
use std::sync::Arc;

use crate::core::bssrdf::BSSRDF;
use crate::core::bxdf::{BSDF, DielectricBxDF, DiffuseBxDF};
use crate::core::interaction::InteractionTrait;
use crate::core::interaction::{Interaction, ShadingGeometry, SurfaceInteraction};
use crate::core::pbrt::Float;
use crate::core::texture::{FloatTexture, SpectrumTexture, TextureEvalContext, TextureEvaluator};
use crate::geometry::{Frame, Normal3f, Point2f, Point3f, Vector2f, Vector3f, VectorLike};
use crate::image::{Image, WrapMode, WrapMode2D};
use crate::spectra::{SampledWavelengths, Spectrum, SpectrumTrait};
use crate::utils::hash::hash_float;
use crate::utils::scattering::TrowbridgeReitzDistribution;

#[derive(Clone, Debug)]
pub struct MaterialEvalContext {
    pub texture: TextureEvalContext,
    pub wo: Vector3f,
    pub ns: Normal3f,
    pub dpdus: Vector3f,
}

impl Deref for MaterialEvalContext {
    type Target = TextureEvalContext;

    fn deref(&self) -> &Self::Target {
        &self.texture
    }
}

impl From<&SurfaceInteraction> for MaterialEvalContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            texture: TextureEvalContext::from(si),
            wo: si.common.wo,
            ns: si.shading.n,
            dpdus: si.shading.dpdu,
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct NormalBumpEvalContext {
    p: Point3f,
    uv: Point2f,
    n: Normal3f,
    shading: ShadingGeometry,
    dpdx: Vector3f,
    dpdy: Vector3f,
    // All 0
    dudx: Float,
    dudy: Float,
    dvdx: Float,
    dvdy: Float,
    face_index: usize,
}

impl From<&SurfaceInteraction> for NormalBumpEvalContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            p: si.p(),
            uv: si.uv,
            n: si.n(),
            shading: si.shading.clone(),
            dudx: si.dudx,
            dudy: si.dudy,
            dvdx: si.dvdx,
            dvdy: si.dvdy,
            dpdx: si.dpdx,
            dpdy: si.dpdy,
            face_index: si.face_index,
        }
    }
}

impl From<&NormalBumpEvalContext> for TextureEvalContext {
    fn from(ctx: &NormalBumpEvalContext) -> Self {
        Self {
            p: ctx.p,
            uv: ctx.uv,
            n: ctx.n,
            dpdx: ctx.dpdx,
            dpdy: ctx.dpdy,
            dudx: ctx.dudx,
            dudy: ctx.dudy,
            dvdx: ctx.dvdx,
            dvdy: ctx.dvdy,
            face_index: ctx.face_index,
        }
    }
}

pub fn normal_map(normal_map: &Image, ctx: &NormalBumpEvalContext) -> (Vector3f, Vector3f) {
    let wrap = WrapMode2D::from(WrapMode::Repeat);
    let uv = Point2f::new(ctx.uv[0], 1. - ctx.uv[1]);
    let r = normal_map.bilerp_channel_with_wrap(uv, 0, wrap);
    let g = normal_map.bilerp_channel_with_wrap(uv, 1, wrap);
    let b = normal_map.bilerp_channel_with_wrap(uv, 2, wrap);
    let mut ns = Vector3f::new(2.0 * r - 1.0, 2.0 * g - 1.0, 2.0 * b - 1.0);
    ns = ns.normalize();
    let frame = Frame::from_xz(ctx.shading.dpdu.normalize(), Vector3f::from(ctx.shading.n));
    ns = frame.from_local(ns);
    let ulen = ctx.shading.dpdu.norm();
    let vlen = ctx.shading.dpdv.norm();
    let dpdu = ctx.shading.dpdu.gram_schmidt(ns).normalize() * ulen;
    let dpdv = ctx.shading.dpdu.cross(dpdu).normalize() * vlen;
    (dpdu, dpdv)
}

pub fn bump_map<T: TextureEvaluator>(
    tex_eval: &T,
    displacement: &FloatTexture,
    ctx: &NormalBumpEvalContext,
) -> (Vector3f, Vector3f) {
    debug_assert!(tex_eval.can_evaluate(&[displacement], &[]));
    let mut du = 0.5 * (ctx.dudx.abs() + ctx.dudy.abs());
    if du == 0.0 {
        du = 0.0005;
    }

    let mut dv = 0.5 * (ctx.dvdx.abs() + ctx.dvdy.abs());
    if dv == 0.0 {
        dv = 0.0005;
    }
    let mut shifted_ctx = TextureEvalContext::from(ctx);
    shifted_ctx.p = ctx.p + ctx.shading.dpdu * du;
    shifted_ctx.uv = ctx.uv + Vector2f::new(du, 0.0);
    let u_displace = tex_eval.evaluate_float(displacement, &shifted_ctx);
    shifted_ctx.p = ctx.p + ctx.shading.dpdv * dv;
    shifted_ctx.uv = ctx.uv + Vector2f::new(0.0, dv);
    let v_displace = tex_eval.evaluate_float(displacement, &shifted_ctx);
    let center_ctx = TextureEvalContext::from(ctx);
    let displace = tex_eval.evaluate_float(displacement, &center_ctx);

    let d_displace_du = (u_displace - displace) / du;
    let d_displace_dv = (v_displace - displace) / dv;
    let n_vec = Vector3f::from(ctx.shading.n);
    let dndu_vec = Vector3f::from(ctx.shading.dndu);
    let dndv_vec = Vector3f::from(ctx.shading.dndv);

    let dpdu = ctx.shading.dpdu + n_vec * d_displace_du + dndu_vec * displace;
    let dpdv = ctx.shading.dpdv + n_vec * d_displace_dv + dndv_vec * displace;
    (dpdu, dpdv)
}

#[enum_dispatch]
pub trait MaterialTrait: Send + Sync + std::fmt::Debug {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
        scratch: &'a Bump,
    ) -> BSDF<'a>;

    fn get_bssrdf<'a, T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>>;

    fn can_evaluate_textures(&self, tex_eval: &dyn TextureEvaluator) -> bool;
    fn get_normal_map(&self) -> Option<&Image>;
    fn get_displacement(&self) -> Option<FloatTexture>;
    fn has_surface_scattering(&self) -> bool;
}

#[derive(Clone, Debug)]
#[enum_dispatch(MaterialTrait)]
pub enum Material {
    CoatedDiffuse(CoatedDiffuseMaterial),
    CoatedConductor(CoatedConductorMaterial),
    Conductor(ConductorMaterial),
    Dielectric(DielectricMaterial),
    Diffuse(DiffuseMaterial),
    DiffuseTransmission(DiffuseTransmissionMaterial),
    Hair(HairMaterial),
    Measured(MeasuredMaterial),
    Subsurface(SubsurfaceMaterial),
    ThinDielectric(ThinDielectricMaterial),
    Mix(MixMaterial),
}

#[derive(Clone, Debug)]
pub struct CoatedDiffuseMaterial;
impl MaterialTrait for CoatedDiffuseMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }

    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }

    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}

#[derive(Clone, Debug)]
pub struct CoatedConductorMaterial;
impl MaterialTrait for CoatedConductorMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }
    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }

    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct ConductorMaterial;
impl MaterialTrait for ConductorMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }
    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct DielectricMaterial {
    normal_map: Option<Arc<Image>>,
    displacement: FloatTexture,
    u_roughness: FloatTexture,
    v_roughness: FloatTexture,
    remap_roughness: bool,
    eta: Spectrum,
}

impl MaterialTrait for DielectricMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
        scratch: &'a Bump,
    ) -> BSDF<'a> {
        let mut sampled_eta = self.eta.evaluate(lambda[0]);
        if !self.eta.is_constant() {
            lambda.terminate_secondary();
        }

        if sampled_eta == 0.0 {
            sampled_eta = 1.0;
        }

        let mut u_rough = tex_eval.evaluate_float(&self.u_roughness, ctx);
        let mut v_rough = tex_eval.evaluate_float(&self.v_roughness, ctx);

        if self.remap_roughness {
            u_rough = TrowbridgeReitzDistribution::roughness_to_alpha(u_rough);
            v_rough = TrowbridgeReitzDistribution::roughness_to_alpha(v_rough);
        }

        let distrib = TrowbridgeReitzDistribution::new(u_rough, v_rough);
        let bxdf = scratch.alloc(DielectricBxDF::new(sampled_eta, distrib));

        BSDF::new(ctx.ns, ctx.dpdus, Some(bxdf))
    }

    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        None
    }

    fn can_evaluate_textures(&self, tex_eval: &dyn TextureEvaluator) -> bool {
        tex_eval.can_evaluate(&[&self.u_roughness, &self.v_roughness], &[])
    }

    fn get_normal_map(&self) -> Option<&Image> {
        self.normal_map.as_deref()
    }

    fn get_displacement(&self) -> Option<FloatTexture> {
        Some(self.displacement.clone())
    }

    fn has_surface_scattering(&self) -> bool {
        false
    }
}
#[derive(Clone, Debug)]
pub struct DiffuseMaterial {
    normal_map: Option<Arc<Image>>,
    displacement: FloatTexture,
    reflectance: SpectrumTexture,
}

impl MaterialTrait for DiffuseMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
        scratch: &'a Bump,
    ) -> BSDF<'a> {
        let r = tex_eval.evaluate_spectrum(&self.reflectance, ctx, lambda);
        let bxdf = scratch.alloc(DiffuseBxDF::new(r));
        BSDF::new(ctx.ns, ctx.dpdus, Some(bxdf))
    }

    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, tex_eval: &dyn TextureEvaluator) -> bool {
        tex_eval.can_evaluate(&[], &[&self.reflectance])
    }

    fn get_normal_map(&self) -> Option<&Image> {
        self.normal_map.as_deref()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        Some(self.displacement.clone())
    }
    fn has_surface_scattering(&self) -> bool {
        false
    }
}
#[derive(Clone, Debug)]
pub struct DiffuseTransmissionMaterial;
impl MaterialTrait for DiffuseTransmissionMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct HairMaterial;
impl MaterialTrait for HairMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}

#[derive(Clone, Debug)]
pub struct MeasuredMaterial;
impl MaterialTrait for MeasuredMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct SubsurfaceMaterial;
impl MaterialTrait for SubsurfaceMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }

    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct ThinDielectricMaterial;
impl MaterialTrait for ThinDielectricMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
        _scratch: &'a Bump,
    ) -> BSDF<'a> {
        todo!()
    }
    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        todo!()
    }
    fn can_evaluate_textures(&self, _tex_eval: &dyn TextureEvaluator) -> bool {
        todo!()
    }
    fn get_normal_map(&self) -> Option<&Image> {
        todo!()
    }
    fn get_displacement(&self) -> Option<FloatTexture> {
        todo!()
    }
    fn has_surface_scattering(&self) -> bool {
        todo!()
    }
}
#[derive(Clone, Debug)]
pub struct MixMaterial {
    pub amount: FloatTexture,
    pub materials: [Box<Material>; 2],
}

impl MixMaterial {
    pub fn choose_material<T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
    ) -> &Material {
        let amt = tex_eval.evaluate_float(&self.amount, ctx);

        if amt <= 0.0 {
            return &self.materials[0];
        }
        if amt >= 1.0 {
            return &self.materials[1];
        }

        let u = hash_float(&(ctx.p, ctx.wo));

        if amt < u {
            &self.materials[0]
        } else {
            &self.materials[1]
        }
    }
}

impl MaterialTrait for MixMaterial {
    fn get_bxdf<'a, T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
        scratch: &'a Bump,
    ) -> BSDF<'a> {
        let chosen_mat = self.choose_material(tex_eval, ctx);
        chosen_mat.get_bxdf(tex_eval, ctx, lambda, scratch)
    }

    fn get_bssrdf<'a, T>(
        &self,
        _tex_eval: &T,
        _ctx: &MaterialEvalContext,
        _lambda: &SampledWavelengths,
    ) -> Option<BSSRDF<'a>> {
        None
    }

    fn can_evaluate_textures(&self, tex_eval: &dyn TextureEvaluator) -> bool {
        tex_eval.can_evaluate(&[&self.amount], &[])
    }

    fn get_normal_map(&self) -> Option<&Image> {
        None
    }

    fn get_displacement(&self) -> Option<FloatTexture> {
        None
        // panic!(
        //     "MixMaterial::get_displacement() shouldn't be called. \
        //        Displacement is not supported on Mix materials directly."
        // );
    }

    fn has_surface_scattering(&self) -> bool {
        false
    }
}