//#![feature(portable_simd)]

use clap::Parser;
use image::{ImageBuffer, Rgb};
use std::path::PathBuf;

pub mod transform;

#[derive(Parser)]
pub struct Args {
    #[clap(short, long, default_value_t = 3840)]
    pub width: u32,
    #[clap(short, long, default_value_t = 2160)]
    pub height: u32,

    #[clap(short, long, default_value_t = 100)]
    pub maxiter: usize,

    #[clap(short, long, default_value="out.png")]
    pub output: PathBuf,

    pub transform: Vec<String>
}

pub type Float = f64;
pub type Point = Complex<Float>;
use num::complex::Complex;
pub type Transform = [(Point, Point); 3];

trait ComplexExt<T> {
    fn norm2(self) -> T;
}

impl ComplexExt<Float> for Complex<Float> {
    fn norm2(self) -> Float {
        self.im * self.im + self.re * self.re
    }
}

fn transform_point(transform: &Transform, point: Point) -> (Point, Point) {
    let start = transform[0].0 * point.re + transform[1].0 * point.im + transform[2].0;
    let delta = transform[0].1 * point.re + transform[1].1 * point.im + transform[2].1;
    (start, delta)
}

pub struct Config {
    pub transform: Transform,
    pub width: u32,
    pub height: u32,
    pub maxiter: usize,
    // We represent escape as the square of its magnitude
}



impl Config {
    fn render_point(&self, (base, delta): (Point, Point)) -> Float {
        const ESCAPE: Float = 2.;
        const ESCAPE2: Float = ESCAPE * ESCAPE;

        let result  =
            std::iter::successors(Some(base), |i| Some(*i * *i + delta))
                .take(self.maxiter)
                .enumerate()
                .skip_while(|(_, i)| i.norm2() <= ESCAPE2)
                .next()
                .map_or(0.,
                        |(iter, val)|
                            iter as Float + 1. - val.norm().ln().ln() / ESCAPE.ln()
                            );
        result as Float
    }

    pub fn render_image(&self) -> ImageBuffer<Rgb<u8>, Vec<u8>> {
        let scale = 2.0 / std::cmp::min(self.width, self.height) as Float;

        ImageBuffer::<image::Rgb<u8>, _>::from_par_fn(self.width, self.height, |x, y| {
            let x = (x as Float) * scale - 1.;
            let y = (y as Float) * scale - 1.;
            let point = transform_point(&self.transform, Complex::new(x, y));
            map_palette(self.render_point(point) / self.maxiter as Float)
        })
    }
}

pub fn map_palette(val: Float) -> image::Rgb<u8> {
    let magnitude = Float::clamp(val * 255., 0., 255.) as u8;
    image::Rgb([magnitude, magnitude, magnitude])
}
pub fn main() -> anyhow::Result<()> {

    // let transform = [
    //     (Point::new(0., 0.), Point::new(1., 0.)),
    //     (Point::new(0., 0.), Point::new(0., 1.)),
    //     (Point::new(0., 0.), Point::new(0., 0.)),
    // ];

    let args = <Args as clap::Parser>::parse();

    let config = Config {
        transform: transform::parse_transforms(args.transform.iter())?,
        width: args.width,
        height: args.height,
        maxiter: args.maxiter,
    };

    let img = config.render_image();

    img.save(&args.output)?;
    Ok(())
}