smartcrop2 0.4.0

use std::{num::NonZeroU32, ops::RangeInclusive};

use crate::{math, Crop, Error, Image, ResizableImage, Score, ScoredCrop, RGB};

const PRESCALE: bool = true;
const PRESCALE_MIN: f64 = 400.00;
const MIN_SCALE: f64 = 1.0;
const MAX_SCALE: f64 = 1.0;
// STEP * minscale rounded down to the next power of two should be good
const STEP: f64 = 8.0;
const SCALE_STEP: f64 = 0.1;

const SCORE_DOWN_SAMPLE: f64 = 8.0;

const SKIN_WEIGHT: f64 = 1.8;
const DETAIL_WEIGHT: f64 = 0.2;

const SKIN_BRIGHTNESS_RANGE: RangeInclusive<f64> = 0.2..=1.0;
const SKIN_THRESHOLD: f64 = 0.8;
const SKIN_BIAS: f64 = 0.01;

const SATURATION_BRIGHTNESS_RANGE: RangeInclusive<f64> = 0.05..=0.9;
const SATURATION_THRESHOLD: f64 = 0.4;
const SATURATION_BIAS: f64 = 0.2;
const SATURATION_WEIGHT: f64 = 0.1;

#[derive(Debug)]
struct ImageMap {
    width: u32,
    height: u32,

    pixels: Vec<Vec<RGB>>,
}

impl ImageMap {
    fn new(width: u32, height: u32) -> ImageMap {
        let white = RGB::new(255, 255, 255);
        let pixels = vec![vec![white; height as usize]; width as usize];
        ImageMap {
            width,
            height,
            pixels,
        }
    }

    fn set(&mut self, x: u32, y: u32, color: RGB) {
        self.pixels[x as usize][y as usize] = color
    }

    fn get(&self, x: u32, y: u32) -> RGB {
        self.pixels[x as usize][y as usize]
    }

    fn down_sample(self, factor: u32) -> Self {
        let width = (self.width as f64 / factor as f64).floor() as u32;
        let height = (self.height as f64 / factor as f64).floor() as u32;
        let mut output = ImageMap::new(width, height);
        //        let data = output.data;
        let ifactor2: f64 = 1.0 / (factor as f64 * factor as f64);

        let max = |a, b| {
            if a > b {
                a
            } else {
                b
            }
        };

        for y in 0..height {
            for x in 0..width {
                let mut r: f64 = 0.0;
                let mut g: f64 = 0.0;
                let mut b: f64 = 0.0;

                let mut mr: f64 = 0.0;
                let mut mg: f64 = 0.0;

                for v in 0..factor {
                    for u in 0..factor {
                        let ix = x * factor + u;
                        let iy = y * factor + v;
                        let icolor = self.get(ix, iy);

                        r += icolor.r as f64;
                        g += icolor.g as f64;
                        b += icolor.b as f64;
                        mr = max(mr, icolor.r as f64);
                        mg = max(mg, icolor.g as f64);
                    }
                }

                // this is some funky magic to preserve detail a bit more for
                // skin (r) and detail (g). saturation (b) does not get this boost.
                output.set(
                    x,
                    y,
                    RGB::new(
                        (r * ifactor2 * 0.5 + mr * 0.5).round() as u8,
                        (g * ifactor2 * 0.7 + mg * 0.3).round() as u8,
                        (b * ifactor2).round() as u8,
                    ),
                )
            }
        }

        output
    }
}

/// Analyze the image and find the best crop of the given aspect ratio
pub fn find_best_crop<I: Image + ResizableImage<RI>, RI: Image>(
    img: &I,
    width: NonZeroU32,
    height: NonZeroU32,
) -> Result<ScoredCrop, Error> {
    if img.width() == 0 || img.height() == 0 {
        return Err(Error::ZeroSizedImage);
    }

    let width = width.get() as f64;
    let height = height.get() as f64;

    let scale = f64::min((img.width() as f64) / width, (img.height() as f64) / height);

    // resize image for faster processing
    if PRESCALE {
        let f = PRESCALE_MIN / f64::min(img.width() as f64, img.height() as f64);
        let prescalefactor = f.min(1.0);

        let crop_width = (width * scale * prescalefactor).max(1.0).round() as u32;
        let crop_height = (height * scale * prescalefactor).max(1.0).round() as u32;
        let real_min_scale = calculate_real_min_scale(scale);

        let new_width = ((img.width() as f64) * prescalefactor).round() as u32;
        let new_height = (prescalefactor * img.height() as f64).round() as u32;

        let old_width = img.width() as f64;
        let old_height = img.height() as f64;

        let img = img.resize(new_width, new_height);

        assert!(img.width() == crop_width || img.height() == crop_height);
        let top_crop = analyse(
            &img,
            NonZeroU32::new(crop_width).unwrap(),
            NonZeroU32::new(crop_height).unwrap(),
            real_min_scale,
        );

        let post_scale_w = img.width() as f64 / old_width;
        let post_scale_h = img.height() as f64 / old_height;
        let post_scale_factor = f64::max(post_scale_w, post_scale_h);

        Ok(top_crop.scale(1.0 / post_scale_factor))
    } else {
        let crop_width = (width * scale).round() as u32;
        let crop_height = (height * scale).round() as u32;
        let real_min_scale = calculate_real_min_scale(scale);

        assert!(img.width() == crop_width || img.height() == crop_height);
        let top_crop = analyse(
            img,
            NonZeroU32::new(crop_width).unwrap(),
            NonZeroU32::new(crop_height).unwrap(),
            real_min_scale,
        );
        Ok(top_crop)
    }
}

fn calculate_real_min_scale(scale: f64) -> f64 {
    (1.0 / scale).clamp(MIN_SCALE, MAX_SCALE)
}

fn analyse<I: Image>(
    img: &I,
    crop_width: NonZeroU32,
    crop_height: NonZeroU32,
    real_min_scale: f64,
) -> ScoredCrop {
    assert!(img.width() >= crop_width.get());
    assert!(img.height() >= crop_height.get());

    let mut o = ImageMap::new(img.width(), img.height());

    edge_detect(img, &mut o);

    skin_detect(img, &mut o);

    saturation_detect(img, &mut o);

    let cs: Vec<Crop> = crops(&o, crop_width.get(), crop_height.get(), real_min_scale);
    assert!(!cs.is_empty());
    let score_output = o.down_sample(SCORE_DOWN_SAMPLE as u32);
    let top_crop: Option<ScoredCrop> = cs
        .iter()
        .map(|crop| ScoredCrop {
            crop: crop.clone(),
            score: score(&score_output, crop),
        })
        .fold(None, |result, scored_crop| {
            Some(match result {
                None => scored_crop,
                Some(result) => {
                    if result.score.total > scored_crop.score.total {
                        result
                    } else {
                        scored_crop
                    }
                }
            })
        });

    top_crop.unwrap()
}

fn edge_detect<I: Image>(i: &I, o: &mut ImageMap) {
    //TODO check type casts if those are safe

    let w = i.width() as usize;
    let h = i.height() as usize;
    let cies = make_cies(i);

    for y in 0..h {
        for x in 0..w {
            let color = i.get(x as u32, y as u32);

            let lightness = if x == 0 || x >= w - 1 || y == 0 || y >= h - 1 {
                cies[y * w + x]
            } else {
                cies[y * w + x] * 4.0
                    - cies[x + (y - 1) * w]
                    - cies[x - 1 + y * w]
                    - cies[x + 1 + y * w]
                    - cies[x + (y + 1) * w]
            };

            let g = math::bounds(lightness);

            let nc = RGB { g, ..color };
            o.set(x as u32, y as u32, nc)
        }
    }
}

fn make_cies<I: Image>(img: &I) -> Vec<f64> {
    //TODO `cies()` can probably be made RGB member that will make this function redundant
    let w = img.width();
    let h = img.height();
    let size = w as u64 * h as u64;

    let size = if size > usize::MAX as u64 {
        None
    } else {
        Some(size as usize)
    };

    //TODO error handling
    let mut cies = Vec::with_capacity(size.expect("Too big image dimensions"));

    let mut i: usize = 0;
    for y in 0..h {
        for x in 0..w {
            let color = img.get(x, y);
            cies.insert(i, color.cie());
            i += 1;
        }
    }

    cies
}

fn crops(i: &ImageMap, crop_width: u32, crop_height: u32, real_min_scale: f64) -> Vec<Crop> {
    let mut crops: Vec<Crop> = vec![];
    let width = i.width as f64;
    let height = i.height as f64;

    let min_dimension = f64::min(width, height);

    let crop_w = if crop_width != 0 {
        crop_width as f64
    } else {
        min_dimension
    };
    let crop_h = if crop_height != 0 {
        crop_height as f64
    } else {
        min_dimension
    };

    let y_step = STEP.min(height);
    let x_step = STEP.min(width);

    let mut scale = MAX_SCALE;
    loop {
        if scale < real_min_scale {
            break;
        };

        let stepping = |step| (0..).map(f64::from).map(move |i| i * step);

        for y in stepping(y_step).take_while(|y| y + crop_h * scale <= height) {
            for x in stepping(x_step).take_while(|x| x + crop_w * scale <= width) {
                crops.push(Crop {
                    x: x.round() as u32,
                    y: y.round() as u32,
                    width: (crop_w * scale).round() as u32,
                    height: (crop_h * scale).round() as u32,
                });
            }
        }

        scale -= SCALE_STEP;
    }

    crops
}

fn score(o: &ImageMap, crop: &Crop) -> Score {
    let height = o.height as f64;
    let width = o.width as f64;

    let down_sample = SCORE_DOWN_SAMPLE;
    let inv_down_sample = 1.0 / down_sample;
    let output_height_down_sample = height * down_sample;
    let output_width_down_sample = width * down_sample;

    let mut skin = 0.0;
    let mut detail = 0.0;
    let mut saturation = 0.0;

    for y in (0..)
        .map(|i: u32| i as f64 * SCORE_DOWN_SAMPLE)
        .take_while(|&y| y < output_height_down_sample)
    {
        for x in (0..)
            .map(|i: u32| i as f64 * SCORE_DOWN_SAMPLE)
            .take_while(|&x| x < output_width_down_sample)
        {
            let orig_x = (x * inv_down_sample).round() as u32;
            let orig_y = (y * inv_down_sample).round() as u32;

            let color = o.get(orig_x, orig_y);

            let imp = math::importance(crop, x.round() as u32, y.round() as u32);
            let det = color.g as f64 / 255.0;

            skin += color.r as f64 / 255.0 * (det + SKIN_BIAS) * imp;
            detail += det * imp;
            saturation += color.b as f64 / 255.0 * (det + SATURATION_BIAS) * imp;
        }
    }

    let total = (detail * DETAIL_WEIGHT + skin * SKIN_WEIGHT + saturation * SATURATION_WEIGHT)
        / crop.width as f64
        / crop.height as f64;

    Score {
        skin,
        detail,
        saturation,
        total,
    }
}

fn skin_detect<I: Image>(i: &I, o: &mut ImageMap) {
    let w = i.width();
    let h = i.height();

    for y in 0..h {
        for x in 0..w {
            let lightness = i.get(x, y).cie() / 255.0;
            let skin = math::skin_col(i.get(x, y));

            let nc = if skin > SKIN_THRESHOLD && SKIN_BRIGHTNESS_RANGE.contains(&lightness) {
                let r = (skin - SKIN_THRESHOLD) * (255.0 / (1.0 - SKIN_THRESHOLD));
                let RGB { r: _, g, b } = o.get(x, y);

                RGB {
                    r: math::bounds(r),
                    g,
                    b,
                }
            } else {
                let RGB { r: _, g, b } = o.get(x, y);
                RGB { r: 0, g, b }
            };

            o.set(x, y, nc);
        }
    }
}

