crystal_ball 0.3.0

use std::sync::Arc;

use nanorand::tls::TlsWyRand;

use crate::color::{Color, Texture};
use crate::materials::{compute_color, compute_normal, Material};
use crate::math::{Hit, Point2, Ray, Vec3};
use crate::prelude::IOR;
use crate::util::random_float;

/// A general purpose PBR material.
///
/// For more information see the [glTF Spec](https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html#metallic-roughness-material)
/// as well as the following extensions:
/// - [KHR_materials_emissive_strength](https://github.com/KhronosGroup/glTF/blob/main/extensions/2.0/Khronos/KHR_materials_emissive_strength/README.md),
/// - [KHR_materials_ior](https://github.com/KhronosGroup/glTF/blob/main/extensions/2.0/Khronos/KHR_materials_ior/README.md)
/// - [KHR_materials_transmission](https://github.com/KhronosGroup/glTF/blob/main/extensions/2.0/Khronos/KHR_materials_transmission/README.md)
#[derive(Clone)]
pub struct PbrMaterial {
    /// The material's base color.
    ///
    /// Defaults to [`Color::WHITE`].
    pub base_color: Color,
    /// The materials's base color texture.
    ///
    /// If this is None, `base_color` remains unchanged.
    /// Otherwise each pixel value is multiplied with `base_color`.
    ///
    /// Defaults to [`None`].
    pub base_color_texture: Option<Arc<dyn Texture>>,
    /// The material's metallicness.
    ///
    /// Values should be in the range \[0, 1\].
    ///
    /// Defaults to `0.0`.
    pub metallic: f64,
    /// The material's roughness.
    ///
    /// Values should be in the range \[0, 1\].
    ///
    /// Defaults to `0.0`.
    pub roughness: f64,
    /// The material's metallic roughness texture.
    ///
    /// If this is None, `metallic` and `roughness` remain unchanged.
    /// Otherwise each the B channel of each pixel is multiplied with `metallic`
    /// and the G channel of each pixel is multiplied wie the `roughness`.
    ///
    /// Defaults to [`None`].
    pub metallic_roughness_texture: Option<Arc<dyn Texture>>,
    /// The material's transmission.
    ///
    /// Values should be in the range \[0, 1\].
    ///
    /// Defaults to `0.0`.
    pub transmission: f64,
    /// The material's index of refraction.
    ///
    /// Defaults to [`IOR::GLASS`].
    pub ior: f64,
    /// The material's emissive color.
    ///
    /// If this is black, the material is not emissive.
    ///
    /// Defaults to [`Color::BLACK`].
    pub emissive_color: Color,
    /// The materials's emissive texture.
    ///
    /// If this is None, `emissive_color` remains unchanged.
    /// Otherwise each pixel value is multiplied with `emissive_color`.
    ///
    /// Defaults to [`None`].
    pub emissive_texture: Option<Arc<dyn Texture>>,
    /// The material's emissive strength.
    ///
    /// If this is black, the material is not emissive.
    /// Otherwise the resulting color of `emissive_color` and `emissive_texture`
    /// is multiplied with `emissive_strength`.
    ///
    /// Defaults to `1.0`.
    pub emissive_strength: f64,
    /// The material's normal texture.
    ///
    /// This will only be used if the corresponding [`Shape`](crate::shapes::Shape) has tangents.
    ///
    /// Defaults to [`None`].
    pub normal_texture: Option<Arc<dyn Texture>>,
}

impl Default for PbrMaterial {
    fn default() -> Self {
        Self {
            base_color: Color::WHITE,
            base_color_texture: None,
            metallic: 0.0,
            roughness: 0.0,
            metallic_roughness_texture: None,
            transmission: 0.0,
            ior: IOR::GLASS,
            emissive_color: Color::BLACK,
            emissive_texture: None,
            emissive_strength: 1.0,
            normal_texture: None,
        }
    }
}

impl PbrMaterial {
    /// Calculate the normal for a given [`Hit`].
    ///
    /// If `normal_texture` is None or the corresponding [`Shape`](crate::shapes::Shape) has no tangents,
    /// the [`Hit`]'s normal is returned.
    /// Otherwise the normal is calculated using the `normal_texture`.
    fn normal(&self, hit: Hit) -> Vec3 {
        match hit.tangent {
            None => hit.normal,
            Some(tangent) => {
                let normal = hit.normal;
                let bitangent = Vec3::cross(normal, tangent.xyz()) * tangent.w;

                compute_normal(
                    normal,
                    tangent.xyz(),
                    bitangent,
                    &self.normal_texture,
                    hit.uv,
                )
            }
        }
    }

    /// Generate a new [`Ray`] for diffuse reflection.
    fn next_ray_diffuse(&self, hit: Hit, rng: &mut TlsWyRand) -> Option<Ray> {
        let direction = self.normal(hit) + Vec3::random_unit_vector(rng);
        Some(Ray::new(hit.position, direction, IOR::AIR))
    }

    /// Generate a new [`Ray`] for metallic reflection.
    fn next_ray_metallic(&self, ray: Ray, hit: Hit, rng: &mut TlsWyRand) -> Option<Ray> {
        let direction = ray.direction.reflect(self.normal(hit));

        let roughness = compute_color(
            Color::WHITE,
            &self.metallic_roughness_texture,
            self.roughness,
            hit.uv,
        )
        .g as f64;

        Some(Ray::new(
            hit.position,
            direction + roughness * Vec3::random_in_unit_sphere(rng),
            IOR::AIR,
        ))
    }

    /// Generate a new [`Ray`] for refraction.
    fn next_ray_refractive(&self, ray: Ray, hit: Hit, rng: &mut TlsWyRand) -> Option<Ray> {
        let (n1, n2) = (ray.last_ior, self.ior);
        let direction = ray.direction.refract(self.normal(hit), n1, n2, rng);

        Some(Ray::new(
            hit.position,
            direction + self.roughness * Vec3::random_in_unit_sphere(rng),
            self.ior,
        ))
    }
}

impl Material for PbrMaterial {
    fn next_ray(&self, ray: Ray, hit: Hit, rng: &mut TlsWyRand) -> Option<Ray> {
        let metallic = compute_color(
            Color::WHITE,
            &self.metallic_roughness_texture,
            self.metallic,
            hit.uv,
        )
        .b as f64;

        if random_float(rng, 0.0, 1.0) < metallic {
            self.next_ray_metallic(ray, hit, rng)
        } else if random_float(rng, 0.0, 1.0) < self.transmission {
            self.next_ray_refractive(ray, hit, rng)
        } else {
            self.next_ray_diffuse(hit, rng)
        }
    }

    fn get_color(&self, uv: Point2) -> (Color, Color) {
        (
            compute_color(self.base_color, &self.base_color_texture, 1.0, uv),
            compute_color(
                self.emissive_color,
                &self.emissive_texture,
                self.emissive_strength,
                uv,
            ),
        )
    }
}