use anyhow::{bail, Result};
use rayon::prelude::*;
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};

pub trait MetricElem: Send + Sync + 'static + Serialize + for<'a> Deserialize<'a> {
    fn dist(&self, _: &Self) -> f64;
}

impl MetricElem for f64 {
    fn dist(&self, b: &f64) -> f64 {
        (self - b).abs()
    }
}

pub trait EmbedderT: Send + Sync {
    type Embedding: MetricElem;
    const NAME: &'static str;

    fn embed(&self, _: &Path) -> Result<Self::Embedding>;
}

pub trait BatchEmbedder: Send + Sync {
    type Embedding: MetricElem;
    const NAME: &'static str;

    fn embeds(&mut self, _: &[PathBuf]) -> Result<Vec<Self::Embedding>>;
}

impl<T: EmbedderT> BatchEmbedder for T {
    type Embedding = T::Embedding;
    const NAME: &'static str = T::NAME;

    fn embeds(&mut self, paths: &[PathBuf]) -> Result<Vec<Self::Embedding>> {
        paths.par_iter()
        .map(|p| self.embed(p))
        .collect::<Vec<_>>()
        .into_iter()
        .try_collect()
    }
}

pub struct BrightnessEmbedder;
impl EmbedderT for BrightnessEmbedder {
    type Embedding = f64;
    const NAME: &'static str = "Brightness";

    fn embed(&self, path: &Path) -> Result<f64> {
        let im = image::open(path)?;
        let num_bytes = 3 * (im.height() * im.width());

        if num_bytes == 0 {
            bail!("Encountered NaN brightness, due to an empty image");
        }

        Ok(im.to_rgb8()
            .iter()
            .map(|e| *e as u64)
            .sum::<u64>() as f64 / num_bytes as f64)
    }
}

#[repr(transparent)]
#[derive(Serialize, Deserialize)]
pub struct Hue (f64);
impl MetricElem for Hue {
    fn dist(&self, b: &Hue) -> f64 {
        let d = self.0.dist(&b.0);
        d.min(6.-d)
    }
}

pub struct HueEmbedder;
impl EmbedderT for HueEmbedder {
    type Embedding = Hue;
    const NAME: &'static str = "Hue";

    fn embed(&self, path: &Path) -> Result<Hue> {
        let im = image::open(path)?;
        let num_pixels = im.height() * im.width();
        let [sr, sg, sb] = im
            .to_rgb8()
            .pixels()
            .fold([0, 0, 0], |[or, og, ob], n| {
                let [nr, ng, nb] = n.0;
                [or + nr as u64, og + ng as u64, ob + nb as u64]
            })
            .map(|e| e as f64 / 255. / num_pixels as f64);

        let hue =
            if sr >= sg && sr >= sb {
                (sg - sb) / (sr - sg.min(sb))
            }
            else if sg >= sb {
                2. + (sb - sr) / (sg - sr.min(sb))
            }
            else {
                4. + (sr - sg) / (sb - sr.min(sg))
            };

        if hue.is_nan() {
            bail!("Encountered NaN hue, possibly because of a colorless or empty image");
        }

        Ok(Hue(hue))
    }
}

impl MetricElem for (f64, f64, f64) {
    fn dist(&self, o: &(f64, f64, f64)) -> f64 {
        let (dr, dg, db) =
            ((self.0 - o.0), (self.1 - o.1), (self.2 - o.2));
        (dr*dr + dg*dg + db*db).sqrt()
    }
}

pub struct ColorEmbedder;
impl EmbedderT for ColorEmbedder {
    type Embedding = (f64, f64, f64);
    const NAME: &'static str = "Color";

    fn embed(&self, path: &Path) -> Result<(f64, f64, f64)> {
        let im = image::open(path)?;
        let num_pixels = im.height() * im.width();
        let [sr, sg, sb] = im
            .to_rgb8()
            .pixels()
            .fold([0, 0, 0], |[or, og, ob], n| {
                let [nr, ng, nb] = n.0;
                [or + nr as u64, og + ng as u64, ob + nb as u64]
            })
            .map(|e| e as f64 / num_pixels as f64);

        Ok((sr, sg, sb))
    }
}