use std::{env, fs::File, io::Read};

use image::{imageops::FilterType, DynamicImage, GenericImage, GenericImageView, Pixel, Rgb};

fn main() {
    let mut args: Vec<String> = env::args().skip(1).collect();
    if args.is_empty() {
        println!("no input file specified");
        return;
    }
    dbg!(&args);

    let inpath = args.pop().unwrap();
    let mut data = Vec::new();
    File::open(inpath).unwrap().read_to_end(&mut data).unwrap();
    let image = image::load_from_memory(&data).unwrap();
    println!("loaded");

    let (w, h) = image.dimensions();
    let image = image.resize(w / 8, h / 8, FilterType::CatmullRom);
    image.save("scaled.png").unwrap();
    println!("resized");

    const COLOR_COUNT: usize = 6;

    let quantized_image = quantize_image(image.clone(), COLOR_COUNT);
    quantized_image.save("q.png").unwrap();

    let image = dither(image, COLOR_COUNT);
    image.save("out.png").unwrap();
    println!("saved");
}

fn generate_palette(image: &DynamicImage, count: usize) -> Vec<Rgb<u8>> {
    assert!(count > 0);

    let mut buckets: Vec<Vec<Rgb<u8>>> = vec![image
        .pixels()
        .map(|(_x, _y, a)| a.to_rgb()) //
        .collect()];

    struct Spread {
        range: u8,
        channel: usize,
        mid_point: u8,
    }
    // divide buckets count-1 times
    for _i in 0..(count - 1) {
        // spread amount and channel for each bucket
        let mut spreads: Vec<Spread> = Vec::new();
        // calculate where each bucket would do its division if it's chosen
        // todo: only calculate this when the bucket is created/modified
        for bucket in &buckets {
            let mut min = [255; 3];
            let mut max = [0; 3];
            for pixel in bucket {
                for channel in 0..3 {
                    min[channel] = min[channel].min(pixel.0[channel]);
                    max[channel] = max[channel].max(pixel.0[channel]);
                }
            }
            let range_r = max[0] - min[0];
            let range_g = max[1] - min[1];
            let range_b = max[2] - min[2];
            let mut widest_channel = 0;
            let mut widest_amount = range_r;
            let mut widest_middle = min[0] + range_r / 2;
            if range_g > widest_amount {
                widest_amount = range_g;
                widest_channel = 1;
                widest_middle = min[1] + range_g / 2;
            }
            if range_b > widest_amount {
                widest_amount = range_b;
                widest_channel = 2;
                widest_middle = min[2] + range_b / 2;
            }
            spreads.push(Spread {
                range: widest_amount,
                channel: widest_channel,
                mid_point: widest_middle,
            });
        }

        let mut most_spread_bucket = 0;
        let mut highest_spread = 0;
        for (i, spread) in spreads.iter().enumerate() {
            if spread.range > highest_spread {
                most_spread_bucket = i;
                highest_spread = spread.range;
            }
        }

        // divide the most spread bucket
        let bucket = &mut buckets[most_spread_bucket];
        if bucket.len() < 2 {
            continue;
        }

        let channel = spreads[most_spread_bucket].channel;
        bucket.sort_unstable_by_key(|pixel| pixel.0[channel]);
        let mid_value = spreads[most_spread_bucket].mid_point;
        let pop_split = bucket.len() / 2;
        let pos_split = bucket
            .iter()
            .position(|c| c.0[channel] >= mid_value)
            .unwrap();
        let split_index = if spreads[most_spread_bucket].range > 150 {
            pos_split
        } else {
            pop_split
        };
        let new_bucket = bucket[split_index..].to_owned();
        bucket.truncate(split_index);
        buckets.push(new_bucket);
    }
    let mut colors = Vec::new();
    for bucket in buckets {
        let mut avg = [0u128; 3];
        for p in &bucket {
            for channel in 0..3 {
                avg[channel] += p.0[channel] as u128;
            }
        }
        let num = bucket.len() as u128;
        colors.push(Rgb([
            (avg[0] / num) as u8,
            (avg[1] / num) as u8,
            (avg[2] / num) as u8,
        ]))
    }
    colors
}

fn color_dist(a: Rgb<u8>, b: Rgb<u8>) -> f32 {
    let r = a.0[0] as f32 - b.0[0] as f32;
    let g = a.0[1] as f32 - b.0[1] as f32;
    let b = a.0[2] as f32 - b.0[2] as f32;
    (r * r + g * g + b * b).sqrt()
}

fn quantize_image(input: DynamicImage, count: usize) -> DynamicImage {
    let colors = generate_palette(&input, count);
    let mut out = input.clone();

    for (x, y, color) in input.pixels() {
        let color = color.to_rgb();
        let mut closest_index = 0;
        let mut closest_dist = color_dist(color, colors[0]);
        for (i, &c) in colors.iter().enumerate() {
            let d = color_dist(color, c);
            if d < closest_dist {
                closest_dist = d;
                closest_index = i;
            }
        }
        out.put_pixel(x, y, colors[closest_index].to_rgba());
    }

    out
}

const BAYER_4X4: [[u8; 4]; 4] = [[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]];

fn bayer_f(x: usize, y: usize) -> f32 {
    (BAYER_4X4[y][x] as f32 / 16.) - (15. / 16.) / 2.
}

fn dither(input: DynamicImage, count: usize) -> DynamicImage {
    let mut colors = generate_palette(&input, count);
    // let mut colors = vec![
    //     Rgb([177, 138, 129]),
    //     Rgb([155, 68, 52]),
    //     Rgb([234, 130, 70]),
    //     Rgb([110, 75, 72]),
    //     Rgb([60, 43, 41]),
    //     Rgb([0, 0, 0]),
    // ];
    let (w, h) = input.dimensions();
    let mut out = DynamicImage::new(w, h, input.color());

    for (x, y, color) in input.pixels() {
        let real_col = color.to_rgb();
        colors.sort_unstable_by_key(|&c| (color_dist(c, real_col) * 100.) as u32);

        let best_col = colors[0];
        let second_col = colors[1];

        let best_dist = color_dist(real_col, best_col);
        let second_dist = color_dist(real_col, second_col);

        let ratio = second_dist / (second_dist + best_dist) * 2. - 1.;

        let bx = x as usize % 4;
        let by = y as usize % 4;
        let bayer = bayer_f(bx, by);

        let ratio = ratio + bayer;
        let out_col = if ratio > 0. { best_col } else { second_col };

        out.put_pixel(x, y, out_col.to_rgba());
    }
    out
}