#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2021-12-07 21:37:58

import sys
import math
import torch
import numpy as np
import scipy.ndimage as snd
from scipy.special import softmax
from scipy.interpolate import interp2d

import torch.nn.functional as F


from . import util_image
from ResizeRight.resize_right import resize

def modcrop(im, sf):
    h, w = im.shape[:2]
    h -= (h % sf)
    w -= (w % sf)
    return im[:h, :w,]

#--------------------------------------------Kernel-----------------------------------------------
def sigma2kernel(sigma, k_size=21, sf=3, shift=False):
    '''
    Generate Gaussian kernel according to cholesky decomposion.
    Input:
        sigma: N x 1 x 2 x 2 torch tensor, covariance matrix
        k_size: integer, kernel size
        sf: scale factor
    Output:
        kernel: N x 1 x k x k torch tensor
    '''
    try:
        sigma_inv = torch.inverse(sigma)
    except:
        sigma_disturb = sigma + torch.eye(2, dtype=sigma.dtype, device=sigma.device).unsqueeze(0).unsqueeze(0) * 1e-5
        sigma_inv = torch.inverse(sigma_disturb)

    # Set expectation position (shifting kernel for aligned image)
    if shift:
        center = k_size // 2 + 0.5 * (sf - k_size % 2)                         # + 0.5 * (sf - k_size % 2)
    else:
        center = k_size // 2

    # Create meshgrid for Gaussian
    X, Y = torch.meshgrid(torch.arange(k_size), torch.arange(k_size))
    Z = torch.stack((X, Y), dim=2).to(device=sigma.device, dtype=sigma.dtype).view(1, -1, 2, 1)      # 1 x k^2 x 2 x 1

    # Calcualte Gaussian for every pixel of the kernel
    ZZ = Z - center                                                        # 1 x k^2 x 2 x 1
    ZZ_t = ZZ.permute(0, 1, 3, 2)                                          # 1 x k^2 x 1 x 2
    ZZZ = -0.5 * ZZ_t.matmul(sigma_inv).matmul(ZZ).squeeze(-1).squeeze(-1) # N x k^2
    kernel = F.softmax(ZZZ, dim=1)                                         # N x k^2

    return kernel.view(-1, 1, k_size, k_size)                # N x 1 x k x k

def shifted_anisotropic_Gaussian(k_size=21, sf=4, lambda_1=1.2, lambda_2=5., theta=0, shift=True):
    '''
    # modified version of https://github.com/cszn/USRNet/blob/master/utils/utils_sisr.py
    '''
    # set covariance matrix
    Lam = np.diag([lambda_1, lambda_2])
    U = np.array([[np.cos(theta), -np.sin(theta)],
                    [np.sin(theta),  np.cos(theta)]])
    sigma = U @ Lam @ U.T                                 # 2 x 2
    inv_sigma = np.linalg.inv(sigma)[None, None, :, :]    # 1 x 1 x 2 x 2

    # set expectation position (shifting kernel for aligned image)
    if shift:
        center = k_size // 2 + 0.5*(sf - k_size % 2)
    else:
        center = k_size // 2

    # Create meshgrid for Gaussian
    X, Y = np.meshgrid(range(k_size), range(k_size))
    Z = np.stack([X, Y], 2).astype(np.float32)[:, :, :, None]                  # k x k x 2 x 1

    # Calcualte Gaussian for every pixel of the kernel
    ZZ = Z - center
    ZZ_t = ZZ.transpose(0,1,3,2)
    ZZZ = -0.5 * np.squeeze(ZZ_t @ inv_sigma  @ ZZ).reshape([1, -1])
    kernel = softmax(ZZZ, axis=1).reshape([k_size, k_size]) # k x k

    # The convariance of the marginal distributions along x and y axis
    s1, s2 = sigma[0, 0], sigma[1, 1]
    # Pearson corrleation coefficient
    rho = sigma[0, 1] / (math.sqrt(s1) * math.sqrt(s2))
    kernel_infos = np.array([s1, s2, rho])   # (3,)

    return kernel, kernel_infos

#------------------------------------------Degradation-------------------------------------------
def imconv_np(im, kernel, padding_mode='reflect', correlate=False):
    '''
    Image convolution or correlation.
    Input:
        im: h x w x c numpy array
        kernel: k x k numpy array
        padding_mode: 'reflect', 'constant' or 'wrap'
    '''
    if kernel.ndim != im.ndim: kernel = kernel[:, :, np.newaxis]

    if correlate:
        out = snd.correlate(im, kernel, mode=padding_mode)
    else:
        out = snd.convolve(im, kernel, mode=padding_mode)

    return out

def conv_multi_kernel_tensor(im_hr, kernel, sf, downsampler):
    '''
    Degradation model by Pytorch.
    Input:
        im_hr: N x c x h x w
        kernel: N x 1 x k x k
        sf: scale factor
    '''
    im_hr_pad = F.pad(im_hr, (kernel.shape[-1] // 2,)*4, mode='reflect')
    im_blur = F.conv3d(im_hr_pad.unsqueeze(0), kernel.unsqueeze(1), groups=im_hr.shape[0])
    if downsampler.lower() == 'direct':
        im_blur = im_blur[0, :, :, ::sf, ::sf]      # N x c x ...
    elif downsampler.lower() == 'bicubic':
        im_blur = resize(im_blur, scale_factors=1/sf)
    else:
        sys.exit('Please input the corrected downsampler: Direct or Bicubic!')

    return im_blur

def tidy_kernel(kernel, expect_size=21):
    '''
    Input:
        kernel: p x p numpy array
    '''
    k_size = kernel.shape[-1]
    kernel_new = np.zeros([expect_size, expect_size], dtype=kernel.dtype)
    if expect_size >= k_size:
        start_ind = expect_size // 2 - k_size // 2
        end_ind = start_ind + k_size
        kernel_new[start_ind:end_ind, start_ind:end_ind] = kernel
    elif expect_size < k_size:
        start_ind = k_size // 2 - expect_size // 2
        end_ind = start_ind + expect_size
        kernel_new = kernel[start_ind:end_ind, start_ind:end_ind]
        kernel_new /= kernel_new.sum()

    return kernel_new

def shift_pixel(x, sf, upper_left=True):
    """shift pixel for super-resolution with different scale factors
    Args:
        x: WxHxC or WxH
        sf: scale factor
        upper_left: shift direction
    """
    h, w = x.shape[:2]
    shift = (sf-1)*0.5
    xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
    if upper_left:
        x1 = xv + shift
        y1 = yv + shift
    else:
        x1 = xv - shift
        y1 = yv - shift

    x1 = np.clip(x1, 0, w-1)
    y1 = np.clip(y1, 0, h-1)

    if x.ndim == 2:
        x = interp2d(xv, yv, x)(x1, y1)
    if x.ndim == 3:
        for i in range(x.shape[-1]):
            x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)

    return x

#-----------------------------------------Transform--------------------------------------------
class Bicubic:
    def __init__(self, scale=0.25):
        self.scale = scale

    def __call__(self, im, scale=None, out_shape=None):
        scale = self.scale if scale is None else scale
        out = resize(im, scale_factors=scale, out_shape=None)
        return out