import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class Loss:
    def __new__(cls, criterion, args):
        if criterion == 'CE':
            return nn.CrossEntropyLoss()
        elif criterion == 'FL':
            return FocalLoss()
        elif criterion == 'BCE':
            weight = args.BCE_weight
            weight = torch.tensor(weight)
            # print(weight)
            return nn.BCEWithLogitsLoss(pos_weight=weight)
        elif criterion == 'MSE':
            return nn.MSELoss()
        elif criterion == 'CD':
            return Contrast_depth_loss(args)
        elif criterion == 'BCELOSS':
            return nn.BCELoss()


class FocalLoss(nn.Module):
    def __init__(self, gamma=1, alpha=None, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        #if isinstance(alpha,(float,int,long)): self.alpha = torch.Tensor([alpha,1-alpha])
        #if isinstance(alpha,list): self.alpha = torch.Tensor(alpha)
        self.size_average = size_average

    def forward(self, input, target):
        if input.dim() > 2:
            # N,C,H,W => N,C,H*W
            input = input.view(input.size(0), input.size(1), -1)
            input = input.transpose(1, 2)    # N,C,H*W => N,H*W,C
            input = input.contiguous().view(-1, input.size(2))   # N,H*W,C => N*H*W,C
        target = target.view(-1, 1)

        logpt = F.log_softmax(input, dim=-1)
        logpt = logpt.gather(1, target)
        logpt = logpt.view(-1)
        pt = Variable(logpt.data.exp())

        if self.alpha is not None:
            if self.alpha.type() != input.data.type():
                self.alpha = self.alpha.type_as(input.data)
            at = self.alpha.gather(0, target.data.view(-1))
            logpt = logpt * Variable(at)

        loss = -1 * (1-pt)**self.gamma * logpt
        if self.size_average:
            return loss.mean()
        else:
            return loss.sum()


# Pearson range [-1, 1] so if < 0, abs|loss| ; if >0, 1- loss
class Contrast_depth_loss(nn.Module):
    def __init__(self, args):
        super(Contrast_depth_loss, self).__init__()
        self.device = args.device
        return

    def forward(self, out, label):
        '''
        compute contrast depth in both of (out, label),
        then get the loss of them
        tf.atrous_convd match tf-versions: 1.4
        '''
        contrast_out = self.contrast_depth_conv(out)
        contrast_label = self.contrast_depth_conv(label)

        criterion_MSE = nn.MSELoss().to(self.device)

        loss = criterion_MSE(contrast_out, contrast_label)
        #loss = torch.pow(contrast_out - contrast_label, 2)
        #loss = torch.mean(loss)

        return loss


    def contrast_depth_conv(self, input): 
        ''' compute contrast depth in both of (out, label) '''
        '''
            input  32x32
            output 8x32x32
        '''

        kernel_filter_list =[
                            [[1,0,0],[0,-1,0],[0,0,0]], [[0,1,0],[0,-1,0],[0,0,0]], [[0,0,1],[0,-1,0],[0,0,0]],
                            [[0,0,0],[1,-1,0],[0,0,0]], [[0,0,0],[0,-1,1],[0,0,0]],
                            [[0,0,0],[0,-1,0],[1,0,0]], [[0,0,0],[0,-1,0],[0,1,0]], [[0,0,0],[0,-1,0],[0,0,1]]
                            ]
        
        kernel_filter = np.array(kernel_filter_list, np.float32)

        kernel_filter = torch.from_numpy(kernel_filter.astype(np.float)).float().to(self.device)
        # weights (in_channel, out_channel, kernel, kernel)
        kernel_filter = kernel_filter.unsqueeze(dim=1)

        input = input.unsqueeze(dim=1).expand(input.shape[0], 8, input.shape[1], input.shape[2])

        contrast_depth = F.conv2d(input, weight=kernel_filter, groups=8)  # depthwise conv

        return contrast_depth