import torch
from torch import nn
import numpy as np

def cc(img1, img2):
    eps = torch.finfo(torch.float32).eps
    """Correlation coefficient for (N, C, H, W) image; torch.float32 [0.,1.]."""
    N, C, _, _ = img1.shape
    img1 = img1.reshape(N, C, -1)
    img2 = img2.reshape(N, C, -1)
    img1 = img1 - img1.mean(dim=-1, keepdim=True)
    img2 = img2 - img2.mean(dim=-1, keepdim=True)
    cc = torch.sum(img1 * img2, dim=-1) / (eps + torch.sqrt(torch.sum(
        img1 **2, dim=-1)) * torch.sqrt(torch.sum(img2**2, dim=-1)))
    cc = torch.clamp(cc, -1., 1.)
    return cc.mean()


def gradient_calllback(network):
    for name, param in network.named_parameters():
        if param.grad is not None:
            # print("Gradient of {}: {}".format(name, param.grad.abs().mean()))

            if param.grad.abs().mean() == 0:
                print("Gradient of {} is 0".format(name))

        else:
            print("Gradient of {} is None".format(name))


def bright(x, a,b):
    # input datatype np.uint8
    x = np.array(x, dtype='float')
    x = x/(b-a) - 255*a/(b-a)
    x[x>255.0] = 255.0
    x[x<0.0] = 0.0
    x = x.astype(np.uint8)
    return x

def trunc(x):
    # input datatype float
    x[x>255.0] = 255.0
    x[x<0.0] = 0.0
    return x





def gradient_calllback(network):
    for name, param in network.named_parameters():
        if param.grad is not None:
            if param.grad.abs().mean() == 0:
                print("Gradient of {} is 0".format(name))

        else:
            print("Gradient of {} is None".format(name))

class Frequency_Loss(nn.Module):
    def __init__(self):
        super(Frequency_Loss, self).__init__()
        self.cri = nn.L1Loss()
        self.cri_sum = nn.L1Loss(reduction="sum")

    def forward(self, x, y, mask=None):
        x = torch.fft.fftshift(torch.fft.fft2(x))       # rfft2
        y = torch.fft.fftshift(torch.fft.fft2(y))



        # def apply_tofre(x_start, mask):
        #     # B, C, H, W = x_start.shape
        #     kspace = fftshift(fft2(x_start, norm=None, dim=(-2, -1)), dim=(-2, -1))  # Default: all dimensions
        #     mask = mask.to(kspace.device)
        #     return kspace, mask

        x_mag = torch.abs(x)
        y_mag = torch.abs(y)
        x_ang = torch.angle(x)
        y_ang = torch.angle(y)
        if isinstance(mask, type(None)):
            return self.cri(x_mag,y_mag) + self.cri(x_ang, y_ang)

        k = (1 - mask.to(x.device)).detach()
        # W = x.shape[-1]
        # k = k[..., :W // 2 + 1]
        k_total = torch.sum(k)

        x_mag = x_mag * k
        y_mag = y_mag * k
        x_ang = x_ang * k
        y_ang = y_ang * k

        # Compute L1 loss between magnitudes
        return self.cri_sum(x_mag, y_mag) / k_total + self.cri_sum(x_ang, y_ang) / k_total