palette.rs

use std::cmp::{min, Ordering};
use std::collections::hash_map::RandomState;
use std::collections::{HashMap, HashSet};
use std::fmt::{Formatter, LowerHex, UpperHex};
use std::io::Write;
use std::path::Path;

use deltae::{Delta, LabValue, DE2000};
use image::{GenericImage, Pixel, Rgba};
use num_traits::ToPrimitive;

use crate::format::PaletteFormat;
use crate::utils::{new_image, BasicRgba};

pub type Palette = Vec<BasicRgba>;

pub fn palette(image: &impl GenericImage) -> anyhow::Result<Palette> {
	let mut colours = HashSet::new();
	for (_, _, pixel) in image.pixels() {
		let pixel = pixel.to_rgba();
		colours.insert(Rgba::from([
			pixel.0[0].to_u8().unwrap(),
			pixel.0[1].to_u8().unwrap(),
			pixel.0[2].to_u8().unwrap(),
			pixel.0[3].to_u8().unwrap(),
		]));
	}

	Ok(colours.iter().map(BasicRgba::from).collect())
}

struct HexStringValue(String);
trait HexString {
	fn as_hex_string(&self) -> HexStringValue;
}
impl HexString for Rgba<u8> {
	fn as_hex_string(&self) -> HexStringValue {
		HexStringValue(format!(
			"{:02X}{:02X}{:02X}{:02X}",
			self.0[0], self.0[1], self.0[2], self.0[3]
		))
	}
}
impl<T: HexString + Clone> HexString for &T {
	fn as_hex_string(&self) -> HexStringValue {
		(*self).clone().as_hex_string()
	}
}
impl UpperHex for HexStringValue {
	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
		f.write_str(&self.0.to_uppercase())
	}
}
impl LowerHex for HexStringValue {
	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
		f.write_str(&self.0.to_lowercase())
	}
}

pub fn write_palette<T: AsRef<Path>>(
	colours: Palette,
	format: PaletteFormat,
	outpath: T,
) -> anyhow::Result<()> {
	#[allow(clippy::redundant_clone)]
	let mut sorted = colours.clone();
	sorted.sort_by(|pa, pb| {
		let hue_a = pa.hue();
		let hue_b = pb.hue();

		println!("A: {} vs B: {}", hue_a, hue_b);

		if hue_a > hue_b {
			Ordering::Greater
		} else if hue_b > hue_a {
			Ordering::Less
		} else {
			Ordering::Equal
		}
	});

	match format {
		PaletteFormat::Png => {
			let num_colours = sorted.len();
			let image_width = min(16, num_colours);
			let image_height = if num_colours % 16 > 0 {
				num_colours / image_width + 1
			} else {
				num_colours / image_width
			};

			let mut out_image = new_image(image_width as u32, image_height as u32);
			for (idx, colour) in sorted.iter().enumerate() {
				out_image.put_pixel(
					(idx % image_width) as u32,
					(idx / image_width) as u32,
					Rgba::from(colour),
				);
			}

			out_image.save(outpath)?;
		}
		PaletteFormat::Txt => {
			let mut file = std::fs::File::create(outpath)?;
			for colour in sorted.iter() {
				let line = format!("#{:X}\n", colour);
				file.write_all(line.as_bytes())?;
			}
		}
	}

	Ok(())
}

pub type ColourMapping = HashMap<BasicRgba, BasicRgba>;

pub fn calculate_mapping(from: &Palette, to: &Palette) -> ColourMapping {
	let colour_labs = Vec::from_iter(to.iter().map(LabValue::from));

	let to_palette_vectors: HashMap<usize, &BasicRgba, RandomState> =
		HashMap::from_iter(to.iter().enumerate());
	let mut out_map: ColourMapping = HashMap::with_capacity(from.len());

	for colour in from {
		let closest = to_palette_vectors
			.keys()
			.fold(None, |lowest, idx| match lowest {
				Some(num) => {
					let current = colour_labs[*idx];
					let previous: LabValue = colour_labs[num];

					if colour.delta(current, DE2000) < colour.delta(previous, DE2000) {
						Some(*idx)
					} else {
						Some(num)
					}
				}
				None => Some(*idx),
			});

		match closest {
			Some(idx) => match to_palette_vectors.get(&idx) {
				Some(col) => {
					out_map.insert(*colour, **col);
				}
				None => {
					println!("No matching vec for {} with col {:?}", idx, &colour);

					out_map.insert(
						*colour,
						BasicRgba {
							r: 0,
							g: 0,
							b: 0,
							a: 0,
						},
					);
				}
			},
			None => {
				println!("No closest for {:?}", &colour);
				out_map.insert(
					*colour,
					BasicRgba {
						r: 0,
						g: 0,
						b: 0,
						a: 0,
					},
				);
			}
		}
	}

	out_map
}