//#![feature(portable_simd)]

use std::path::PathBuf;
use clap::{Arg, Parser};
use image::{GrayImage, ImageBuffer, Luma};
use image::buffer::ConvertBuffer;
use rayon::iter::ParallelIterator;

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>
}

type Float = f32;
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)) -> f32 {
        const ESCAPE: Float = 2.;
        const ESCAPE2: Float = ESCAPE * ESCAPE;

        let mut base = base;
        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 f32
    }
}

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 scale = 2.0 / std::cmp::min(config.width, config.height) as Float;

    let mut img = ImageBuffer::<image::Luma<f32>, _>::from_par_fn(config.width, config.height, |x, y| {
        let x = (x as Float) * scale - 1.;
        let y = (y as Float) * scale - 1.;
        let point = transform_point(&config.transform, Complex::new(x, y));
        image::Luma([config.render_point(point)])
    });
    let img_norm = img.pixels()
        .map(|p: &Luma<f32>| p.0[0])
        .reduce(f32::max)
        .unwrap_or(1.);
    img.par_pixels_mut().for_each(|Luma([p])| *p /= img_norm);
    let img: GrayImage = img.convert();
    img.save(&args.output)?;
    Ok(())
}