Ray tracer

raymarcher.rs

use cgmath::prelude::*;
use cgmath::{vec2, vec3, ElementWise, EuclideanSpace, InnerSpace, SquareMatrix};
use crossbeam::thread;
use image::{ImageBuffer, RgbImage};
use reduce::Reduce;
use std::cmp::{PartialOrd};
use std::f64;
use rand::Rng;
use std::sync::Mutex;

type Scalar = f64;
type Vec2 = cgmath::Vector2<Scalar>;
type Vec3 = cgmath::Vector3<Scalar>;
type Vec4 = cgmath::Vector4<Scalar>;
type Basis3 = cgmath::Basis3<Scalar>;
type Matrix3 = cgmath::Matrix3<Scalar>;
type Point3 = cgmath::Point3<Scalar>;

const THREADS: u16 = 16;
const EPSILON: f64 = 0.001;
const BLOCK_SIZE: u16 = 16;
const MAXIMUM_RAY_MARCH_DISTANCE: Scalar = 1000.0;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DebugId {
    NONE,
    SPHERE
}

#[derive(Copy, Clone)]
struct Ray {
    origin: Point3,
    direction: Vec3,
}

struct Material {
    diffuse_color: Vec3,
    radiant_exitance: Vec3
}

struct LightSample {
    direction: Vec3,
    color: Vec3,
    pdf: Scalar
}

trait Object: Send + Sync {
    fn debug_id(&self) -> DebugId;
    fn as_object(&self) -> &dyn Object;
    fn material_at(&self, point: Point3, eye: Vec3) -> Material;
    fn tangent_space_at(&self, point: Point3) -> Matrix3;
}

trait AnalyticObject: Object {
    fn intersection(&self, ray: Ray) -> Option<Scalar>;
}
trait NumericalObject: Object {
    fn signed_distance(&self, point: Point3) -> Scalar;
}

trait Light {
    fn sample_incident(&self, target_point: Vec3, tangent_space: &Basis3)
        -> LightSample;
}

pub struct Sun {
    orientation: Matrix3,
    inverted_orientation: Matrix3,
    sun_color: Vec3,
    sky_color: Vec3,
    angular_diameter: Scalar,
    debug_id: DebugId,

    sun_probability: Scalar,
    sky_probability: Scalar,
    sun_pdf: Scalar,
    sky_pdf: Scalar,
}

fn luminance_for_rgb(rgb: Vec3) -> Scalar {
    0.2126*rgb.x + 0.7152*rgb.y + 0.0722*rgb.z
}

impl Sun {
    // Orientation is given so that z points outward from center of sun toward
    // scene, and y would typically point "up" in the general direction of
    // zenith for a sun at the horizon but rotate appropriately when sun is at
    // zenith.
    fn new(orientation: Matrix3,
        sun_color: Vec3,
        sky_color: Vec3,
        angular_diameter: f64,
        debug_id: DebugId,
    ) -> Sun {
        // Calculate total intensity of sun vs total intensity of sky sphere.
        // https://en.wikipedia.org/wiki/Spherical_cap
        // We assume unit sphere for all calculations and it'll cancel out
        // in the end.
        let sun_area = 2.0*f64::consts::PI*(1.0-(2.0*angular_diameter).cos());
        let sky_area = 4.0*f64::consts::PI - sun_area;
        let sun_brightness = sun_area*luminance_for_rgb(sun_color);
        let sky_brightness = sky_area*luminance_for_rgb(sky_color);
        let sun_probability = sun_brightness/(sun_brightness+sky_brightness);
        let sky_probability = 1.0 - sun_probability;

        Sun {
            orientation,
            inverted_orientation: orientation.invert().unwrap(),
            sun_color,
            sky_color,
            angular_diameter,
            debug_id,

            sun_probability,
            sky_probability,
            sun_pdf: sun_probability/sun_area,
            sky_pdf: sky_probability/sky_area
        }
    }
}

