bs-trace 0.3.0

use super::prelude::*;
use crate::image::prelude::*;
use crate::trace::material::Material;
use rand::RngCore;

/// Represents an object that a ray can hit.
pub trait Hittable {
    /// Computes the hit parameters for an input ray. Returns `None` if the ray does not hit this
    /// object.
    fn intersection(&self, rng: &mut Box<dyn RngCore>, ray: &Ray) -> Option<HitRecord>;
}

/// The relevant parameters for a ray hitting a hittable object.
pub struct HitRecord {
    pub(self) t: f64,
    pub(self) intersection: Vec3,
    pub(self) normal: Vec3,
    pub(self) attenuation: Colour,
    pub(self) bounce_direction: Option<Vec3>,
}

/// Real solution(s) for a quadratic equation.
enum QuadSoln {
    NoRoots,
    OneRoot(f64),
    TwoRoots(f64, f64),
}

/// A world, consisting of a set of objects that can be hit by rays.
pub struct World {
    objects: Vec<Box<dyn Hittable + Send + Sync>>,
}

/// A sphere, at the given centre with the given radius.
pub struct Sphere {
    centre: Vec3,
    radius: f64,
    material: Box<dyn Material>,
}

impl World {
    /// Creates a new world.
    pub fn new() -> Self {
        Self {
            objects: Vec::new(),
        }
    }

    /// Adds the given boxed hittable object to this world.
    pub fn add_object(&mut self, object: Box<dyn Hittable + Send + Sync>) {
        self.objects.push(object);
    }

    /// Traces a ray into the world, and returns the resultant colour.
    pub fn trace(&self, rng: &mut Box<dyn RngCore>, ray: &Ray) -> Colour {
        self.do_trace(rng, ray, 0)
    }

    fn do_trace(&self, rng: &mut Box<dyn RngCore>, ray: &Ray, bounce_count: usize) -> Colour {
        if bounce_count > 20 {
            Colour::BLACK
        } else {
            let hit = self
                .objects
                .iter()
                .flat_map(|obj| obj.intersection(rng, ray))
                .filter(|x| !x.t.is_nan())
                .filter(|x| x.t > 1e-3)
                .min_by(|x1, x2| x1.t.partial_cmp(&x2.t).unwrap()); // unwrap OK because partial_cmp is None iff nan

            match hit {
                Some(hit) => {
                    // this is normal-based colouring
                    // Colour::from(hit.normal.into_map(|x| (x + 1.0) / 2.0))
                    let new_origin = hit.intersection;
                    let new_direction = hit.bounce_direction;
                    if let Some(new_direction) = new_direction {
                        let new_ray = Ray::new(new_origin, new_direction);
                        hit.attenuation * self.do_trace(rng, &new_ray, bounce_count + 1)
                    } else {
                        Colour::BLACK
                    }
                }
                None => {
                    let t = (ray.direction.y() + 1.0) / 2.0;
                    let c0 = Vec3::new([1.0, 1.0, 1.0]);
                    let c1 = Vec3::new([0.5, 0.7, 1.0]);
                    Colour::from(Vec3::pure(1.0 - t) * c0 + Vec3::pure(t) * c1)
                }
            }
        }
    }
}

impl Sphere {
    /// Constructs a new sphere at the given centre with the given radius.
    pub const fn new(centre: Vec3, radius: f64, material: Box<dyn Material>) -> Self {
        Self {
            centre,
            radius,
            material,
        }
    }
}

impl Hittable for Sphere {
    fn intersection(&self, rng: &mut Box<dyn RngCore>, ray: &Ray) -> Option<HitRecord> {
        let co = self.centre - ray.origin;
        let a = ray.direction.dot(ray.direction);
        let b = -2.0 * (ray.direction.dot(co));
        let c = co.mag_sq() - self.radius * self.radius;
        let discriminant = b * b - 4.0 * a * c;

        let roots = if discriminant < 0.0 {
            QuadSoln::NoRoots
        } else if discriminant == 0.0 {
            let root = -b / (2.0 * a);
            QuadSoln::OneRoot(root)
        } else {
            let pos_root = (-b + discriminant.sqrt()) / (2.0 * a);
            let neg_root = (-b - discriminant.sqrt()) / (2.0 * a);
            QuadSoln::TwoRoots(pos_root, neg_root)
        };

        let nearest_pos_root = match roots {
            QuadSoln::NoRoots => None,
            QuadSoln::OneRoot(root) => {
                if root < 0.0 {
                    None
                } else {
                    Some(root)
                }
            }
            QuadSoln::TwoRoots(r1, r2) => {
                let min = r1.min(r2);
                let max = r1.max(r2);

                if max < 0.0 {
                    None
                } else if min < 0.0 {
                    Some(max)
                } else {
                    Some(min)
                }
            }
        }?;

        let intersection = ray.at(nearest_pos_root);
        let normal = (intersection - self.centre).normalised();
        let attenuation = self.material.colour();
        let bounce_direction = self.material.bounce(rng, ray.direction, normal);
        Some(HitRecord {
            t: nearest_pos_root,
            intersection,
            normal,
            attenuation,
            bounce_direction,
        })
    }
}