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

mce_loss = nn.MSELoss()


def channel_1toN(img, num_channel):
    T = torch.LongTensor(num_channel, img.shape[1], img.shape[2]).zero_()
    mask = torch.LongTensor(img.shape[1], img.shape[2]).zero_()
    for i in range(num_channel):
        T[i] = T[i] + i
        layer = T[i] - img
        T[i] = torch.from_numpy(np.logical_not(np.logical_xor(layer.numpy(), mask.numpy())).astype(int))
    return T.float()


class WeightedBCEWithLogitsLoss(nn.Module):
    
    def __init__(self, size_average=True):
        super(WeightedBCEWithLogitsLoss, self).__init__()
        self.size_average = size_average
        
    def weighted(self, input, target, weight, alpha, beta, weight_region):
        if not (target.size() == input.size()):
            raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))

        max_val = (-input).clamp(min=0)
        loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log()
                
        if weight is not None:
            loss = alpha * loss + beta * loss * weight


        if weight_region is not None:
            if not (loss.size() == weight_region.size()):
                raise ValueError("loss size ({}) must be the same as weight_region size ({})".format(loss.size(), weight_region.size()))
            loss = loss * weight_region
           # num = len(np.where(weight_region.detach().cpu().numpy()==1)[0])
          #  num = len(np.where(weight_region.detach().cpu().numpy()>0.95)[0])
         #   print('weight_region loss function', weight_region)
        #    print('weight_region loss function.shape', weight_region.shape)
            return loss.sum()/weight_region.sum()  
        
        # if self.size_average:
            # return loss.mean()
        # else:
            # return loss.sum()
    
    def forward(self, input, target, weight, alpha, beta, weight_region = None):
        if weight is not None:
            return self.weighted(input, target, weight, alpha, beta, weight_region)
        else:
            return self.weighted(input, target, None, alpha, beta, weight_region)


###yuanlai focal loss
class CrossEntropy2d(nn.Module):
    
    def __init__(self, class_num, alpha=None, gamma=2, size_average=True, ignore_label=255):
        super(CrossEntropy2d, self).__init__()
        if alpha is None:
            self.alpha = Variable(torch.ones(class_num, 1))
        else:
            if isinstance(alpha, Variable):
                self.alpha = alpha
            else:
                self.alpha = Variable(alpha)
        self.gamma = gamma
        self.class_num = class_num
        self.size_average = size_average
        self.ignore_label = ignore_label

    def forward(self, predict, target):
        N, C, H, W = predict.size()
        sm = nn.Softmax2d()
        
        P = sm(predict)
        P = torch.clamp(P, min = 1e-9, max = 1-(1e-9))
        
        target_mask = (target >= 0) * (target != self.ignore_label)
        target = target[target_mask].view(1, -1)
        predict = P[target_mask.view(N, 1, H, W).repeat(1, C, 1, 1)].view(C, -1)
        probs = torch.gather(predict, dim = 0, index = target)
        log_p = probs.log()
        batch_loss = -(torch.pow((1-probs), self.gamma))*log_p 

        if self.size_average:
            loss = batch_loss.mean()
        else:
            
            loss = batch_loss.sum()
        return loss

## my cross entropy loss      

class CrossEntropyLoss2d(nn.Module):
    def __init__(self, weight=None, reduction = 'mean'): # size_average=True
        super(CrossEntropyLoss2d, self).__init__()
       # self.nll_loss = nn.NLLLoss2d(weight, size_average)
        self.nll_loss = nn.NLLLoss(weight, reduction) #size_average

    def forward(self, inputs, targets):
        return self.nll_loss(F.log_softmax(inputs, dim=1), targets)