rsrt/src/raytracer.rs

use rayon::iter::{IntoParallelIterator, ParallelIterator};

use crate::{
	builder::RaytracerBuilder,
	hittable::{Hittable, HittableList, Object},
	interval::WORLD,
	material::Material,
	ray::Ray,
	vec3::{Color, Point3, Vec3},
};

pub struct Raytracer {
	world: HittableList,

	samples_per_pixel: usize,
	max_depth: usize,
	pixel_samples_scale: f64,

	image_width: usize,
	image_height: usize,

	center: Point3,
	pixel00_loc: Point3,
	pixel_delta_u: Vec3,
	pixel_delta_v: Vec3,

	defocus_angle: f64,
	defocus_disk_u: Vec3,
	defocus_disk_v: Vec3,
}

impl From<RaytracerBuilder> for Raytracer {
	fn from(config: RaytracerBuilder) -> Self {
		let aspect_ratio: f64 = config.image_width as f64 / config.image_height as f64;
		let theta = config.vertical_fov.to_radians();
		let h = (theta / 2.0).tan();
		let viewport_height = 2.0 * h * config.focus_dist;
		let viewport_width = viewport_height as f64 * aspect_ratio;

		let center = config.look_from;

		let w = (config.look_from - config.look_at).unit_vector();
		let u = config.vup.cross(&w).unit_vector();
		let v = w.cross(&u);

		let viewport_u = viewport_width * u;
		let viewport_v = viewport_height * -v;

		let pixel_delta_u = viewport_u / config.image_width as f64;
		let pixel_delta_v = viewport_v / config.image_height as f64;

		let viewport_upper_left =
			center - (config.focus_dist * w) - (viewport_u / 2.0) - (viewport_v / 2.0);
		let pixel00_loc = viewport_upper_left + 0.5 * (pixel_delta_u + pixel_delta_v);

		let defocus_radius = config.focus_dist * (config.defocus_angle / 2.0).to_radians().tan();
		let defocus_disk_u = u * defocus_radius;
		let defocus_disk_v = v * defocus_radius;

		Self {
			world: config.world,
			samples_per_pixel: config.samples_per_pixel,
			max_depth: config.max_depth,
			pixel_samples_scale: 1.0 / config.samples_per_pixel as f64,
			image_width: config.image_width,
			image_height: config.image_height,
			center,
			pixel00_loc,
			pixel_delta_u,
			pixel_delta_v,
			defocus_angle: config.defocus_angle,
			defocus_disk_u,
			defocus_disk_v,
		}
	}
}

impl Raytracer {
	pub fn builder() -> RaytracerBuilder {
		RaytracerBuilder::default()
	}

	pub fn add(&mut self, object: Object) {
		self.world.add(object);
	}

	fn sample_square() -> Vec3 {
		Vec3::new(random_f64() - 0.5, random_f64() - 0.5, 0.0)
	}

	fn defocus_disk_sample(&self) -> Point3 {
		let p = Point3::random_in_unit_disk();
		self.center + (p[0] * self.defocus_disk_u) + (p[1] * self.defocus_disk_v)
	}

	fn get_ray(&self, i: usize, j: usize) -> Ray {
		let offset = Self::sample_square();
		let pixel_sample = self.pixel00_loc
			+ ((i as f64 + offset.x()) * self.pixel_delta_u)
			+ ((j as f64 + offset.y()) * self.pixel_delta_v);

		let origin = match self.defocus_angle <= 0.0 {
			true => self.center,
			false => self.defocus_disk_sample(),
		};
		let direction = pixel_sample - origin;

		Ray::new(origin, direction)
	}

	fn ray_color(&self, depth: usize, ray: &Ray) -> Color {
		if depth == 0 {
			return Color::default();
		}

		if let Some(hit_record) = self.world.hit(ray, &WORLD) {
			if let Some((attenuation, scattered)) = hit_record.mat.scatter(ray, &hit_record) {
				return attenuation * self.ray_color(depth - 1, &scattered);
			}
			return Color::default();
		}

		let unit_direction = ray.direction().unit_vector();
		let a = 0.5 * (unit_direction.y() + 1.0);
		((1.0 - a) * Color::new(1.0, 1.0, 1.0)) + (a * Color::new(0.5, 0.7, 1.0))
	}

	pub fn render_multithread_rayon(&self) -> Vec<Color> {
		(0..self.image_height)
			.into_par_iter()
			.flat_map(|y| (0..self.image_width).into_par_iter().map(move |x| (x, y)))
			.map(|(j, i)| {
				let mut pixel_color = Vec3::default();

				for _ in 0..self.samples_per_pixel {
					let ray = self.get_ray(i, j);
					pixel_color += self.ray_color(self.max_depth, &ray);
				}

				pixel_color * self.pixel_samples_scale
			})
			.collect()
	}

	pub fn render(&self) -> Vec<Color> {
		let mut out = Vec::with_capacity(self.image_height * self.image_width);

		for j in 0..self.image_height {
			tracing::debug!("Scanlines remaining: {}", self.image_height - j);
			for i in 0..self.image_width {
				let mut pixel_color = Vec3::default();

				for _ in 0..self.samples_per_pixel {
					let ray = self.get_ray(i, j);
					pixel_color += self.ray_color(self.max_depth, &ray);
				}

				out.push(pixel_color * self.pixel_samples_scale);
			}
		}

		out
	}
}

#[inline(always)]
pub fn random_f64() -> f64 {
	fastrand::f64()
}

#[inline(always)]
pub fn random_f64_range(min: f64, max: f64) -> f64 {
	min + (max - min) * random_f64()
}