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


def focal_loss(input_values, gamma):
    """Computes the focal loss"""
    p = torch.exp(-input_values)
    loss = (1 - p) ** gamma * input_values
    return loss.mean()


class FocalLoss(nn.Module):
    def __init__(self, gamma=0.):
        super(FocalLoss, self).__init__()
        assert gamma >= 0
        self.gamma = gamma

    def forward(self, input, target):
        return focal_loss(F.cross_entropy(input, target, reduction='none'), self.gamma)

class PHuberCE(nn.Module):
    def __init__(self, tau=10):
        super(PHuberCE, self).__init__()
        self.tau = tau

        # Probability threshold for the clipping
        self.prob_thresh = 1 / self.tau
        # Negative of the Fenchel conjugate of base loss at tau
        self.boundary_term = math.log(self.tau) + 1

        self.softmax = nn.Softmax(dim=1)

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        p = self.softmax(input)
        p = p[torch.arange(p.shape[0]), target]

        loss = torch.empty_like(p)
        clip = p <= self.prob_thresh
        loss[clip] = -self.tau * p[clip] + self.boundary_term
        loss[~clip] = -torch.log(p[~clip])

        return torch.mean(loss)

def loss_sce(y, labels_one_hot, alpha, beta):

    pred = F.softmax(y, dim=1)

    pred = torch.clamp(pred, min=1e-7, max=1.0)
    label_one_hot = torch.clamp(labels_one_hot, min=1e-4, max=1.0)

    ce = (-1 * torch.sum(label_one_hot * torch.log(pred), dim=1)).mean()
    rce = (-1 * torch.sum(pred * torch.log(label_one_hot), dim=1)).mean()
    loss = alpha * ce + beta * rce

    return loss