impl Light for Sun {
    fn sample_incident(&self, target_point: Vec3, tangent_space: &Basis3) -> LightSample {
        // https://fasiha.github.io/post/spherical-cap/
        // https://stackoverflow.com/questions/38997302/create-random-unit-vector-inside-a-defined-conical-region/39003745#39003745
        let mut rng = &rand::thread_rng();
        let hit_sun = rng.gen::<Scalar>() < self.sun_probability;
        let z = if hit_sun {
            // Select random z on sun
            let min_z = (self.angular_diameter/2.0).cos();
            rng.gen::<Scalar>() * (1.0 - min_z) + min_z
        } else {
            // Select random z on sky.
            let max_z = (self.angular_diameter/2.0).cos();
            rng.gen::<Scalar>() * max_z
        };
        let phi = rng.gen::<Scalar>() * 2.0 * f64::consts::PI;
        let x = (1.0-z*z).sqrt() * phi.cos();
        let y = (1.0-z*z).sqrt() * phi.sin();
        let sample_direction = vec3(x, y, z);

        let d: Vec3 = self.inverted_orientation*sample_direction;
        LightSample {
            direction: d,
            color: self.sun_color,
            pdf: self.sun_pdf
        }
    }
}

impl Object for Sun {
    fn debug_id(&self) -> DebugId {self.debug_id}
    fn as_object(&self) -> &dyn Object {
        self
    }

    fn material_at(&self, point: Point3, eye: Vec3) -> Material {
        let angle = eye.dot(self.orientation.z).acos();
        //println!("{} > {} ({})", angle/f64::consts::PI*180.0, self.angular_diameter/f64::consts::PI*180.0, eye.magnitude());
        let color = if angle > self.angular_diameter {
            self.sky_color
        } else {
            //println!("sun!");
            self.sun_color
        };
        Material {
            diffuse_color: vec3(0.0, 0.0, 0.0),
            radiant_exitance: color,
        }
    }
    fn tangent_space_at(&self, point: Point3) -> Matrix3 {
        Matrix3::new(
            0.0, 0.0, 0.0,
            0.0, 0.0, 0.0,
            0.0, 0.0, 0.0)
    }
}

struct Scene {
    analytic_objects: Vec<Box<dyn AnalyticObject>>,
    numerical_objects: Vec<Box<dyn NumericalObject>>,
    lights: Vec<Box<dyn Light>>,
}

struct Sphere {
    origin: Point3,
    radius: Scalar,
    debug_id: DebugId
}

impl Object for Sphere {
    fn debug_id(&self) -> DebugId {self.debug_id}
    fn material_at(&self, point: Point3, eye: Vec3) -> Material {
        Material {
            diffuse_color: vec3(1.0, 1.0, 1.0),
            radiant_exitance: vec3(0.0, 0.0, 0.0)
        }
    }
    fn tangent_space_at(&self, point: Point3) -> Matrix3 {
        tangent_space_from_normal((point - self.origin).normalize())
    }
    fn as_object(&self) -> &dyn Object {
        self
    }
}

impl AnalyticObject for Sphere {
    fn intersection(&self, ray: Ray) -> Option<Scalar> {
        let radius2 = self.radius * self.radius;

        let L = self.origin - ray.origin;
        let tca = L.dot(ray.direction);
        if tca < 0.0 {
            return None;
        }
        let d2 = L.dot(L) - tca * tca;
        if d2 > radius2 {
            return None;
        }
        let thc = (radius2 - d2).sqrt();
        let mut t0 = tca - thc;
        let mut t1 = tca + thc;
        if t0 > t1 {
            let t = t0;
            t0 = t1;
            t1 = t;
        }
        if t0 < 0.0 {
            t0 = t1;
            if t0 < 0.0 {
                return None;
            }
        }
        Some(t0)
    }
}

impl NumericalObject for Sphere {
    fn signed_distance(&self, point: Point3) -> Scalar {
        (point - self.origin).magnitude() - self.radius
    }
}

struct Metaball {
    balls: Vec<Sphere>,
    sharpness: f64,
    debug_id: DebugId,
}

impl Object for Metaball {
    fn debug_id(&self) -> DebugId {self.debug_id}
    fn material_at(&self, point: Point3, eye: Vec3) -> Material {
        Material {
            diffuse_color: vec3(1.0, 1.0, 1.0),
            radiant_exitance: vec3(0.0, 0.0, 0.0)
        }
    }
    fn tangent_space_at(&self, point: Point3) -> Matrix3 {
        tangent_space_from_normal(calculate_normal(self, point))
    }
    fn as_object(&self) -> &dyn Object {
        self
    }
}

impl NumericalObject for Metaball {
    fn signed_distance(&self, point: Point3) -> Scalar {
        let mut sum = 0.0;
        for ball in self.balls.iter() {
            let dist = (point - ball.origin).magnitude() - ball.radius;
            sum += (-self.sharpness * dist).exp2();
        }
        -sum.log2() / self.sharpness
    }
}

