palette.rs

use clap::Parser;
use std::cmp::{min, Ordering};
use std::collections::hash_map::RandomState;
use std::collections::{HashMap, HashSet};

use anyhow::Error;
use deltae::{Delta, LabValue, DE2000};
use image::{GenericImage, Pixel, Rgba};
use num_traits::ToPrimitive;
use serde::{Deserialize, Serialize};
use std::io::Write;

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

pub type PixelPalette = Vec<BasicRgba>;
pub type ColourMapping = HashMap<BasicRgba, BasicRgba>;

pub fn sort_by_hue(palette: &mut PixelPalette) {
	palette.sort_by(|pa, pb| {
		let hue_a = pa.hue();
		let hue_b = pb.hue();

		if hue_a.is_nan() && hue_b.is_nan() {
			Ordering::Equal
		} else if hue_a.is_nan() {
			Ordering::Less
		} else if hue_b.is_nan() || hue_a > hue_b {
			Ordering::Greater
		} else if hue_b > hue_a {
			Ordering::Less
		} else {
			Ordering::Equal
		}
	});
}

/// Create a palette file containing every distinct colour from the input image
#[derive(Parser, Clone, Serialize, Deserialize, Debug)]
#[clap(author, version = "0.9.0")]
pub struct Palette {
	/// The path to the image file
	#[serde(default)]
	pub input: String,
	/// The path to write the palette data
	#[serde(default)]
	pub output: String,

	/// The format for the palette file. PNG will output an image with colours, while TXT will
	/// output a text file with one hex value per line
	#[clap(short, long, value_enum, default_value_t)]
	#[serde(default)]
	format: PaletteFormat,
}

impl Palette {
	pub fn run(&self, image: &impl GenericImage) -> anyhow::Result<()> {
		let palette = Palette::extract_from(image)?;
		self.write_palette(palette)
	}

	pub fn extract_from(image: &impl GenericImage) -> anyhow::Result<PixelPalette> {
		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())
	}

	pub fn write_palette(&self, mut colours: PixelPalette) -> anyhow::Result<()> {
		sort_by_hue(&mut colours);

		match self.format {
			PaletteFormat::Png => {
				let num_colours = colours.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 colours.iter().enumerate() {
					out_image.put_pixel(
						(idx % image_width) as u32,
						(idx / image_width) as u32,
						Rgba::from(colour),
					);
				}

				out_image.save(self.output.as_str())?;
			}
			PaletteFormat::Columns => {
				let image_width = 2;
				let image_height = colours.len();
				let mut out_image = new_image(image_width as u32, image_height as u32);

				for (idx, colour) in colours.iter().enumerate() {
					out_image.put_pixel(0, idx as u32, Rgba::from(colour));
					out_image.put_pixel(1, idx as u32, Rgba::from(colour));
				}

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

		Ok(())
	}

	pub fn from_columns(image: &impl GenericImage) -> anyhow::Result<ColourMapping> {
		if image.width() != 2 {
			return Err(Error::msg("Image must have a pixel width of 2"))?;
		}

		let mut mapping = ColourMapping::with_capacity(image.height() as usize);
		for y in 0..image.height() {
			let left = image.get_pixel(0, y);
			let right = image.get_pixel(1, y);
			let left = left.to_rgba();
			let right = right.to_rgba();
			let data = Rgba::from([
				left.0[0].to_u8().unwrap(),
				left.0[1].to_u8().unwrap(),
				left.0[2].to_u8().unwrap(),
				left.0[3].to_u8().unwrap(),
			]);
			let left_pixel = BasicRgba::from(data);
			let data = Rgba::from([
				right.0[0].to_u8().unwrap(),
				right.0[1].to_u8().unwrap(),
				right.0[2].to_u8().unwrap(),
				right.0[3].to_u8().unwrap(),
			]);
			let right_pixel = BasicRgba::from(data);
			mapping.insert(left_pixel, right_pixel);
		}

		Ok(mapping)
	}

	pub fn calculate_mapping(from: &PixelPalette, to: &PixelPalette) -> 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 => {
						log::warn!("No matching set for {} with col {:?}", idx, &colour);

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

		out_map
	}
}