use image::Rgb;
use exoquant::ditherer::FloydSteinberg;
use exoquant::Remapper;
use exoquant::optimizer::WeightedKMeans;
use exoquant::generate_palette;
use exoquant::SimpleColorSpace;
use image::DynamicImage;
use image::GenericImageView;
use exoquant::Color;

const MAX_SIZE_PIXELS: u32 = 160_000;
const RESIZE_TARGET_PIXELS: f32 = 90_000.;
const NUM_COLORS: usize = 1;

pub struct Palette {
	pub raw_colors: [Rgb<u8>; NUM_COLORS],
	color_frequencies: [usize; NUM_COLORS],
}

impl Palette {
	pub fn generate(img: &DynamicImage) -> Self {

		let thumb;

		// Scale the image down if it's too big
		let thumb = if img.width() * img.height() > MAX_SIZE_PIXELS {
			let w = img.width() as f32;
			let h = img.height() as f32;
			let ratio = w / h;

			let new_width = (ratio * RESIZE_TARGET_PIXELS).sqrt();
			let new_height = RESIZE_TARGET_PIXELS / new_width;

			thumb = img.thumbnail(new_width as u32, new_height as u32);
			&thumb
		} else {
			img
		};

		// Convert to exoquant image
		let width = thumb.width();
		let exo_img: Vec<Color> = thumb.to_rgb8()
			.pixels()
			.map(|p| Color {r: p[0], g: p[1], b: p[2], a: 255})
			.collect();

		// Generate histogram
		let histogram = exo_img.iter().copied().collect();

		// Generate raw palette
		let colorspace = SimpleColorSpace::default();
		let palette = generate_palette(
			&histogram,
			&colorspace,
			&WeightedKMeans,
			NUM_COLORS,
		);

		// Quantize the image & count the frequency of each color
		let ditherer = FloydSteinberg::new();
		let remapper = Remapper::new(&palette, &colorspace, &ditherer);
		let remapped = remapper.remap_iter(
			Box::new(exo_img.into_iter()),
			width as usize,
		);
		let color_counts = remapped.fold([0; NUM_COLORS], |mut counts, pixel| {
			counts[pixel as usize] += 1;
			counts
		});

		// Convert to a more comfortable image::Rgba color format
		let palette: Vec<Rgb<u8>> = palette.into_iter()
			.map(|c| Rgb([c.r, c.g, c.b]))
			.collect();
		let palette: [Rgb<u8>; NUM_COLORS] = palette.try_into().unwrap();


		Self {
			raw_colors: palette,
			color_frequencies: color_counts,
		}
	}

	pub fn dominant_color(&self) -> Rgb<u8> {
		let max_index = self.color_frequencies
			.iter()
			.enumerate()
			.max_by_key(|(_i, count)| *count)
			.unwrap()
			.0;

		self.raw_colors[max_index]
	}
}