struct InfinitePlane {
    origin: Point3,
    forward: Vec3,
    right: Vec3,
    up: Vec3,
    debug_id: DebugId
}

impl Object for InfinitePlane {
    fn debug_id(&self) -> DebugId {self.debug_id}
    fn material_at(&self, point: Point3, eye: Vec3) -> Material {
        let point = point.to_vec() - self.origin.to_vec();
        let y = self.forward.dot(point);
        let x = self.right.dot(point);
        let x = x * 2.0;
        let y = y * -2.0;
        let x = x % 2.0;
        let y = y % 2.0;
        let x_flip = if x < 0.0 { x < -1.0 } else { x < 1.0 };
        let y_flip = if y < 0.0 { y < -1.0 } else { y < 1.0 };
        let color = if x_flip ^ y_flip {
            vec3(1.0, 0.5, 0.5)
        } else {
            vec3(1.0, 1.0, 1.0)
        };

        Material {
            diffuse_color: color,
            radiant_exitance: vec3(0.0, 0.0, 0.0)
        }
    }

    fn tangent_space_at(&self, point: Point3) -> Matrix3 {
        Matrix3::from_cols(self.right, self.forward, self.up)
    }

    fn as_object(&self) -> &dyn Object {
        self
    }
}

impl AnalyticObject for InfinitePlane {
    fn intersection(&self, ray: Ray) -> Option<Scalar> {
        // https://samsymons.com/blog/math-notes-ray-plane-intersection/
        let denominator = self.up.dot(ray.direction);

        // 0.0001 is an arbitrary epsilon value. We just want
        // to avoid working with intersections that are almost
        // orthogonal.
        if denominator.abs() > 0.0001 {
            let difference = self.origin - ray.origin;
            let t = difference.dot(self.up) / denominator;

            if t > 0.0001 {
                return Some(t);
            }
        }

        return None;
    }
}

impl NumericalObject for InfinitePlane {
    fn signed_distance(&self, point: Point3) -> Scalar {
        let from = point.to_vec() - self.origin.to_vec();
        from.dot(self.up)
    }
}

fn calculate_normal(object: &dyn NumericalObject, point: Point3) -> Vec3 {
    let small_step = vec2(0.001, 0.0);
    let gradient_x = object.signed_distance(point + small_step.xyy())
        - object.signed_distance(point - small_step.xyy());
    let gradient_y = object.signed_distance(point + small_step.yxy())
        - object.signed_distance(point - small_step.yxy());
    let gradient_z = object.signed_distance(point + small_step.yyx())
        - object.signed_distance(point - small_step.yyx());
    vec3(gradient_x, gradient_y, gradient_z).normalize()
}

fn trace_ray(scene: &Scene, ray: Ray) -> Option<(&dyn Object, Scalar)> {
    // Distance to nearest analytic object.
    let distances = scene
        .analytic_objects
        .iter()
        .filter_map(|obj| obj.intersection(ray).map(|d| (obj.as_object(), d)));
    let nearest_analytic =
        distances.reduce(
            |(o1, d1), (o2, d2)| if d1 < d2 { (o1, d1) } else { (o2, d2) },
        ); //.or(scene.background.as_ref().map(|b| (&**b, f64::INFINITY)));

    let nearest_analytic = nearest_analytic.map(|(o, d)| (o.as_object(), d));
    if scene.numerical_objects.is_empty() {
        return nearest_analytic;
    }

    let objects: Vec<_> =
        scene.numerical_objects.iter().map(|b| &**b).collect();

    let mut marched_distance = 0.0;
    loop {
        if let Some((analytic_obj, analytic_dist)) = nearest_analytic {
            if marched_distance >= analytic_dist {
                return nearest_analytic;
            }
        }
        // TODO actually use ray direction for optimized signed distance
        // (but we still need signed_distance for numerical normal calculation)
        let ray = Ray {
            origin: ray.origin + marched_distance * ray.direction,
            ..ray
        };
        // Here we could cull all objects that have a distance further than the
        // nearest analytical object.
        let (nearest_obj, nearest_dist) = objects
            .iter()
            .map(|o| (o, o.signed_distance(ray.origin)))
            .reduce(
                |(o1, d1), (o2, d2)| if d1 < d2 { (o1, d1) } else { (o2, d2) },
            )
            .unwrap();
        if nearest_dist < EPSILON {
            return Some((nearest_obj.as_object(), marched_distance));
        } else {
            marched_distance += nearest_dist;
        }

        if marched_distance > MAXIMUM_RAY_MARCH_DISTANCE {
            return nearest_analytic;
        }
    }
}