class SCE(nn.Module):
    def __init__(self, alpha=0.5, beta=1.0, num_classes=10):
        super(SCE, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.num_classes = num_classes

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        labels_one_hot = torch.zeros(target.shape[0], self.num_classes).to(input.device).scatter_(1,
                                                                                                 target.unsqueeze(
                                                                                                     1), 1)
        loss = loss_sce(input, labels_one_hot, self.alpha, self.beta)
        return loss



class GCE(nn.Module):
    def __init__(self, device, q=0.7, k=10):
        super(GCE, self).__init__()
        self.q = q
        self.k = k
        self.device = device

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        soft_max = nn.Softmax(dim=1)
        sm_outputs = soft_max(input)
        label_one_hot = nn.functional.one_hot(target, self.k).float().to(self.device)
        sm_out = torch.pow((sm_outputs * label_one_hot).sum(dim=1), self.q)
        target = torch.ones_like(target)
        loss_vec = (target - sm_out) / self.q
        average_loss = loss_vec.mean()
        return average_loss

class AGCELoss(nn.Module):
    def __init__(self, num_classes=10, a=1, q=2, eps=1e-7, scale=1.):
        super(AGCELoss, self).__init__()
        self.a = a
        self.q = q
        self.num_classes = num_classes
        self.eps = eps
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.softmax(pred, dim=1)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        loss = ((self.a+1)**self.q - torch.pow(self.a + torch.sum(label_one_hot * pred, dim=1), self.q)) / self.q
        return loss.mean() * self.scale


class NGCE(torch.nn.Module):
    def __init__(self, num_classes, scale=1.0, q=0.7):
        super(NGCE, self).__init__()
        self.num_classes = num_classes
        self.q = q
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.softmax(pred, dim=1)
        pred = torch.clamp(pred, min=1e-7, max=1.0)
        label_one_hot = torch.nn.functional.one_hot(labels, self.num_classes).float().to(pred.device)
        numerators = 1. - torch.pow(torch.sum(label_one_hot * pred, dim=1), self.q)
        denominators = self.num_classes - pred.pow(self.q).sum(dim=1)
        ngce = numerators / denominators
        return self.scale * ngce.mean()


class TaylorCE(nn.Module):
    def __init__(self, device, series=2):
        super(TaylorCE, self).__init__()
        self.series = series
        self.device = device

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        n = self.series
        k = input.shape[1]
        soft_max = nn.Softmax(dim=1)
        sm_outputs = soft_max(input)
        label_one_hot = nn.functional.one_hot(target, k).float().to(self.device)
        final_outputs = (sm_outputs * label_one_hot).sum(dim=1)

        total_loss = 0
        for i in range(n):  # 0 to n-1
            total_loss += torch.pow(torch.tensor([-1.0]).to(self.device), i + 1) * torch.pow(final_outputs - 1, i + 1) * 1.0 / (
                    i + 1)  # \sum_i=0^n(x-1)^(i+1)*(-1)^(i+1)/(i+1)
        average_loss = total_loss.mean()
        return average_loss



class NCELoss(nn.Module):
    def __init__(self, num_classes, scale=1.0):
        super(NCELoss, self).__init__()
        self.num_classes = num_classes
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.log_softmax(pred, dim=1)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        loss = -1 * torch.sum(label_one_hot * pred, dim=1) / (-pred.sum(dim=1))
        return self.scale * loss.mean()

class NCEandAGCE(torch.nn.Module):
    def __init__(self, alpha=1., beta = 1., num_classes=10, a=3, q=1.5):
        super(NCEandAGCE, self).__init__()
        self.num_classes = num_classes
        self.nce = NCELoss(num_classes=num_classes, scale=alpha)
        self.agce = AGCELoss(num_classes=num_classes, a=a, q=q, scale=beta)

    def forward(self, pred, labels):
        return self.nce(pred, labels) + self.agce(pred, labels)

class MAELoss(nn.Module):
    def __init__(self, num_classes=10, scale=2.0):
        super(MAELoss, self).__init__()
        self.num_classes = num_classes
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.softmax(pred, dim=1)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        loss = 1. - torch.sum(label_one_hot * pred, dim=1)
        return self.scale * loss.mean()

class NCEandMAE(nn.Module):
    def __init__(self, alpha=1., beta=1., num_classes=10):
        super(NCEandMAE, self).__init__()
        self.num_classes = num_classes
        self.nce = NCELoss(num_classes=num_classes, scale=alpha)
        self.mae = MAELoss(num_classes=num_classes, scale=beta)

    def forward(self, pred, labels):
        return self.nce(pred, labels) + self.mae(pred, labels)

class NGCEandMAE(torch.nn.Module):
    def __init__(self, alpha, beta, num_classes, q=0.7):
        super(NGCEandMAE, self).__init__()
        self.num_classes = num_classes
        self.ngce = NGCE(num_classes=num_classes, scale=alpha, q=q)
        self.mae = MAELoss(scale=beta, num_classes=num_classes)

    def forward(self, pred, labels):
        return self.ngce(pred, labels) + self.mae(pred, labels)


class AUELoss(nn.Module):
    def __init__(self, num_classes=10, a=5.5, q=3, scale=1.0):
        super(AUELoss, self).__init__()
        self.num_classes = num_classes
        self.a = a
        self.q = q
        self.eps = 1e-7
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.softmax(pred, dim=1)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        loss = (torch.pow(self.a - torch.sum(label_one_hot * pred, dim=1), self.q) - (self.a-1)**self.q)/ self.q
        return loss.mean() * self.scale

class AExpLoss(torch.nn.Module):
    def __init__(self, num_classes=10, a=2.5, scale=1.0):
        super(AExpLoss, self).__init__()
        self.num_classes = num_classes
        self.a = a
        self.scale = scale

    def forward(self, pred, labels):
        pred = F.softmax(pred, dim=1)
        label_one_hot = F.one_hot(labels, self.num_classes).float().to(pred.device)
        loss = torch.exp(-torch.sum(label_one_hot * pred, dim=1) / self.a)
        return loss.mean() * self.scale



class CoresLoss(torch.nn.Module):
    def __init__(self, device):
        super(CoresLoss, self).__init__()
        self.device = device

    def forward(self, pred, labels, epoch):
        beta = f_beta(epoch)
        loss = F.cross_entropy(pred, labels, reduce=False)
        loss_numpy = loss.data.cpu().numpy()
        num_batch = len(loss_numpy)
        loss_v = np.zeros(num_batch)
        loss_ = -torch.log(F.softmax(pred) + 1e-8)
        # sel metric
        loss_sel = loss - torch.mean(loss_, 1)
        loss = loss - beta * torch.mean(loss_, 1)

        loss_div_numpy = loss_sel.data.cpu().numpy()

        for i in range(len(loss_numpy)):
            if epoch <= 60:
                loss_v[i] = 1.0
            elif loss_div_numpy[i] <= 0:
                loss_v[i] = 1.0
        loss_v = loss_v.astype(np.float32)
        loss_v_var = Variable(torch.from_numpy(loss_v)).to(self.device)
        loss_ = loss_v_var * loss
        if sum(loss_v) == 0.0:
            return torch.mean(loss_) / 100000000
        else:
            return torch.sum(loss_) / sum(loss_v)
# def loss_cores(epoch, y, t, class_list, ind, noise_or_not, loss_all, loss_div_all, noise_prior=None):

def f_beta(epoch):
    beta1 = np.linspace(0.0, 0.0, num=20)
    beta2 = np.linspace(0.0, 2, num=60)
    beta3 = np.linspace(2, 2, num=120)

    beta = np.concatenate((beta1, beta2, beta3), axis=0)
    return beta[epoch]