import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable, Function

class CrossEntropyLoss(nn.Module):
    def __init__(self, weight=None, ignore_index=-100, reduction='mean'):
        super(CrossEntropyLoss, self).__init__()
        self.weight = weight
        if self.weight is not None:
            self.weight = torch.FloatTensor(self.weight)
        self.ignore_index = ignore_index
        self.reduction = reduction
    
    def forward(self, input, target):
        return F.cross_entropy(input, target, 
                               weight=self.weight,
                               ignore_index=self.ignore_index, 
                               reduction=self.reduction)

class FocalLoss(nn.Module):
    def __init__(self, gamma=2, alpha=0.25, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        if isinstance(alpha,(float,int)): 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:
            input = input.view(input.size(0),input.size(1),-1)  # N,C,H,W => N,C,H*W
            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()
        
class TVLoss(nn.Module):
    """
    Total variation loss.
    """
    def __init__(self):
        super(TVLoss, self).__init__()

    def forward(self, yhat, y):
        bsize, chan, height, width = y.size()
        errors = []
        dy = torch.abs(y[:,:,1:,:] - y[:,:,:-1,:])
        dyhat = torch.abs(yhat[:,:,1:,:] - yhat[:,:,:-1,:])
        error = torch.norm(dy - dyhat, 1)
        return error / height