fn normal_to_rgb(normal: Vec3) -> Vec3 {
    let r = normal.x * 0.5 + 0.5;
    let g = normal.y * 0.5 + 0.5;
    let b = normal.z * 0.5 + 0.5;
    vec3(r, g, b)
}
fn noop() {}

// https://stackoverflow.com/questions/33757842/how-to-calculate-coordinate-system-from-one-vector/33758795#33758795
// Not sure I get this off the top of my head, but should be possible to work
// out a 90 degree rotation in one plane using the equation v.dot(x) = 0
// See also https://www.gamedev.net/forums/topic/357797-rotate-a-vector-by-90-degrees/?tab=comments#comment-3348469
// I think the trick is to do 2D rotation and set the other component to zero
// fn make_coordinate_system(v: Vec3) -> (Vec3, Vec3) {
//     if v.z < v.x && v.z < v.y {
//         (
//             vec3(v.y, -v.x, 0.0).normalize(),
//             vec3(-v.x*v.z, -v.y*v.z, v.x*v.x + v.y*v.y).normalize()
//         )
//     } else if(v.y < v.x && v.y < v.z) {
//         (
//             vec3(v.z, 0.0, v.x)
//         )
//     }
//     let mut smallest = v.x;
//     if v.y < smallest {
//         smallest = v.y;
//     }
//     if v.z < smallest {
//         smallest = v.z;
//     }

//     let r = vec3()
// }

fn make_coordinate_system(v: Vec3) -> (Vec3, Vec3) {
    // https://www.gamedev.net/forums/topic/357797-rotate-a-vector-by-90-degrees/?tab=comments#comment-3348469
    // Rotate v 90 degrees clockwise

    let (ax, ay, az) = (v.x.abs(), v.y.abs(), v.z.abs());

    let s =
        if ax < ay {
            // y is not smallest.
            if ax < az {
                // x is smallest.
                vec3(0.0, -v.z, v.y)
            } else {
                // y is not smallest, x is not smallest, so z is.
                vec3(v.y, -v.x, 0.0)
            }
        } else {
            // x is not smallest.
            if ay < az {
                // y is smallest.
                vec3(-v.z, 0.0, v.x)
            } else {
                // x is not smallest, y is not smallest, so z is.
                vec3(v.y, -v.x, 0.0)
            }
        };
    let s = s.normalize();
    // TODO this could be optimized by moving it into the if statement
    // and taking the zero component into account.
    let t = s.cross(v);
    (s, t)
}

fn tangent_space_from_normal(normal: Vec3) -> Matrix3 {
    let (u, v) = make_coordinate_system(normal);
    Matrix3::from_cols(u, v, normal)
}

fn shade_gi(
    object: &dyn Object,
    position: Point3,
    eye: Vec3,
    scene: &Scene
) -> Vec3 {

    object.tangent_space_at()

    for light in scene.lights.iter() {
        light.sample_incident()
    }
    // let pdf = 1.0 / (2.0 * f64::consts::PI);
    // let normal = object.normal_at(position);
    // let (s, t) = make_coordinate_system(normal);
    // // println!("{} {} {}", normal.magnitude(), s.magnitude(), t.magnitude());
    // //let ld = random_uniform_hemisphere_point();
    // let ld = random_importance_sampled_hemisphere_point();
    // let direction = ld.x*s + ld.y*t + ld.z*normal;
    // // println!("{}", direction.dot(normal));

    // let ray = Ray {
    //     origin: position + EPSILON * normal,
    //     direction: direction
    // };
    // if object.debug_id() == DebugId::SPHERE {
    //     noop();
    // }
    // //let light=
    // let light = if let Some((object, distance)) = trace_ray(scene, ray) {
    //     let point = ray.origin + distance*ray.direction;
    //     object.material_at(point, ray.direction).radiant_exitance
    // } else {
    //     scene.background.as_ref().map_or(vec3(0.0, 0.0, 0.0), |b| {
    //         b.material_at(ray.origin, ray.direction).radiant_exitance
    //     })
    // };

    // //let light_angle = normal.dot(direction);
    // //let light = (light_angle/pdf)*light;
    // let light = light/pdf;
    // object.material_at(position, eye).diffuse_color.mul_element_wise(light)
}

