use std::collections::BinaryHeap;

use empfindung::cie00;
use image::RgbImage;
use lab::Lab;
use ordered_float::NotNan;
use rand::Rng;

pub fn euclid_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
    let mut result = 0.0;
    let weights = [1.641, 5.764, 2.088];
    for (i, weight) in weights.iter().enumerate() {
        let delta = x[i] as f32 - y[i] as f32;
        result += weight * delta * delta;
    }
    result
}

pub fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
    let mut ds: [f32; 3] = [0.0; 3];
    for (i, d) in ds.iter_mut().enumerate() {
        let delta = x[i] as f32 - y[i] as f32;
        *d = delta * delta;
    }
    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
    let result = (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2];
    // normalize value to the range 0.0 - 1.0
    result / 584971.0
}

fn rgb_to_ycocg(x: [u8; 3]) -> [f32; 3] {
    let [r, g, b] = x.map(|x| x as f32);
    let co = r - b;
    let tmp = b + co/2.0;
    let cg = g - tmp;
    let y = tmp + cg/2.0;
    [y, co, cg]
}

pub fn ycocg_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
    let x = rgb_to_ycocg(x);
    let y = rgb_to_ycocg(y);
    let mut result = 0.0;
    let weights = [15.03, 1.982, 7.702];
    for (i, weight) in weights.iter().enumerate() {
        let delta = (x[i] as f32 - y[i] as f32)/255.0;
        result += weight * delta * delta;
    }
    result
}

type Rgb = [u8; 3];

const PICO8_PALETTE: [Rgb; 32] = [
    // 0-15: standard palette
    [0x00, 0x00, 0x00],
    [0x1d, 0x2b, 0x53],
    [0x7e, 0x25, 0x53],
    [0x00, 0x87, 0x51],
    [0xab, 0x52, 0x36],
    [0x5f, 0x57, 0x4f],
    [0xc2, 0xc3, 0xc7],
    [0xff, 0xf1, 0xe8],
    [0xff, 0x00, 0x4d],
    [0xff, 0xa3, 0x00],
    [0xff, 0xec, 0x27],
    [0x00, 0xe4, 0x36],
    [0x29, 0xad, 0xff],
    [0x83, 0x76, 0x9c],
    [0xff, 0x77, 0xa8],
    [0xff, 0xcc, 0xaa],
    // 16-31: secret palette
    // In PICO-8, these indexes start at 128
    [0x29, 0x18, 0x14],
    [0x11, 0x1d, 0x35],
    [0x42, 0x21, 0x36],
    [0x12, 0x53, 0x59],
    [0x74, 0x2f, 0x29],
    [0x49, 0x33, 0x3b],
    [0xa2, 0x88, 0x79],
    [0xf3, 0xef, 0x7d],
    [0xbe, 0x12, 0x50],
    [0xff, 0x6c, 0x24],
    [0xa8, 0xe7, 0x2e],
    [0x00, 0xb5, 0x43],
    [0x06, 0x5a, 0xb5],
    [0x75, 0x46, 0x65],
    [0xff, 0x6e, 0x59],
    [0xff, 0x9d, 0x81],
];

#[derive(Clone, Copy)]
pub struct PicoColor(u8);

impl PicoColor {
    pub fn to_rgb(self) -> Rgb {
        let idx = ((self.0 & 0x80) >> 3) | (self.0 & 0x0F);
        PICO8_PALETTE[idx as usize]
    }
}

pub struct PicoImage {
    pub palette: [PicoColor; 16],
    pub width: u32,
    pub height: u32,
    pub data: Vec<u8>,
}

impl PicoImage {
    pub fn new(width: u32, height: u32) -> Self {
        let mut palette = [PicoColor(0); 16];
        for (i, color) in palette.iter_mut().enumerate() {
            *color = PicoColor(i as u8);
        }
        Self {
            palette,
            width,
            height,
            data: vec![0; (width * height) as usize],
        }
    }

    pub fn to_indexed(img: &RgbImage) -> Self {
        // Brute-force nearest-color conversion
        // TODO begin using this everywhere
        let mut result = Self::new(img.width(), img.height());
        let mut discrepancies: Vec<(f32, (Rgb, Rgb), f32)> = vec![];
        for y in 0..img.height() {
            let offset = (y * img.width()) as usize;
            for x in 0..img.width() {
                let original = img.get_pixel(x, y).0;
                let best = (0..16)
                    .min_by_key(|&i| {
                        let candidate = PicoColor(i).to_rgb();
                        // let diff = redmean_sq(original, candidate);
                        // let diff = euclid_sq(original, candidate);
                        let diff = ycocg_sq(original, candidate);
                        // let diff =
                        //     cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&candidate));
                        NotNan::new(diff).unwrap()
                    })
                    .unwrap();

                // Log places where our color formula disagrees with CIEDE2000
                let best_rgb = PicoColor(best).to_rgb();
                let ciede2000 = cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&best_rgb));
                let diff = 24.36 * ycocg_sq(original, best_rgb).sqrt();
                let discrepancy = (
                    ciede2000,
                    (original, best_rgb),
                    diff
                );
                if discrepancies.len() < 10 {
                    discrepancies.push(discrepancy);
                } else {
                    let idx = rand::thread_rng().gen_range(0..discrepancies.len());
                    if discrepancies[idx].0 < discrepancy.0 {
                        discrepancies[idx] = discrepancy;
                    }
                }

                // if x % 43 == 22 && y % 43 == 22 {
                //     let best = PicoColor(best).to_rgb();
                //     let ciede2000 = cie00::diff(Lab::from_rgb(&original), Lab::from_rgb(&best));
                //     let redmean = redmean_sq(original, best);
                //     dbg!((original, best));
                //     dbg!((redmean, ciede2000 * ciede2000 / 100_00.0));
                // }
                result.data[offset + x as usize] = best;
            }
        }
        for d in discrepancies {
            dbg!(d);
        }
        result
    }

    pub fn get(&self, x: u32, y: u32) -> u8 {
        self.data[(y * self.width + x) as usize]
    }

    pub fn to_rgb(&self) -> RgbImage {
        let mut result = RgbImage::new(self.width, self.height);
        for y in 0..self.height {
            for x in 0..self.width {
                let index = self.get(x, y);
                let color = self.palette[index as usize];
                result.put_pixel(x, y, image::Rgb(color.to_rgb()));
            }
        }
        result
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_to_rgb() {
        let to_rgb = |i| PicoColor(i).to_rgb();

        // Regular palette
        assert_eq!(to_rgb(0), [0, 0, 0]);
        assert_eq!(to_rgb(1), PICO8_PALETTE[1]);
        assert_eq!(to_rgb(2), to_rgb(0x12));
        assert_eq!(to_rgb(3), to_rgb(0x73));

        // Secret palette
        assert_eq!(to_rgb(128), PICO8_PALETTE[16]);
        assert_eq!(to_rgb(128 + 15), PICO8_PALETTE[31]);
        assert_eq!(to_rgb(128 + 16), PICO8_PALETTE[16]);
    }

    #[test]
    fn test_quantization_round_trip() {
        for i in 0..16 {
            let mut img = RgbImage::new(1, 1);
            let input_color = PICO8_PALETTE[i];
            img.put_pixel(0, 0, image::Rgb(input_color));
            let indexed = PicoImage::to_indexed(&img);
            let img = indexed.to_rgb();
            let output_color = img.get_pixel(0, 0).0;
            assert_eq!(input_color, output_color);
        }
    }
}
