bs-trace 0.3.0

use super::prelude::*;
use crate::linalg::vector::Vector;
use const_default::ConstDefault;
use std::io::Write;
use std::ops::{Mul, MulAssign};

/// A colour, with [r], [g] and [b] values.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Colour(pub Vector<f64, 3>);

/// An image with compile-time-known width and height, written `Image<W, H>`.
///
/// The bitmap data is represented as a grid of [Colour]s.
pub struct Image<const W: usize, const H: usize> {
    pub rows: Vec<Vec<Colour>>,
}

impl Colour {
    pub const WHITE: Colour = Colour::new(1.0, 1.0, 1.0);
    pub const BLACK: Colour = Colour::new(0.0, 0.0, 0.0);

    pub const fn new(r: f64, g: f64, b: f64) -> Self {
        debug_assert!(0.0 <= r && r <= 1.0);
        debug_assert!(0.0 <= g && g <= 1.0);
        debug_assert!(0.0 <= b && b <= 1.0);
        Self(Vector::new([r, g, b]))
    }

    /// Returns the red value for this colour.
    pub const fn r(&self) -> f64 {
        *self.0.x()
    }

    /// Returns the green value for this colour.
    pub const fn g(&self) -> f64 {
        *self.0.y()
    }

    /// Returns the blue value for this colour.
    pub const fn b(&self) -> f64 {
        *self.0.z()
    }

    pub fn map<F>(self, f: F) -> Self
    where
        F: Fn(f64) -> f64,
    {
        Colour(self.0.map(f))
    }

    pub fn to_pixel(&self) -> Pixel {
        // TODO these clamps may be unnecessary if I lock down new further
        let r = self.r().clamp(0.0, 1.0);
        let g = self.g().clamp(0.0, 1.0);
        let b = self.b().clamp(0.0, 1.0);

        let r = (r * 255.0) as u8;
        let g = (g * 255.0) as u8;
        let b = (b * 255.0) as u8;
        Pixel::new(r, g, b)
    }
}

impl ConstDefault for Colour {
    const DEFAULT: Self = Colour::new(0.0, 0.0, 0.0);
}

impl<const W: usize, const H: usize> Image<W, H> {
    /// Constructs an image from the given row data.
    ///
    /// This method will panic if there are not precisely `H` rows, or any of the rows do not contain precisely `W` colours.
    pub fn new(rows: Vec<Vec<Colour>>) -> Self {
        assert_eq!(rows.len(), H);
        for row in rows.iter() {
            assert_eq!(row.len(), W);
        }

        unsafe { Self::new_unchecked(rows) }
    }

    /// Constructs an image from the given row data.
    ///
    /// # Safety
    ///
    /// It is the responsibility of the caller to ensure that `rows` contains precisely `H`
    /// rows, and that each row contains precisely `W` colours.
    pub unsafe fn new_unchecked(rows: Vec<Vec<Colour>>) -> Self {
        Self { rows }
    }

    /// Generates an image using the given function in each pixel's (x,y) coordinates -- the
    /// top-left corner of the image is (0,0), x grows to the right in the positive direction,
    /// and y grows downwards in the positive direction.
    pub fn gen<F>(mut f: F) -> Self
    where
        F: FnMut(usize, usize) -> Colour,
    {
        let mut rows = Vec::<Vec<Colour>>::with_capacity(H);
        for y in 0..H {
            let mut row_data = Vec::<Colour>::with_capacity(W);
            for x in 0..W {
                row_data.push(f(x, y));
            }
            rows.push(row_data);
        }

        unsafe { Self::new_unchecked(rows) }
    }

    /// Converts this image to PPM format.
    pub fn to_ppm(&self) -> Ppm<W, H> {
        let mut pixel_rows = Vec::<Vec<Pixel>>::with_capacity(H);
        for colour_row in self.rows.iter() {
            let mut pixel_row = Vec::<Pixel>::with_capacity(W);
            for colour in colour_row.iter() {
                pixel_row.push(colour.to_pixel());
            }
            pixel_rows.push(pixel_row);
        }
        Ppm::new(pixel_rows)
    }

    pub fn write_as_png<OUT: Write>(&self, writer: OUT) {
        let mut encoder = png::Encoder::new(writer, W as u32, H as u32);
        encoder.set_color(png::ColorType::Rgb);
        encoder.set_depth(png::BitDepth::Eight);
        let mut png_writer = encoder.write_header().unwrap();
        let image_data: Vec<u8> = self
            .rows
            .iter()
            .flat_map(|row| row.iter().flat_map(|colour| colour.to_pixel().into_iter()))
            .collect();
        png_writer.write_image_data(&image_data).unwrap();
    }
}

impl From<Vector<f64, 3>> for Colour {
    fn from(v: Vector<f64, 3>) -> Self {
        Self::new(*v.x(), *v.y(), *v.z()) // TODO is this optimised to be equivalent to Self(v)?
    }
}

impl Mul<Colour> for Colour {
    type Output = Colour;

    fn mul(self, rhs: Colour) -> Self::Output {
        Colour(self.0 * rhs.0)
    }
}

impl MulAssign<Colour> for Colour {
    fn mul_assign(&mut self, rhs: Colour) {
        self.0 *= rhs.0;
    }
}