fn shade_recursive(
    object: &dyn Object,
    position: Point3,
    eye: Vec3,
    scene: &Scene,
) -> Vec3 {
    let normal = object.normal_at(position);
    let mut light = vec3(0.2, 0.2, 0.2);

    for sun in scene.sun.iter() {
        let ray = Ray {
            origin: position + EPSILON * normal, // TODO use ray direction or some such instead of normal?
            direction: -sun.orientation.z,
        };
        if let Some(_) = trace_ray(scene, ray) {
            continue;
        }

        let angle = normal.dot(-sun.orientation.z) + sun.angular_diameter;
        let angle = clamp(angle, 0.0, 1.0);
        light += sun.sun_color * angle;
    }
    let material = object.material_at(position, eye);
    material.diffuse_color.mul_element_wise(light)
}

fn shade(
    object: &dyn Object,
    position: Point3,
    eye: Vec3,
    scene: &Scene,
) -> Vec3 {
    let normal = object.normal_at(position);
    let mut light = vec3(0.2, 0.2, 0.2);
    for sun in scene.sun.iter() {
        let angle = normal.dot(-sun.orientation.z) + sun.angular_diameter;
        let angle = clamp(angle, 0.0, 1.0);
        light += sun.sun_color * angle;
    }
    let material = object.material_at(position, eye);
    material.diffuse_color.mul_element_wise(light)
}

fn trace_ray_and_shade(scene: &Scene, ray: Ray) -> Option<Vec3> {
    if let Some((object, distance)) = trace_ray(scene, ray) {
        let point = ray.origin + distance * ray.direction;
        Some(shade_gi(object, point, ray.direction, scene))
    } else {
        scene.background.as_ref().map(|background| {
            background.material_at(ray.origin, ray.direction).radiant_exitance
        })
    }
}

fn trace_ray_shade_by_normal(scene: &Scene, ray: Ray) -> Option<Vec3> {
    if let Some((object, distance)) = trace_ray(scene, ray) {
        let point = ray.origin + distance * ray.direction;
        let normal = object.normal_at(point);
        Some(vec3(0.5, 0.5, 0.5) + 0.5 * normal)
    } else {
        None
    }
}

fn clamp<T: PartialOrd>(v: T, min: T, max: T) -> T {
    if v < min {
        min
    } else if max < v {
        max
    } else {
        v
    }
}

struct HdrImage {
    image: Vec<Vec3>,
    width: u16,
    height: u16,
}

impl HdrImage {
    fn new(width: u16, height: u16) -> HdrImage {
        let w32 = width as u32;
        let h32 = height as u32;
        HdrImage {
            image: vec![vec3(0.0, 0.0, 0.0); (w32 * h32) as usize],
            width,
            height,
        }
    }
    fn put_rgb(&mut self, x: u16, y: u16, color: Vec3) {
        let (x, y) = (x as u32, y as u32);
        let stride = self.width as u32;
        self.image[(y * stride + x as u32) as usize] = color;
    }
    fn get_rgb(&self, x: u16, y: u16) -> Vec3 {
        let (x, y) = (x as u32, y as u32);
        let w = self.width as u32;
        self.image[(y * w + x) as usize].xyz()
    }
    fn develop_rgb(&self) -> RgbImage {
        let mut max = 0.0;
        for c in self.image.iter() {
            for i in 0..3 {
                if c[i] > max {
                    max = c[i]
                }
            }
        }
        let multiplier = 1.0 / max;

        let mut image: RgbImage =
            ImageBuffer::new(self.width as u32, self.height as u32);
        for y in 0..self.height {
            for x in 0..self.width {
                let hdr = self.get_rgb(x, y) * multiplier;
                let r = clamp((hdr.x * 256.0) as i32, 0, 255) as u8;
                let g = clamp((hdr.y * 256.0) as i32, 0, 255) as u8;
                let b = clamp((hdr.z * 256.0) as i32, 0, 255) as u8;
                image.put_pixel(x as u32, y as u32, image::Rgb([r, g, b]))
            }
        }
        image
    }
}


// TODO other sampler with better distribution
fn random_uniform_point_on_circle<R: Rng>(rng: &mut R) -> Vec2 {
    let a = rng.gen::<Scalar>() * 2.0 * f64::consts::PI;
    let r = rng.gen::<Scalar>().sqrt();

    vec2(r*a.cos(), r*a.sin())
}

