import numpy as np
import math
from skimage import color as sc
import cv2

####################
# metric
####################
def calc_psnr_ssim(img1, img2, crop_border, test_Y=True):
    #
    # img1 = img1 / 255.
    # img2 = img2 / 255.
    h, w = img1.shape[:2]

    if test_Y and img1.shape[2] == 3:  # evaluate on Y channel in YCbCr color space
        im1_in = sc.rgb2ycbcr(img1)[:, :, 0]
        im2_in = sc.rgb2ycbcr(img2)[:, :, 0]
    else:
        im1_in = img1
        im2_in = img2
    if im1_in.ndim == 3:
        cropped_im1 = im1_in[
            crop_border : h - crop_border, crop_border : w - crop_border, :
        ]
        cropped_im2 = im2_in[
            crop_border : h - crop_border, crop_border : w - crop_border, :
        ]
    elif im1_in.ndim == 2:
        cropped_im1 = im1_in[
            crop_border : h - crop_border, crop_border : w - crop_border
        ]
        cropped_im2 = im2_in[
            crop_border : h - crop_border, crop_border : w - crop_border
        ]
    else:
        raise ValueError(
            "Wrong image dimension: {}. Should be 2 or 3.".format(im1_in.ndim)
        )

    psnr = calc_psnr(cropped_im1, cropped_im2)
    ssim = calc_ssim(cropped_im1, cropped_im2)
    return psnr, ssim


def calc_psnr(img1, img2):
    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    mse = np.mean((img1 - img2) ** 2)
    if mse == 0:
        return float("inf")
    return 20 * math.log10(255.0 / math.sqrt(mse))


def ssim(img1, img2):

    C1 = (0.01 * 255) ** 2
    C2 = (0.03 * 255) ** 2

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1 ** 2
    mu2_sq = mu2 ** 2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
        (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    )
    return ssim_map.mean()


def calc_ssim(img1, img2):

    """calculate SSIM
    the same outputs as MATLAB's
    img1, img2: [0, 255]
    """
    if not img1.shape == img2.shape:
        raise ValueError("Input images must have the same dimensions.")
    if img1.ndim == 2:
        return ssim(img1, img2)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(3):
                ssims.append(ssim(img1, img2))
            return np.array(ssims).mean()
        elif img1.shape[2] == 1:
            return ssim(np.squeeze(img1), np.squeeze(img2))
    else:
        raise ValueError("Wrong input image dimensions.")
