mandelbrot/src/mandelbrot.rs

use crate::{TEXTURE_HEIGHT, TEXTURE_WIDTH};
use num_complex::Complex;
use rayon::prelude::*;

const DEFAULT_ITERATION_LIMIT: u32 = 64;
const DEFAULT_BOUNDS: Bounds = Bounds::new((-2.5, 1.0), (-1.0, 1.0));

pub struct Coordinate(usize, usize);

impl Coordinate {
    pub fn new(x: usize, y: usize) -> Self {
        Self(x, y)
    }

    pub fn x(&self) -> usize {
        self.0
    }

    pub fn y(&self) -> usize {
        self.1
    }
}

impl Coordinate {
    pub fn into_complex(self, bounds: Bounds) -> Complex<f64> {
        Complex::new(
            Mandelbrot::scale_width(self.x(), bounds.x()),
            Mandelbrot::scale_height(self.y(), bounds.y()),
        )
    }
}

#[derive(Debug, Clone, Copy)]
pub struct Bounds {
    x: (f64, f64),
    y: (f64, f64),
}

impl Bounds {
    pub const fn new(x: (f64, f64), y: (f64, f64)) -> Self {
        Self { x, y }
    }
    pub fn x(&self) -> (f64, f64) {
        self.x
    }

    pub fn y(&self) -> (f64, f64) {
        self.y
    }
}

#[derive(Debug, Clone)]
pub struct Mandelbrot {
    frame_buffer: Box<[u8; (TEXTURE_WIDTH * TEXTURE_HEIGHT) * 4]>,
}

impl Default for Mandelbrot {
    fn default() -> Self {
        Mandelbrot {
            frame_buffer: Box::new([0; (TEXTURE_WIDTH * TEXTURE_HEIGHT) * 4]),
        }
    }
}

impl Mandelbrot {
    pub fn image(&mut self) -> &[u8] {
        let limit = DEFAULT_ITERATION_LIMIT as f64;

        self.frame_buffer
            .par_chunks_mut(4)
            .enumerate()
            .for_each(|(i, buf)| {
                let iters = Self::escape_time(i, DEFAULT_BOUNDS, DEFAULT_ITERATION_LIMIT);

                buf.copy_from_slice(&Self::default_colours(iters, limit))
            });

        self.frame_buffer.as_ref()
    }

    pub fn scaled_image(&mut self, bounds: Bounds, limit: u32) -> &[u8] {
        self.frame_buffer
            .par_chunks_mut(4)
            .enumerate()
            .for_each(|(i, buf)| {
                let iters = Self::escape_time(i, bounds, limit);
                let limit = limit as f64;

                buf.copy_from_slice(&Self::default_colours(iters, limit))
            });

        self.frame_buffer.as_ref()
    }

    fn olc_colours(iters: f64) -> [u8; 4] {
        let a = 0.1;

        let red = (0.5 * (a * iters).sin() + 0.5) * 255.0;
        let green = (0.5 * (a * iters + 2.094).sin() + 0.5) * 255.0;
        let blue = (0.5 * (a * iters + 4.188).sin() + 0.5) * 255.0;

        [red as u8, green as u8, blue as u8, 0xFF]
    }

    fn default_colours(iters: f64, limit: f64) -> [u8; 4] {
        let normal = iters / limit;

        let h = normal * 350.0;
        let v = if (iters - limit).abs() < f64::EPSILON {
            0.0
        } else {
            1.0
        };

        Self::hsv_to_rgb(h, 1.0, v)
    }

    fn escape_time(i: usize, bounds: Bounds, limit: u32) -> f64 {
        let point = Coordinate::new(i % TEXTURE_WIDTH, i / TEXTURE_WIDTH);
        let c = point.into_complex(bounds);

        Self::count_iterations(c, limit)
    }

    fn hsv_to_rgb(h: f64, s: f64, v: f64) -> [u8; 4] {
        let c = v * s;
        let x = c * (1.0 - (((h / 60.0) % 2.0) - 1.0).abs());
        let m = v - c;

        let (r_prime, g_prime, b_prime) = {
            if h < 60.0 {
                (c, x, 0.0)
            } else if h < 120.0 {
                (x, c, 0.0)
            } else if h < 180.0 {
                (0.0, c, x)
            } else if h < 240.0 {
                (0.0, x, c)
            } else if h < 300.0 {
                (x, 0.0, c)
            } else if h < 360.0 {
                (c, 0.0, x)
            } else {
                unreachable!()
            }
        };

        let r = ((r_prime + m) * 255.0) as u8;
        let g = ((g_prime + m) * 255.0) as u8;
        let b = ((b_prime + m) * 255.0) as u8;
        let a = 0xFF;

        [r, g, b, a]
    }

    fn count_iterations(c: Complex<f64>, limit: u32) -> f64 {
        let mut z: Complex<f64> = Complex::new(0.0, 0.0);
        let mut count: u32 = 0;

        loop {
            if count == limit {
                break;
            }

            if z.norm_sqr() > 4.0 {
                break;
            }

            z = (z * z) + c;
            count += 1;
        }

        // For a performance booth (with the consequence of a less smooth gradient)
        // just return count instead of doing all this math on it
        if count < limit {
            (count as f64 + 1.0) - z.norm().ln().ln() / 2f64.ln()
        } else {
            count as f64
        }
    }

    fn scale_width(x: usize, bounds: (f64, f64)) -> f64 {
        // const X_MIN: f64 = -2.5;
        // const X_MAX: f64 = 1.0;

        bounds.0 + ((bounds.1 - bounds.0) * (x as f64 - 0.0)) / TEXTURE_WIDTH as f64 - 0.0
    }

    fn scale_height(y: usize, bounds: (f64, f64)) -> f64 {
        // const Y_MIN: f64 = -1.0;
        // const Y_MAX: f64 = 1.0;

        bounds.0 + ((bounds.1 - bounds.0) * (y as f64 - 0.0)) / TEXTURE_HEIGHT as f64 - 0.0
    }
}