fn random_uniform_hemisphere_point() -> Vec3 {
    // Float z = u[0];
    // Float r = std::sqrt(std::max((Float)0, (Float)1. - z * z));
    // Float phi = 2 * Pi * u[1];
    // return Vector3f(r * std::cos(phi), r * std::sin(phi), z);

    let rng = &mut rand::thread_rng();
    let z: Scalar = rng.gen();
    let u: Scalar = rng.gen();
    let r = (1.0 - z*z).sqrt();
    let phi = 2.0*f64::consts::PI * u;
    vec3(r*phi.cos(), r*phi.sin(), z)

    // let rng = &mut rand::thread_rng();
    // let r1: Scalar = rng.gen();
    // let r2: Scalar = rng.gen();
    // let sin_theta = (1.0 - r1*r1);
    // let phi = 2.0*f64::consts::PI*r2;
    // let x = sin_theta * phi.cos();
    // let z = sin_theta * phi.sin();
    // vec3(x, r1, z)
}
fn random_importance_sampled_hemisphere_point() -> Vec3 {
    let rng = &mut rand::thread_rng();
    let sin2_theta: Scalar = rng.gen();
    let cos2_theta = 1.0 - sin2_theta;
    let sin_theta = sin2_theta.sqrt();
    let cos_theta = cos2_theta.sqrt();
    let orientation = rng.gen::<Scalar>()*2.0*f64::consts::PI;
    vec3(sin_theta*orientation.cos(), sin_theta*orientation.sin(), cos_theta)
}
// Vector RandomDirectionLambertian:
//     sin2_theta = uniform random number between 0 and 1 // the uniform random number is equal to sin^2(theta)
//     cos2_theta = 1 - sin2_theta // cos^2(x) + sin^2(x) = 1
//     sin_theta = sqrt(sin2_theta)
//     cos_theta = sqrt(cos2_theta)
//     orientation = uniform random number between 0 and 2pi
//     return Vector(sin_theta * cos(orientation),
//                     cos_theta,
//                     sin_theta * sin(orientation))

struct PinholeCamera {
    origin: Point3,
    right: Vec3,
    up: Vec3,
    forward: Vec3,
    width: Scalar,
    height: Scalar,
    width_pixels: u16,
    height_pixels: u16,
    focal_length: Scalar,
}

// https://en.wikipedia.org/wiki/Angle_of_view
// For example, for 35 mm film which is 36 mm wide and 24 mm high, d=36 mm
// would be used to obtain the horizontal angle of view and d = 24 mm for the
// vertical angle.
trait Camera: Send + Sync {
    fn get_ray(&self, pixel_x: Scalar, pixel_y: Scalar) -> Ray;
    fn raster_dimensions(&self) -> (u16, u16);
}

impl Camera for PinholeCamera {
    fn get_ray(&self, pixel_x: Scalar, pixel_y: Scalar) -> Ray {
        // Calculate position of pixel on film, with film center as (0, 0)
        let x = (pixel_x/self.width_pixels as Scalar) - 0.5;
        let y = (pixel_y/self.height_pixels as Scalar) - 0.5;
        let x = x*self.width;
        let y = y*self.height;

        // 3D position on film
        let sample_position = self.origin + x*self.right + y*self.up - self.focal_length*self.forward;
        let direction = (self.origin - sample_position).normalize();
        Ray {
            origin: self.origin,
            direction: direction
        }
    }
    fn raster_dimensions(&self) -> (u16, u16) {
        (self.width_pixels, self.height_pixels)
    }
}

struct ThinLensCamera {
    origin: Point3,
    right: Vec3,
    up: Vec3,
    forward: Vec3,
    width: Scalar,
    height: Scalar,
    width_pixels: u16,
    height_pixels: u16,
    focal_length: Scalar,
    aperture: Scalar,
    focal_distance: Scalar,
}

