raytracers/rtiow/renderer/src/material.rs

use std::{fmt::Debug, sync::Arc};

use rand::{self, Rng};

use crate::{
    hitable::HitRecord,
    ray::Ray,
    texture::Texture,
    vec3::{dot, Vec3},
};

fn random_in_unit_sphere() -> Vec3 {
    let mut rng = rand::thread_rng();
    let v = Vec3::new(1., 1., 1.);
    loop {
        let p = 2. * Vec3::new(rng.gen(), rng.gen(), rng.gen()) - v;
        if p.squared_length() < 1. {
            return p;
        }
    }
}

#[derive(Default, Debug)]
pub struct ScatterResponse {
    pub scattered: Ray,
    pub attenutation: Vec3,
    pub reflected: bool,
}

pub trait Material: Send + Sync + Debug {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse;
    fn emitted(&self, _u: f32, _v: f32, _p: Vec3) -> Vec3 {
        Vec3::new(0., 0., 0.)
    }
}

impl Material for Arc<dyn Material> {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        (**self).scatter(r_in, rec)
    }
    fn emitted(&self, u: f32, v: f32, p: Vec3) -> Vec3 {
        (**self).emitted(u, v, p)
    }
}

impl Material for Box<dyn Material> {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        (**self).scatter(r_in, rec)
    }
    fn emitted(&self, u: f32, v: f32, p: Vec3) -> Vec3 {
        (**self).emitted(u, v, p)
    }
}

#[derive(Clone, Debug)]
pub struct Isotropic<T>
where
    T: Texture,
{
    albedo: T,
}

impl<T> Isotropic<T>
where
    T: Texture,
{
    pub fn new(texture: T) -> Isotropic<T> {
        Isotropic { albedo: texture }
    }
}

impl<T> Material for Isotropic<T>
where
    T: Texture,
{
    fn scatter(&self, _r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        let (u, v) = rec.uv;
        ScatterResponse {
            attenutation: self.albedo.value(u, v, rec.p),
            scattered: Ray::new(rec.p, random_in_unit_sphere(), 0.),
            reflected: true,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Lambertian<T>
where
    T: Texture,
{
    albedo: T,
}

impl<T> Lambertian<T>
where
    T: Texture,
{
    pub fn new(texture: T) -> Lambertian<T> {
        Lambertian { albedo: texture }
    }
}

impl<T> Material for Lambertian<T>
where
    T: Texture,
{
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        let target = rec.p + rec.normal + random_in_unit_sphere();
        let (u, v) = rec.uv;
        ScatterResponse {
            attenutation: self.albedo.value(u, v, rec.p),
            scattered: Ray::new(rec.p, target - rec.p, r_in.time),
            reflected: true,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Metal {
    albedo: Vec3,
    fuzzy: f32,
}

fn reflect(v: Vec3, n: Vec3) -> Vec3 {
    v - 2. * dot(v, n) * n
}

impl Metal {
    pub fn new<V>(albedo: V, fuzzy: f32) -> Metal
    where
        V: Into<Vec3>,
    {
        let fuzzy = fuzzy.min(1.);
        Metal {
            albedo: albedo.into(),
            fuzzy,
        }
    }
}

impl Material for Metal {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        let reflected = reflect(r_in.direction.unit_vector(), rec.normal);
        let scattered = Ray::new(
            rec.p,
            reflected + self.fuzzy * random_in_unit_sphere(),
            r_in.time,
        );
        ScatterResponse {
            scattered,
            attenutation: self.albedo,
            reflected: dot(scattered.direction, rec.normal) > 0.,
        }
    }
}

fn refract(v: Vec3, n: Vec3, ni_over_nt: f32) -> Option<Vec3> {
    let uv = v.unit_vector();
    let dt = dot(uv, n);
    let discriminant = 1. - ni_over_nt * ni_over_nt * (1. - dt * dt);
    if discriminant > 0. {
        return Some(ni_over_nt * (uv - n * dt) - n * discriminant.sqrt());
    }
    None
}

fn schlick(cosine: f32, ref_idx: f32) -> f32 {
    let mut r0 = (1. - ref_idx) / (1. + ref_idx);
    r0 = r0 * r0;
    r0 + (1. - r0) * (1. - cosine).powf(5.)
}

#[derive(Clone, Debug)]
pub struct Dielectric {
    ref_idx: f32,
}

impl Dielectric {
    pub fn new(ref_idx: f32) -> Dielectric {
        Dielectric { ref_idx }
    }
}

impl Material for Dielectric {
    fn scatter(&self, r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        let reflected = reflect(r_in.direction, rec.normal);
        let (outward_normal, ni_over_nt, cosine) = if dot(r_in.direction, rec.normal) > 0. {
            (
                -rec.normal,
                self.ref_idx,
                self.ref_idx * dot(r_in.direction, rec.normal) / r_in.direction.length(),
            )
        } else {
            (
                rec.normal,
                1. / self.ref_idx,
                -dot(r_in.direction, rec.normal) / r_in.direction.length(),
            )
        };
        let scattered = if let Some(refracted) = refract(r_in.direction, outward_normal, ni_over_nt)
        {
            let mut rng = rand::thread_rng();
            if rng.gen::<f32>() < schlick(cosine, self.ref_idx) {
                Ray::new(rec.p, reflected, r_in.time)
            } else {
                Ray::new(rec.p, refracted, r_in.time)
            }
        } else {
            Ray::new(rec.p, reflected, r_in.time)
        };

        ScatterResponse {
            attenutation: Vec3::new(1., 1., 1.),
            scattered,
            reflected: true,
        }
    }
}

#[derive(Debug)]
pub struct DiffuseLight<T>
where
    T: Texture,
{
    emit: T,
}

impl<T> DiffuseLight<T>
where
    T: Texture,
{
    pub fn new(emit: T) -> DiffuseLight<T> {
        DiffuseLight { emit }
    }
}

impl<T> Material for DiffuseLight<T>
where
    T: Texture,
{
    fn scatter(&self, _r_in: &Ray, _rec: &HitRecord) -> ScatterResponse {
        ScatterResponse {
            scattered: Default::default(),
            attenutation: Default::default(),
            reflected: false,
        }
    }

    fn emitted(&self, u: f32, v: f32, p: Vec3) -> Vec3 {
        self.emit.value(u, v, p)
    }
}

#[derive(Debug)]
pub struct DebugMaterial {}

impl Material for DebugMaterial {
    fn scatter(&self, _r_in: &Ray, rec: &HitRecord) -> ScatterResponse {
        let dir = Vec3::new(0., -1., -1.).unit_vector();
        ScatterResponse {
            scattered: Ray::new(rec.p, dir, 0.),
            attenutation: [1., 1., 1.].into(),
            reflected: true,
        }
    }
}

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

    #[test]
    fn arc_material() {
        let white: Arc<dyn Material> =
            Arc::new(Lambertian::new(ConstantTexture::new([0.73, 0.73, 0.73])));

        fn material_fn<M>(m: M)
        where
            M: Material,
        {
            let _ = m;
        }

        material_fn(Arc::clone(&white));
        material_fn(Arc::clone(&white));
    }
}