use image::{ImageBuffer, Rgb, RgbImage};
use num_complex::Complex;
use rayon::prelude::*;

const MAX_ITERATIONS: u16 = 256;
const IMG_WIDTH: usize = 1280;
const IMG_HEIGHT: usize = 720;

#[derive(Debug, Clone)]
pub struct Mandelbrot {
    iterations: Vec<u16>,
}

enum Colour {
    White,
    Black,
    Gray,
}

impl Mandelbrot {
    pub fn new() -> Self {
        Mandelbrot {
            iterations: vec![0; IMG_WIDTH * IMG_HEIGHT],
        }
    }

    pub fn generate_image(&mut self) -> Vec<u8> {
        // Create the image
        // Pixel Format can be Rgba8UnormSrgb
        self.escape_time();

        let mut texture_buffer = vec![0; (IMG_WIDTH * IMG_HEIGHT) * 4];

        for i in 0..texture_buffer.len() {
            if i % 4 == 0 {
                let iter_i = i / 4;
                match Self::find_colour(self.iterations[iter_i]) {
                    Colour::White => {
                        texture_buffer[i] = 0x00;
                        texture_buffer[i + 1] = 0x00;
                        texture_buffer[i + 2] = 0x00;
                        texture_buffer[i + 3] = 0xFF;
                    }
                    Colour::Black => {
                        texture_buffer[i] = 0xFF;
                        texture_buffer[i + 1] = 0xFF;
                        texture_buffer[i + 2] = 0xFF;
                        texture_buffer[i + 3] = 0xFF;
                    }
                    Colour::Gray => {
                        texture_buffer[i] = 0xAA;
                        texture_buffer[i + 1] = 0xAA;
                        texture_buffer[i + 2] = 0xAA;
                        texture_buffer[i + 3] = 0xFF;
                    }
                }
            }
        }

        texture_buffer
    }

    pub fn generate_scaled_image(&mut self, x_bounds: (f64, f64), y_bounds: (f64, f64)) -> Vec<u8> {
        let mut texture_buffer = vec![0; (IMG_WIDTH * IMG_HEIGHT) * 4];
        self.iterations
            .par_iter_mut()
            .enumerate()
            .for_each(|(i, value)| {
                let px = i % IMG_WIDTH;
                let py = i / IMG_WIDTH;

                let c = Self::coords_to_complex(px, py, x_bounds, y_bounds);

                *value = Self::calc_num_iterations(c);
            });

        for i in 0..texture_buffer.len() {
            if i % 4 == 0 {
                let iter_i = i / 4;
                match Self::find_colour(self.iterations[iter_i]) {
                    Colour::White => {
                        texture_buffer[i] = 0x00;
                        texture_buffer[i + 1] = 0x00;
                        texture_buffer[i + 2] = 0x00;
                        texture_buffer[i + 3] = 0xFF;
                    }
                    Colour::Black => {
                        texture_buffer[i] = 0xFF;
                        texture_buffer[i + 1] = 0xFF;
                        texture_buffer[i + 2] = 0xFF;
                        texture_buffer[i + 3] = 0xFF;
                    }
                    Colour::Gray => {
                        texture_buffer[i] = 0xAA;
                        texture_buffer[i + 1] = 0xAA;
                        texture_buffer[i + 2] = 0xAA;
                        texture_buffer[i + 3] = 0xFF;
                    }
                }
            }
        }

        texture_buffer
    }

    fn escape_time(&mut self) {
        self.iterations
            .par_iter_mut()
            .enumerate()
            .for_each(|(i, value)| {
                let px = i % IMG_WIDTH;
                let py = i / IMG_WIDTH;

                let c = Self::coords_to_complex(px, py, (-2.5, 1.0), (-1.0, 1.0));

                *value = Self::calc_num_iterations(c);
            });
    }

    pub fn escape_time_image() {
        const IMG_WIDTH: usize = 1280;
        const IMG_HEIGHT: usize = 720;

        let mut img: RgbImage = ImageBuffer::new(IMG_WIDTH as u32, IMG_HEIGHT as u32);
        let mut mandelbrot: Vec<u16> = vec![0; IMG_HEIGHT * IMG_WIDTH];

        mandelbrot
            .par_iter_mut()
            .enumerate()
            .for_each(|(i, value)| {
                let px = i % IMG_WIDTH;
                let py = i / IMG_WIDTH;

                let c = Self::coords_to_complex(px, py, (-2.5, 1.0), (-1.0, 1.0));

                *value = Self::calc_num_iterations(c);
            });

        for (px, py, pixel) in img.enumerate_pixels_mut() {
            let i = px as usize + IMG_WIDTH * py as usize;
            let num_iterations = mandelbrot[i];

            match Self::find_colour(num_iterations) {
                Colour::White => *pixel = Rgb([0x00, 0x00, 0x00]),
                Colour::Black => *pixel = Rgb([0xFF, 0xFF, 0xFF]),
                Colour::Gray => *pixel = Rgb([0xAA, 0xAA, 0xAA]),
            }
        }

        img.save("assets/test.png").unwrap();
    }

    fn find_colour(iteration: u16) -> Colour {
        match iteration {
            MAX_ITERATIONS => Colour::Black,
            other => match other % 2 {
                0 => Colour::White,
                1 => Colour::Gray,
                _ => unreachable!(),
            },
        }
    }

    fn calc_num_iterations(c: Complex<f64>) -> u16 {
        let mut z: Complex<f64> = Complex::new(0.0, 0.0);
        let mut num_iterations: u16 = 0;

        for i in 0..MAX_ITERATIONS {
            if z.norm_sqr() > 4.0 {
                num_iterations = i;
                break;
            }

            z = (z * z) + c;
        }

        num_iterations
    }

    fn coords_to_complex(
        px: usize,
        py: usize,
        x_bounds: (f64, f64),
        y_bounds: (f64, f64),
    ) -> Complex<f64> {
        Complex::new(
            Self::scale_width(px, x_bounds),
            Self::scale_height(py, y_bounds),
        )
    }

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

        x_bounds.0 + ((x_bounds.1 - x_bounds.0) * (px as f64 - 0.0)) / IMG_WIDTH as f64 - 0.0
    }

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

        y_bounds.0 + ((y_bounds.1 - y_bounds.0) * (py as f64 - 0.0)) / IMG_HEIGHT as f64 - 0.0
    }
}