#[macro_use]
extern crate lazy_static;
extern crate image;

use image::{imageops::FilterType, GenericImage, ImageResult, Pixel, Rgba};

lazy_static! {
    /// This values is 20 * log10(255)
    pub static ref TWENTY_LOG10_MAX: f64 = 20.0 * (std::u8::MAX as f64).log10();
}

/// Compute a MSE between an image and its color.
///
/// Will panic an error if the two images don't have the same size.
pub fn mse_color<P: Pixel<Subpixel = u8>, T: GenericImage<Pixel = P>>(
    img: &T,
    color: &Rgba<u8>,
) -> f64 {
    let mut mse = 0.0;

    for x in 0..img.dimensions().0 {
        for y in 0..img.dimensions().1 {
            let p1 = img.get_pixel(x, y).to_rgba();

            let dx = p1.0[0] as f64 - color.0[0] as f64;
            let dy = p1.0[1] as f64 - color.0[1] as f64;
            let dz = p1.0[2] as f64 - color.0[2] as f64;
            let da = p1.0[3] as f64 - color.0[3] as f64;

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

            mse += dx + dy + dz + da;
        }
    }

    mse / (img.dimensions().0 * img.dimensions().1 * 4) as f64
}

/// Compute a MSE between two images.
///
/// Will panic an error if the two images don't have the same size.
pub fn mse<
    P1: Pixel<Subpixel = u8>,
    T1: GenericImage<Pixel = P1>,
    P2: Pixel<Subpixel = u8>,
    T2: GenericImage<Pixel = P2>,
>(
    img1: &T1,
    img2: &T2,
) -> f64 {
    assert_eq!(img1.dimensions(), img2.dimensions());

    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).to_rgba();
            let p2 = img2.get_pixel(x, y).to_rgba();

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

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

            mse += dx + dy + dz + da;
        }
    }

    mse / (img1.dimensions().0 * img1.dimensions().1 * 4) as f64
}

/// Resizes img1 to the size of img2 and then compute the MSE between the images.
pub fn mse_resized<P: 'static + Pixel<Subpixel = u8>, T: 'static + GenericImage<Pixel = P>>(
    img1: &T,
    img2: &T,
    filter: FilterType,
) -> f64 {
    let img1 = image::imageops::resize(img1, img2.width(), img2.height(), filter);
    mse(&img1, img2)
}

/// Compute the PSNR from the MSE.
/// Returns 10 * log10(255^2) - 10 * log10(mse)
pub fn mse_to_psnr(mse: f64) -> f64 {
    *TWENTY_LOG10_MAX - 10.0 * mse.log10()
}

/// Computes the PSNR between two images of the same sizes.
pub fn psnr<P: Pixel<Subpixel = u8>, T: GenericImage<Pixel = P>>(img1: &T, img2: &T) -> f64 {
    *TWENTY_LOG10_MAX - 10.0 * mse(img1, img2).log10()
}

/// Computes the PSNR between two images of the same size from their paths.
pub fn psnr_files(path1: &str, path2: &str) -> ImageResult<f64> {
    let img1 = image::open(path1)?;
    let img2 = image::open(path2)?;
    Ok(psnr(&img1, &img2))
}

/// Compute a MSE between an image and its color.
///
/// Will panic an error if the two images don't have the same size.
pub fn mse_color_without_alpha<P: Pixel<Subpixel = u8>, T: GenericImage<Pixel = P>>(
    img: &T,
    color: &Rgba<u8>,
) -> f64 {
    let mut mse = 0.0;

    for x in 0..img.dimensions().0 {
        for y in 0..img.dimensions().1 {
            let p1 = img.get_pixel(x, y).to_rgba();

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

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

            mse += dx + dy + dz;
        }
    }

    mse / (img.dimensions().0 * img.dimensions().1 * 3) as f64
}

/// Compute a MSE between two images.
///
/// Will panic an error if the two images don't have the same size.
pub fn mse_without_alpha<
    P1: Pixel<Subpixel = u8>,
    T1: GenericImage<Pixel = P1>,
    P2: Pixel<Subpixel = u8>,
    T2: GenericImage<Pixel = P2>,
>(
    img1: &T1,
    img2: &T2,
) -> f64 {
    assert_eq!(img1.dimensions(), img2.dimensions());

    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).to_rgba();
            let p2 = img2.get_pixel(x, y).to_rgba();

            let dx = p1.0[0] as f64 - p2.0[0] as f64;
            let dy = p1.0[1] as f64 - p2.0[1] as f64;
            let dz = p1.0[2] as f64 - p2.0[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
}

/// Resizes img1 to the size of img2 and then compute the MSE between the images.
pub fn mse_resized_without_alpha<
    P: 'static + Pixel<Subpixel = u8>,
    T: 'static + GenericImage<Pixel = P>,
>(
    img1: &T,
    img2: &T,
    filter: FilterType,
) -> f64 {
    let img1 = image::imageops::resize(img1, img2.width(), img2.height(), filter);
    mse_without_alpha(&img1, img2)
}

/// Computes the PSNR between two images of the same sizes.
pub fn psnr_without_alpha<P: Pixel<Subpixel = u8>, T: GenericImage<Pixel = P>>(
    img1: &T,
    img2: &T,
) -> f64 {
    *TWENTY_LOG10_MAX - 10.0 * mse_without_alpha(img1, img2).log10()
}

/// Computes the PSNR between two images of the same size from their paths.
pub fn psnr_files_without_alpha(path1: &str, path2: &str) -> ImageResult<f64> {
    let img1 = image::open(path1)?;
    let img2 = image::open(path2)?;
    Ok(psnr_without_alpha(&img1, &img2))
}