use async_std::{fs::File, io, prelude::*, task, io::BufWriter};

fn lerp(fraction: f64, from: f64, to: f64) -> f64 {
    from + fraction * (to - from)
}

fn compute_pixel(re: f64, im: f64) -> u32 {
    let mut a = re;
    let mut b = im;
    let mut a2 = a * a;
    let mut b2 = b * b;
    let mut niters = 0;
    while a2 + b2 < 4.0 {
        b = 2.0 * a * b + im;
        a = a2 - b2 + re;
        a2 = a * a;
        b2 = b * b;
        niters += 1;
        if niters >= 10240 { break; }
    }
    niters
}

async fn compute_row(im: f64, re_min: f64, re_max: f64, width: usize) -> Vec<u32> {
    let mut row = Vec::new();
    row.resize(width, 0);
    for i in 0..width {
        row[i] = compute_pixel(lerp((i as f64) / (width as f64), re_min, re_max), im);
    }
    row
}

async fn write_image(fname: &str, rows: &Vec<Vec<u32>>) -> io::Result<()> {
    let width = rows[0].len();
    let height = rows.len();
    let maxval = rows.iter().map(|row| row.iter().max().unwrap()).max().unwrap();

    let mut out_vec: Vec<u8> = Vec::new();
    let mut writer = BufWriter::new(&mut out_vec);
    write!(writer, "P2\n{} {}\n{}\n", width, height, maxval).await?;
    for row in rows {
        for item in row {
            write!(writer, "{} ", item).await?;
        }
        writeln!(writer).await?;
    }

    let mut file = File::create(fname).await?;
    file.write(&out_vec).await?;
    file.flush().await?;

    Ok(())
}

async fn compute_brot(fname: &str, width: usize, height: usize, left_top: (f64, f64), right_bottom: (f64, f64)) -> io::Result<()> {
    println!("Spawning...");
    let mut join_handles = Vec::with_capacity(height);
    for i in 0..height {
        let im = lerp(i as f64 / height as f64, left_top.1, right_bottom.1);
        // Note: spawn to not run within the same thread
        join_handles.push(task::spawn(compute_row(im, left_top.0, right_bottom.0, width)));
    }

    println!("Awaiting...");
    let mut rows = Vec::with_capacity(height);
    for jh in join_handles {
        rows.push(jh.await);
    }

    println!("Writing...");
    write_image(fname, &rows).await
}

fn main() -> io::Result<()> {
    // Note: the following does not work:
    //    task::block_on(task::spawn(compute_brot("out.ppm", 1920, 1080, (-1.5, 1.0), (1.0, -1.0))))
    // You get a scary large futures error about Send.

    task::block_on(async {
        let width = 1920;
        let height = 1080;
        let left = -1.8;
        let right = 1.5;
        let height_im = height as f64 / width as f64 * (right - left);
        let top = height_im / 2.0;
        let bottom = -height_im / 2.0;
        compute_brot("out.ppm", width, height, (left, top), (right, bottom)).await
    })
}

// vim: set sw=4 ts=4 et: