extern crate image;

use std::env::args;
use std::fs;

use image::{ImageResult, ImageError};

const MAX: u32 = 255;

fn mse_files(path1: &str, path2: &str) -> ImageResult<f64> {

    let twenty_log10_max: f64 = 20.0 * (MAX as f64).log10();

    let img1 = image::open(path1)?;
    let img2 = image::open(path2)?;

    use image::GenericImage;

    if img1.dimensions() != img2.dimensions() {
        return Err(ImageError::DimensionError);
    }

    let mut mse = 0.0;

    for x in 0 .. img1.dimensions().0 {
        for y in 0 .. img1.dimensions().1 {
            let p1 = img1.get_pixel(x, y);
            let p2 = img2.get_pixel(x, y);

            let dx = p1[0] as f64 - p2[0] as f64;
            let dy = p1[1] as f64 - p2[1] as f64;
            let dz = p1[2] as f64 - p2[2] as f64;

            let dx = dx * dx;
            let dy = dy * dy;
            let dz = dz * dz;

            mse += dx + dy + dz;
        }
    }

    mse /= (img1.dimensions().0 * img1.dimensions().1 * 3) as f64;

    Ok(twenty_log10_max - 10.0 * mse.log10())
}

fn mse_dir(path1: &str, path2: &str) -> Vec<(String, ImageResult<f64>)> {
    let mut dirs1: Vec<_> = fs::read_dir(path1).unwrap()
        .map(|x|x.unwrap())
        .filter(|x| x.file_name().into_string().unwrap().ends_with(".png"))
        .collect();

    let mut dirs2: Vec<_> = fs::read_dir(path2).unwrap()
        .map(|x|x.unwrap())
        .filter(|x| x.file_name().into_string().unwrap().ends_with(".png"))
        .collect();

    dirs1.sort_by_key(|x| x.file_name());
    dirs2.sort_by_key(|x| x.file_name());

    let mut res = vec!();

    for file in dirs1 {
        let p1 = path1.to_owned() + &file.file_name().to_owned().into_string().unwrap();
        let p2 = path2.to_owned() + &file.file_name().to_owned().into_string().unwrap();
        let mse = mse_files(&p1, &p2);

        let elt = (file.file_name().to_owned().into_string().unwrap(), mse.unwrap());
        println!("{} {}", elt.0, elt.1);
        res.push((elt.0, Ok(elt.1)));
    }

    res
}

fn main() {
    let path1 = args().nth(1).unwrap();
    let path2 = args().nth(2).unwrap();

    let _ = mse_dir(&path1, &path2);
}