use rayon::prelude::*;
use std::path::PathBuf;

pub trait MetricElem {
    fn dist(&self, _: &Self) -> f64;
}

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

pub trait EmbedderT {
    type Embedding: MetricElem;

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

pub struct BrightnessEmbedder;
impl EmbedderT for BrightnessEmbedder {
    type Embedding = f64;

    fn embed(&mut self, paths: &[PathBuf]) -> Result<Vec<f64>, String> {
        paths
            .par_iter()
            .map(|p| {
                let im = image::open(p).map_err(|e| e.to_string())?;
                let num_bytes = 3 * (im.height() * im.width());

                if num_bytes == 0 {
                    Err("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)
            })
            .collect::<Vec<_>>()
            .into_iter()
            .try_collect()
    }
}

#[repr(transparent)]
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;

    fn embed(&mut self, paths: &[PathBuf]) -> Result<Vec<Hue>, String> {
        paths
            .par_iter()
            .map(|p| {
                let im = image::open(p).map_err(|e| e.to_string())?;
                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() {
                    Err("Encountered NaN hue, possibly because of a colorless or empty image")?;
                }

                Ok(Hue(hue))
            })
            .collect::<Vec<_>>()
            .into_iter()
            .try_collect()
    }
}

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);

    fn embed(&mut self, paths: &[PathBuf]) -> Result<Vec<(f64, f64, f64)>, String> {
        paths
            .par_iter()
            .map(|p| {
                let im = image::open(p).map_err(|e| e.to_string())?;
                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))
            })
            .collect::<Vec<_>>()
            .into_iter()
            .try_collect()
    }
}