crystal_ball 0.3.0

use std::sync::{Arc, Mutex};
use std::time::Instant;

use indicatif::{ProgressBar, ProgressStyle};
use nanorand::tls::TlsWyRand;
use nanorand::tls_rng;
use rayon::prelude::*;

use crate::color::{Color, Image, ImageTile, Interpolation};
use crate::materials::compute_color;
use crate::math::{Hit, Ray};
use crate::rendering::{PrecalculatedCamera, Scene};
use crate::shapes::{Object, Sphere, BVH};
use crate::util::random_float;

/// The render engine.
///
/// Rendering is performed by splitting the image into tiles (small subsets of the image).
/// These tiles are then rendered in parallel to increase performance.
/// The best tile size varies depending on your scene,
/// therefore Crystal Ball uses default values,
/// that have proven to be a good compromise for many different scenes.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct RenderEngine {
    /// The sample count per pixel.
    ///
    /// Defaults to `64`.
    pub samples: usize,
    /// The maximum amount of bounces before a [`Ray`] is terminated.
    ///
    /// Defaults to `4`.
    pub max_bounces: usize,
    /// The width of a tile.
    ///
    /// If it's greater than the [`Camera`](crate::rendering::Camera)'s width, the width is clamped.
    /// If the [`Camera`](crate::rendering::Camera)'s width is not a multiple of the tile width,
    /// the last tile will be truncated.
    ///
    /// Defaults to `16`.
    pub tile_size_x: u32,
    /// The height of a tile.
    ///
    /// If it's greater than the [`Camera`](crate::rendering::Camera)'s height, the height is clamped.
    /// If the [`Camera`](crate::rendering::Camera)'s height is not a multiple of the tile height,
    /// the last tile will be truncated.
    ///
    /// Defaults to `16`.
    pub tile_size_y: u32,
}

impl Default for RenderEngine {
    fn default() -> Self {
        Self {
            samples: 64,
            max_bounces: 4,
            tile_size_x: 16,
            tile_size_y: 16,
        }
    }
}

impl RenderEngine {
    /// Return the closest hit of a ray with objects of the scene, if any.
    fn get_closest_hit<'a>(&self, bvh: &'a BVH<Object>, ray: Ray) -> Option<(Hit, &'a Object)> {
        bvh.intersects(ray)
    }

    /// Trace a ray and return the resulting color.
    fn trace_tray(
        &self,
        bvh: &BVH<Object>,
        scene: &Scene,
        ray: Ray,
        bounces: usize,
        rng: &mut TlsWyRand,
    ) -> Color {
        if bounces > self.max_bounces {
            return Color::default();
        }

        let closest_hit = self.get_closest_hit(bvh, ray);

        match closest_hit {
            None => {
                let uv = Sphere::new().uv_map(
                    (scene.background_transform.mat4 * ray.direction)
                        .normalize()
                        .to_point3(),
                );

                compute_color(
                    scene.background_color,
                    &scene.background_texture,
                    scene.background_strength,
                    uv,
                )
            }
            Some((hit, object)) => {
                let (diffuse_color, emissive_color) = object.material.get_color(hit.uv);

                match object.material.next_ray(ray, hit, rng) {
                    None => emissive_color,
                    Some(ray) => {
                        self.trace_tray(bvh, scene, ray, bounces + 1, rng) * diffuse_color
                            + emissive_color
                    }
                }
            }
        }
    }

    /// Render the given scene to an image.
    ///
    /// Before rendering this generates a [`BVH`] for the whole scene.
    pub fn render(&self, scene: &Scene) -> Image {
        let camera: PrecalculatedCamera = scene.camera.into();

        let (width, height) = camera.dimensions();

        println!(
            "Rendering {} object(s) with {} samples ({} bounces) to a {}x{} image.",
            scene.objects.len(),
            self.samples,
            self.max_bounces,
            width,
            height
        );

        let bvh = BVH::init(4, scene.objects.clone());

        let tile_size_x = self.tile_size_x.min(width);
        let tile_size_y = self.tile_size_y.min(height);

        let mut tiles = vec![];
        for x in 0..(width as f32 / tile_size_x as f32).ceil() as u32 {
            for y in 0..(height as f32 / tile_size_y as f32).ceil() as u32 {
                let tile_width = if (x + 1) * tile_size_x > width {
                    width - tile_size_x * x
                } else {
                    tile_size_x
                };
                let tile_height = if (y + 1) * tile_size_y >= height {
                    height - tile_size_y * y
                } else {
                    tile_size_y
                };

                tiles.push(ImageTile::new(
                    x * tile_size_x,
                    y * tile_size_y,
                    Image::new(tile_width, tile_height, Interpolation::Closest),
                ));
            }
        }

        let progress_bar = ProgressBar::new(tiles.len() as u64);
        progress_bar.set_style(
            ProgressStyle::with_template(
                "[{elapsed_precise}] [{bar:60}] {percent:>5}% (ETA: ~{eta_precise})\n{msg}",
            )
            .unwrap()
            .progress_chars("=> "),
        );
        progress_bar.set_position(0);
        let progress_bar = Arc::new(Mutex::new(progress_bar));

        let samples_inverse = 1.0 / self.samples as f32;

        let start_time = Instant::now();

        tiles.par_iter_mut().for_each(|tile| {
            let tile_x = tile.x;
            let tile_y = tile.y;

            tile.image
                .pixels
                .iter_mut()
                .enumerate()
                .for_each(|(i, pixel)| {
                    let x = (i as u32 % self.tile_size_x) + tile_x;
                    let y = (i as u32 / self.tile_size_y) + tile_y;

                    let mut rng = tls_rng();
                    let mut color = Color::default();

                    for _ in 0..self.samples {
                        let ray = camera.get_ray(
                            x as f64 + random_float(&mut rng, 0.0, 1.0),
                            height as f64 - y as f64 + random_float(&mut rng, 0.0, 1.0),
                            &mut rng,
                        );

                        color += self.trace_tray(&bvh, scene, ray, 0, &mut rng);
                    }

                    *pixel = color * samples_inverse;
                    *pixel = Color::new(
                        pixel.r.min(1.0).max(0.0),
                        pixel.g.min(1.0).max(0.0),
                        pixel.b.min(1.0).max(0.0),
                    );
                });

            progress_bar.lock().unwrap().inc(1);
        });

        progress_bar.lock().unwrap().finish_with_message(format!(
            "Rendering finished after {:.2?}",
            Instant::now() - start_time
        ));

        Image::from_tiles(width, height, &tiles, Interpolation::Closest)
    }
}