impl Camera for ThinLensCamera {
    fn get_ray(&self, pixel_x: Scalar, pixel_y: Scalar) -> Ray {
        // Determine sample position on film, with film center as (0, 0)
        let x = (pixel_x/self.width_pixels as Scalar) - 0.5;
        let y = (pixel_y/self.height_pixels as Scalar) - 0.5;
        let x = x*self.width;
        let y = y*self.height;
        let sample_position = self.origin + x*self.right + y*self.up - self.focal_length*self.forward;

        // Compute the desired focus point.
        let pinhole_ray = (self.origin - sample_position).normalize();
        let intersection_distance_along_ray = self.focal_distance/pinhole_ray.z;
        let focus_point = self.origin + intersection_distance_along_ray*pinhole_ray;

        // Pick random point on the lens disk where the ray will "start" (after
        // being refracted through the lens from the sample point on the
        // film).
        let lens_radius = self.focal_length/(self.aperture*2.0);
        let lens_point: Vec2 = lens_radius*random_uniform_point_on_circle(
            &mut rand::thread_rng());
        let lens_point = self.origin + lens_point.x*self.right
            + lens_point.y*self.up;

        Ray {
            origin: lens_point,
            direction: (focus_point - lens_point).normalize()
        }
    }

    fn raster_dimensions(&self) -> (u16, u16) {
        (self.width_pixels, self.height_pixels)
    }
}

trait Integrator: Send + Sync {
    fn integrate_pixel(&self, x: u16, y: u16, camera: &dyn Camera, scene: &Scene) -> Vec3;
}

struct SingleSampleIntegrator {
}

impl Integrator for SingleSampleIntegrator {
    fn integrate_pixel(&self, x: u16, y: u16, camera: &dyn Camera, scene: &Scene) -> Vec3 {
        let ray = camera.get_ray(x as Scalar + 0.5, y as Scalar + 0.5);
        trace_ray_and_shade(scene, ray).unwrap_or(vec3(0.0, 0.0, 0.0))
    }
}

struct RandomSampleIntegrator {
    samples: u32
}

impl Integrator for RandomSampleIntegrator {
    fn integrate_pixel(&self, x: u16, y: u16, camera: &dyn Camera, scene: &Scene) -> Vec3 {
        let rng = &mut rand::thread_rng();
        let mut sum: Vec3 = vec3(0.0, 0.0, 0.0);
        for i in 0..self.samples {
            let offset_x: Scalar = rng.gen();
            let offset_y: Scalar = rng.gen();
            let ray = camera.get_ray(x as Scalar + offset_x, y as Scalar + offset_y);
            if let Some(rgb) = trace_ray_and_shade(scene, ray) {
                sum += rgb;
            }
        }
        sum / self.samples as Scalar
    }
}

fn render(
    integrator: &dyn Integrator,
    camera: &dyn Camera,
    scene: &Scene,
) -> HdrImage {
    let mut queue = Vec::new();
    let (raster_width, raster_height) = camera.raster_dimensions();
    let x_blocks = raster_width / BLOCK_SIZE;
    let y_blocks = raster_height / BLOCK_SIZE;
    for y in (0..y_blocks).rev() {
        for x in (0..x_blocks).rev() {
            queue.push((x, y));
        }
    }

    let mut block_queue = Mutex::new(queue);

    let mut image = HdrImage::new(raster_width, raster_height);
    let image_ref: Mutex<&mut HdrImage> = Mutex::new(&mut image);

    let render_block = |x_block: u16, y_block: u16| {
        let y_start = y_block * BLOCK_SIZE;
        let y_end = (y_start + BLOCK_SIZE).min(raster_height);
        let x_start = x_block * BLOCK_SIZE;
        let x_end = (x_start + BLOCK_SIZE).min(raster_width);
        for y in y_start..y_end {
            for x in x_start..x_end {
                let rgb = integrator.integrate_pixel(x, y, camera, scene);
                image_ref.lock().unwrap().put_rgb(x, y, rgb);
            }
        }
    };
    thread::scope(|s| {
        for thread_index in 0..THREADS {
            s.spawn(|_| loop {
                let (x_block, y_block) = {
                    if let Some(v) = block_queue.lock().unwrap().pop() {
                        v
                    } else {
                        return;
                    }
                };
                render_block(x_block, y_block);
            });
        }
    });

    image
}

// fn metaball_scene() -> Scene {
//     Scene {
//         numerical_objects: vec![Box::new(Metaball {
//             balls: vec![
//                 Sphere {
//                     origin: Point3 {
//                         x: 1.0,
//                         y: 0.0,
//                         z: 5.0,
//                     },
//                     radius: 1.0,
//                     debug_id: DebugId::NONE
//                 },
//                 Sphere {
//                     origin: Point3 {
//                         x: -1.0,
//                         y: 0.0,
//                         z: 5.0,
//                     },
//                     radius: 1.0,
//                     debug_id: DebugId::NONE
//                 },
//             ],
//             sharpness: 4.0,
//             debug_id: DebugId::NONE
//         })],
//         analytic_objects: vec![Box::new(InfinitePlane {
//             origin: Point3 {
//                 x: 0.0,
//                 y: -1.0,
//                 z: 0.0,
//             },
//             up: vec3(0.0, 1.0, 0.0),
//             forward: vec3(0.0, 0.0, 1.0),
//             right: vec3(1.0, 0.0, 0.0),
//             debug_id: DebugId::NONE
//         })],
//         sun: vec![Sun {
//             direction: vec3(1.0, 1.0, -1.0).normalize(),
//             sun_color: vec3(1.0, 1.0, 1.0),
//             sky_color: vec3(0.2, 0.2, 1.0),
//             angular_diameter: 1920.0 / 3600.0 * f64::consts::PI / 180.0,
//             debug_id: DebugId::NONE
//         }],
//         background: Some(Box::new(Sun {
//             direction: vec3(1.0, 1.0, -1.0).normalize(),
//             sun_color: vec3(1.0, 1.0, 1.0),
//             sky_color: vec3(0.2, 0.2, 1.0),
//             angular_diameter: 1920.0 / 3600.0 * f64::consts::PI / 180.0,
//             debug_id: DebugId::NONE
//         })),
//     }
// }

fn sphere_scene() -> Scene {

    let sun = Sun::new(
        orientation: Matrix3::look_at(-vec3(0.5, 0.5, 1.0).normalize(),
            vec3(0.0, 1.0, 0.0)),
        sun_color: vec3(25000.0, 25000.0, 25000.0),
        sky_color: vec3(0.2, 0.2, 1.0),
        angular_diameter: 1920.0 / 3600.0 * f64::consts::PI / 180.0,
        debug_id: DebugId::NONE,
    );
    Scene {
        analytic_objects: vec![
            Box::new(Sphere {
                origin: Point3 {
                    x: 0.0,
                    y: 0.0,
                    z: 5.0,
                },
                radius: 1.0,
                debug_id: DebugId::SPHERE,
            }),
            Box::new(InfinitePlane {
                origin: Point3 {
                    x: 0.0,
                    y: -1.0,
                    z: 0.0,
                },
                up: vec3(0.0, 1.0, 0.0),
                forward: vec3(0.0, 0.0, 1.0),
                right: vec3(1.0, 0.0, 0.0),
                debug_id: DebugId::NONE,
            }),
            Box::new(sun)
        ],
        numerical_objects: vec![],
        lights: vec![Box::new(sun)]
    }
}

fn main() {
    // let objects: Vec<Box<dyn Object>> = vec![
    //     Box::new(Sphere {
    //         origin: vec3(1.0, 0.0, 4.0),
    //         radius: 1.0,
    //     }),
    //     Box::new(Sphere {
    //         origin: vec3(-1.0, 0.0, 4.0),
    //         radius: 1.0
    //     })
    //     ];
    let scene = sphere_scene();
    // let camera = PinholeCamera {
    //     origin: Point3 {
    //         x: 0.0,
    //         y: 0.0,
    //         z: 0.0,
    //     },
    //     right: vec3(1.0, 0.0, 0.0),
    //     up: vec3(0.0, 1.0, 0.0),
    //     forward: vec3(0.0, 0.0, 1.0),
    //     width_pixels: 768,
    //     height_pixels: 512,
    //     focal_length: 0.05,
    //     width: 0.036,
    //     height: 0.024
    // };
    let camera = ThinLensCamera {
        origin: Point3 {
            x: 0.0,
            y: 0.0,
            z: 0.0,
        },
        right: vec3(1.0, 0.0, 0.0),
        up: vec3(0.0, 1.0, 0.0),
        forward: vec3(0.0, 0.0, 1.0),
        width_pixels: 768,
        height_pixels: 512,
        focal_length: 0.05,
        aperture: 2.8,
        width: 0.036,
        height: 0.024,
        focal_distance: 4.0, //4.5,
    };
    // let integrator = SingleSampleIntegrator{};
    let integrator = RandomSampleIntegrator{
        samples: 320*10000
    };
    let image = render(&integrator, &camera, &scene);
    let image = image.develop_rgb();
    image.save("output.png").unwrap();
}