fn saturation_detect<I: Image>(i: &I, o: &mut ImageMap) {
    let w = i.width();
    let h = i.height();

    for y in 0..h {
        for x in 0..w {
            let color = i.get(x, y);
            let lightness = color.cie() / 255.0;
            let saturation = color.saturation();

            let nc = if saturation > SATURATION_THRESHOLD
                && SATURATION_BRIGHTNESS_RANGE.contains(&lightness)
            {
                let b =
                    (saturation - SATURATION_THRESHOLD) * (255.0 / (1.0 - SATURATION_THRESHOLD));
                let RGB { r, g, b: _ } = o.get(x, y);
                RGB {
                    r,
                    g,
                    b: math::bounds(b),
                }
            } else {
                let RGB { r, g, b: _ } = o.get(x, y);
                RGB { r, g, b: 0 }
            };

            o.set(x, y, nc);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const WHITE: RGB = RGB {
        r: 255,
        g: 255,
        b: 255,
    };
    const BLACK: RGB = RGB { r: 0, g: 0, b: 0 };
    const RED: RGB = RGB { r: 255, g: 0, b: 0 };
    const GREEN: RGB = RGB { r: 0, g: 255, b: 0 };
    const BLUE: RGB = RGB { r: 0, g: 0, b: 255 };
    const SKIN: RGB = RGB {
        r: 255,
        g: 200,
        b: 159,
    };

    #[derive(Debug, Clone)]
    struct TestImage {
        w: u32,
        h: u32,
        pixels: Vec<Vec<RGB>>,
    }

    impl TestImage {
        fn new(w: u32, h: u32, pixels: Vec<Vec<RGB>>) -> TestImage {
            TestImage { w, h, pixels }
        }

        fn new_single_pixel(pixel: RGB) -> TestImage {
            TestImage {
                w: 1,
                h: 1,
                pixels: vec![vec![pixel]],
            }
        }

        fn new_from_fn<G>(w: u32, h: u32, generate: G) -> TestImage
        where
            G: Fn(u32, u32) -> RGB,
        {
            let mut pixels = vec![vec![WHITE; h as usize]; w as usize];

            for y in 0..h {
                for x in 0..w {
                    pixels[x as usize][y as usize] = generate(x, y)
                }
            }

            TestImage { w, h, pixels }
        }
    }

    impl ImageMap {
        fn from_image<I: Image>(image: &I) -> ImageMap {
            let mut image_map = ImageMap::new(image.width(), image.height());

            for y in 0..image.height() {
                for x in 0..image.width() {
                    let color = image.get(x, y);
                    image_map.set(x, y, color);
                }
            }

            image_map
        }
    }

    impl Image for TestImage {
        fn width(&self) -> u32 {
            self.w
        }

        fn height(&self) -> u32 {
            self.h
        }

        fn get(&self, x: u32, y: u32) -> RGB {
            self.pixels[x as usize][y as usize]
        }
    }

    impl ResizableImage<TestImage> for TestImage {
        fn resize(&self, width: u32, _height: u32) -> TestImage {
            if width == self.w {
                return self.clone();
            }

            let height = (self.h as f64 * width as f64 / self.w as f64).round() as u32;

            //TODO Implement more or less correct resizing
            TestImage {
                w: width,
                h: height,
                pixels: self.pixels.clone(),
            }
        }

        fn crop_and_resize(&self, _crop: Crop, width: u32, height: u32) -> TestImage {
            self.resize(width, height)
        }
    }

    #[test]
    fn saturation_tests() {
        assert_eq!(0.0, BLACK.saturation());
        assert_eq!(0.0, WHITE.saturation());
        assert_eq!(1.0, RGB::new(255, 0, 0).saturation());
        assert_eq!(1.0, RGB::new(0, 255, 0).saturation());
        assert_eq!(1.0, RGB::new(0, 0, 255).saturation());
        assert_eq!(1.0, RGB::new(0, 255, 255).saturation());
    }

    #[test]
    fn image_map_test() {
        let mut image_map = ImageMap::new(1, 2);

        assert_eq!(image_map.width, 1);
        assert_eq!(image_map.height, 2);

        assert_eq!(image_map.get(0, 0), RGB::new(255, 255, 255));
        assert_eq!(image_map.get(0, 1), RGB::new(255, 255, 255));

        let red = RGB::new(255, 0, 0);
        image_map.set(0, 0, red);
        assert_eq!(image_map.get(0, 0), red);

        let green = RGB::new(0, 255, 0);
        image_map.set(0, 1, green);
        assert_eq!(image_map.get(0, 1), green);
    }

    #[test]
    fn crops_test() {
        let real_min_scale = MIN_SCALE;

        let crops = crops(&ImageMap::new(8, 8), 8, 8, real_min_scale);

        assert_eq!(
            crops[0],
            Crop {
                x: 0,
                y: 0,
                width: 8,
                height: 8
            }
        )
    }

    #[test]
    fn score_test_image_with_single_black_pixel_then_score_is_zero() {
        let mut i = ImageMap::new(1, 1);
        i.set(0, 0, RGB::new(0, 0, 0));

        let s = score(
            &i,
            &Crop {
                x: 0,
                y: 0,
                width: 1,
                height: 1,
            },
        );

        assert_eq!(
            s,
            Score {
                detail: 0.0,
                saturation: 0.0,
                skin: 0.0,
                total: 0.0
            }
        );
    }

    #[test]
    fn score_test_image_with_single_white_pixel_then_score_is_the_same_as_for_js_version() {
        let mut i = ImageMap::new(1, 1);
        i.set(0, 0, RGB::new(255, 255, 255));

        let s = score(
            &i,
            &Crop {
                x: 0,
                y: 0,
                width: 1,
                height: 1,
            },
        );

        let js_version_score = Score {
            detail: -6.404213562373096,
            saturation: -7.685056274847715,
            skin: -6.468255697996827,
            total: -13.692208596353678,
        };

        assert_eq!(s, js_version_score);
    }

    #[test]
    fn skin_detect_single_pixel_test() {
        let detect_pixel = |color: RGB| {
            let image = TestImage::new_single_pixel(color);
            let mut o = ImageMap::new(1, 1);
            o.set(0, 0, color);

            skin_detect(&image, &mut o);
            o.get(0, 0)
        };

        assert_eq!(detect_pixel(WHITE), RGB::new(0, 255, 255));
        assert_eq!(detect_pixel(BLACK), RGB::new(0, 0, 0));
        assert_eq!(detect_pixel(RED), RGB::new(0, 0, 0));
        assert_eq!(detect_pixel(GREEN), RGB::new(0, 255, 0));
        assert_eq!(detect_pixel(BLUE), RGB::new(0, 0, 255));
        assert_eq!(detect_pixel(SKIN), RGB::new(159, 200, 159));
    }

    #[test]
    fn edge_detect_single_pixel_image_test() {
        let edge_detect_pixel = |color: RGB| {
            let image = TestImage::new_single_pixel(color);
            let mut o = ImageMap::new(1, 1);
            o.set(0, 0, color);

            edge_detect(&image, &mut o);

            o.get(0, 0)
        };

        assert_eq!(edge_detect_pixel(BLACK), BLACK);
        assert_eq!(edge_detect_pixel(WHITE), WHITE);
        assert_eq!(edge_detect_pixel(RED), RGB::new(255, 18, 0));
        assert_eq!(edge_detect_pixel(GREEN), RGB::new(0, 182, 0));
        assert_eq!(edge_detect_pixel(BLUE), RGB::new(0, 131, 255));
        assert_eq!(edge_detect_pixel(SKIN), RGB::new(255, 243, 159));
    }

    #[test]
    fn edge_detect_3x3() {
        let image = TestImage::new(
            3,
            3,
            vec![
                vec![RED, GREEN, BLUE],
                vec![GREEN, BLUE, RED],
                vec![BLUE, RED, GREEN],
            ],
        );
        let mut o = ImageMap::new(3, 3);

        edge_detect(&image, &mut o);

        assert_eq!(
            o.get(0, 0),
            RGB {
                r: 255,
                g: 18,
                b: 0
            }
        );
        assert_eq!(
            o.get(0, 0),
            RGB {
                r: 255,
                g: 18,
                b: 0
            }
        );
        assert_eq!(o.get(1, 0), RGB { r: 0, g: 182, b: 0 });
        assert_eq!(
            o.get(2, 0),
            RGB {
                r: 0,
                g: 131,
                b: 255
            }
        );
        assert_eq!(o.get(0, 1), RGB { r: 0, g: 182, b: 0 });
        assert_eq!(
            o.get(1, 1),
            RGB {
                r: 0,
                g: 121,
                b: 255
            }
        );
        assert_eq!(
            o.get(2, 1),
            RGB {
                r: 255,
                g: 18,
                b: 0
            }
        );
        assert_eq!(
            o.get(0, 2),
            RGB {
                r: 0,
                g: 131,
                b: 255
            }
        );
        assert_eq!(
            o.get(1, 2),
            RGB {
                r: 255,
                g: 18,
                b: 0
            }
        );
        assert_eq!(o.get(2, 2), RGB { r: 0, g: 182, b: 0 });
    }

    #[test]
    fn saturation_detect_3x3() {
        let image = TestImage::new(
            3,
            3,
            vec![
                vec![RED, GREEN, BLUE],
                vec![WHITE, SKIN, BLACK],
                vec![BLUE, RED, GREEN],
            ],
        );
        let mut o = ImageMap::from_image(&image);

        saturation_detect(&image, &mut o);

        assert_eq!(
            o.get(0, 0),
            RGB {
                r: 255,
                g: 0,
                b: 255
            }
        );
        assert_eq!(
            o.get(0, 1),
            RGB {
                r: 0,
                g: 255,
                b: 255
            }
        );
        assert_eq!(o.get(0, 2), RGB { r: 0, g: 0, b: 255 });
        assert_eq!(
            o.get(1, 0),
            RGB {
                r: 255,
                g: 255,
                b: 0
            }
        );
        assert_eq!(
            o.get(1, 1),
            RGB {
                r: 255,
                g: 200,
                b: 0
            }
        );
        assert_eq!(o.get(1, 2), RGB { r: 0, g: 0, b: 0 });
        assert_eq!(o.get(2, 0), RGB { r: 0, g: 0, b: 255 });
        assert_eq!(
            o.get(2, 1),
            RGB {
                r: 255,
                g: 0,
                b: 255
            }
        );
        assert_eq!(
            o.get(2, 2),
            RGB {
                r: 0,
                g: 255,
                b: 255
            }
        );
    }

    #[test]
    fn analyze_test() {
        let image = TestImage::new_from_fn(24, 24, |x, y| {
            let center = 8..16;
            if center.contains(&x) && center.contains(&y) {
                SKIN
            } else {
                WHITE
            }
        });

        let crop = analyse(
            &image,
            NonZeroU32::new(8).unwrap(),
            NonZeroU32::new(8).unwrap(),
            1.0,
        );

        assert_eq!(crop.crop.width, 8);
        assert_eq!(crop.crop.height, 8);
        assert_eq!(crop.crop.x, 8);
        assert_eq!(crop.crop.y, 8);
        assert_eq!(crop.score.saturation, 0.0);
        assert_eq!(crop.score.detail, -1.7647058823529413);
        assert_eq!(crop.score.skin, -0.03993215515362048);
        assert_eq!(crop.score.total, -0.006637797746048519);
    }

    #[test]
    fn crop_scale_test() {
        let crop = Crop {
            x: 2,
            y: 4,
            width: 8,
            height: 16,
        };

        let scaled_crop = crop.scale(0.5);

        assert_eq!(1, scaled_crop.x);
        assert_eq!(2, scaled_crop.y);
        assert_eq!(4, scaled_crop.width);
        assert_eq!(8, scaled_crop.height);
    }

    #[test]
    fn down_sample_test() {
        let image = TestImage::new(
            3,
            3,
            vec![
                vec![RED, GREEN, BLUE],
                vec![SKIN, BLUE, RED],
                vec![BLUE, RED, GREEN],
            ],
        );

        let image_map = ImageMap::from_image(&image);

        let result = image_map.down_sample(3);

        assert_eq!(result.width, 1);
        assert_eq!(result.height, 1);
        assert_eq!(result.get(0, 0), RGB::new(184, 132, 103));
    }
}