use anyhow::{bail, Result};
use clap::Parser;
use image::{DynamicImage, Rgb, RgbImage};

#[derive(Debug, Parser)]
#[command()]
struct Args {
    #[arg(short, long, value_name = "FILE")]
    input: String,
    #[arg(short, long, value_name = "FILE")]
    output: String,

    #[arg(short = 'e', long, value_name = "NUM")]
    max_err: Option<f32>,
}

fn main() -> Result<()> {
    let args = Args::parse();
    let img = image::open(&args.input)?;
    let img = match img {
        DynamicImage::ImageRgb8(img) => img,
        _ => bail!("Unsupported type: {:?}", img),
    };
    println!("Opened {} ({}x{})", &args.input, img.width(), img.height());
    let mut output = RgbImage::new(128, 128);
    for x in 0..=127 {
        for y in 0..=127 {
            let pixel = img.get_pixel(x * img.width() / 128, y * img.height() / 128);
            output.put_pixel(x, y, *pixel);
        }
    }
    quadtree_quant(
        &mut output,
        args.max_err.unwrap_or(2000.0) * 128.0 * 128.0,
        [0, 0],
        128,
    );
    output.save(&args.output)?;
    println!("Wrote {}", &args.output);
    Ok(())
}

fn quadtree_quant(img: &mut RgbImage, max_err: f32, pos: [u32; 2], size: u32) {
    // Calculate average color
    let mut avg = [0.0, 0.0, 0.0];
    for y in pos[1]..pos[1] + size {
        for x in pos[0]..pos[0] + size {
            let pixel = img.get_pixel(x, y);
            for (i, component) in avg.iter_mut().enumerate() {
                *component += (pixel.0[i] as f32).powf(2.2);
            }
        }
    }
    for component in &mut avg {
        *component /= (size * size) as f32;
    }
    let avg = avg.map(|x| x.powf(0.4545).round() as u8);

    // Measure squared error
    let mut err = 0.0;
    for y in pos[1]..pos[1] + size {
        for x in pos[0]..pos[0] + size {
            let pixel = img.get_pixel(x, y);
            err += redmean_sq(avg, pixel.0);
        }
    }
    if err > max_err {
        let size = size / 2;
        let max_err = max_err / 4.0;
        let [x, y] = pos;
        quadtree_quant(img, max_err, [x, y], size);
        quadtree_quant(img, max_err, [x + size, y], size);
        quadtree_quant(img, max_err, [x, y + size], size);
        quadtree_quant(img, max_err, [x + size, y + size], size);
    } else {
        for y in pos[1]..pos[1] + size {
            for x in pos[0]..pos[0] + size {
                img.put_pixel(x, y, Rgb(avg));
            }
        }
    }
}

fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
    let mut ds: [f32; 3] = [0.0, 0.0, 0.0];
    for i in 0..2 {
        ds[i] = x[i] as f32 - y[i] as f32;
        ds[i] *= ds[i];
    }
    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
    (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2]
}