import numpy as np
import imageio
from skimage.transform import resize
from scipy.ndimage import gaussian_filter


def mssim(
    x: np.ndarray,
    y: np.ndarray,
) -> float:
    # Standard choice for the parameters
    K1 = 0.01
    K2 = 0.03
    sigma = 1.5
    truncate = 3.5
    m = 1
    C1 = (K1 * m) ** 2
    C2 = (K2 * m) ** 2

    # radius size of the local window (needed for
    # normalizing the standard deviation)
    r = int(truncate * sigma + 0.5)
    win_size = 2 * r + 1
    # use these arguments for the gaussian filtering
    # e.g. filtered = gaussian_filter(x, **filter_args)
    filter_args = {
        'sigma': sigma,
        'truncate': truncate
    }

    # Implement Eq. (9) from assignment sheet
    # S should be an "image" of the SSIM evaluated for a window 
    # centered around the corresponding pixel in the original input image
    S = np.ones_like(x)

    mu_x = gaussian_filter(x, **filter_args)
    mu_y = gaussian_filter(y, **filter_args)

    # Compute local variances and covariance
    n = win_size ** 2 # number of pixels in the window
    # n~ = n/(n-1)
    # E[x^2] −E[x]^2). = (gaussian_filter(x * x, **filter_args) - mu_x * mu_x)
    sigma_x_sq = (gaussian_filter(x * x, **filter_args) - mu_x * mu_x) * (n / (n - 1))
    sigma_y_sq = (gaussian_filter(y * y, **filter_args) - mu_y * mu_y) * (n / (n - 1))
    sigma_xy = (gaussian_filter(x * y, **filter_args) - mu_x * mu_y) * (n / (n - 1))
    
    # Compute SSIM
    S = ((2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)) / \
        ((mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x_sq + sigma_y_sq + C2))
    
    # crop to remove boundary artifacts, return MSSIM
    pad = (win_size - 1) // 2
    return S[pad:-pad, pad:-pad].mean()


def psnr(
    x: np.ndarray,
    y: np.ndarray,
) -> float:
    # Implement Eq. (2) without for loops
    mse = np.mean((x - y) ** 2)
    m = 1
    psnr_value = 10 * np.log10(m ** 2 / mse)
    return psnr_value


def psnr_for(
    x: np.ndarray,
    y: np.ndarray,
) -> float:
    # Implement Eq. (2) using for loops
    m, n = x.shape
    mse = 0
    # Eq. (1) from the assignment sheet for mse
    for i in range(m):
        for j in range(n):
            mse += (x[i, j] - y[i, j]) ** 2
    mse /= (m * n)
    m = 1
    psnr_value = 10 * np.log10(m ** 2 / mse)
    return psnr_value


def interpolation_error():
    x = imageio.imread('./girl.png') / 255.
    shape_lower = (x.shape[0] // 2, x.shape[1] // 2)
    # downsample image to half the resolution
    # and successively upsample to the original resolution
    # using no nearest neighbor, linear and cubic interpolation
    nearest, linear, cubic = [
        resize(resize(
            x, shape_lower, order=order, anti_aliasing=False
        ), x.shape, order=order, anti_aliasing=False)
        for order in [0, 1, 3]
    ]

    for label, rescaled in zip(
        ['nearest', 'linear', 'cubic'],
        [nearest, linear, cubic]
    ):
        print(label)
        print(mssim(x, rescaled))
        print(psnr(x, rescaled))
        print(psnr_for(x, rescaled))
        imageio.imwrite('girl_' + label + '.png', (rescaled * 255).astype(np.uint8))


if __name__ == '__main__':
    interpolation_error()