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


@torch.no_grad()
def mixup_data(X, y, alpha=0.2):
    """Returns mixed inputs, pairs of targets, and lambda
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    def mix(x):
        return lam * x + (1 - lam) * x.flip(dims=(0,))

    if isinstance(X, list):
        mixed_X = [mix(x) for x in X]
    else:
        mixed_X = mix(X)
     
    y_a, y_b = y, y.flip(dims=(0,))
    return mixed_X, y_a, y_b, lam


def mixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)



class FocalLoss(nn.Module):
    def __init__(self, num_classes, alpha=1, gamma=2):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.num_classes = num_classes

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        weight = torch.pow(-pred + 1.0, self.gamma)
        focal = -self.alpha * weight * torch.log(pred)
        loss_tmp = torch.sum(label_one_hot * focal, dim=1)

        return torch.mean(loss_tmp)


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

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        loss_tmp = 1 - torch.sum(label_one_hot * pred, dim=1)

        return torch.mean(loss_tmp)


class LQLoss(nn.Module):
    def __init__(self, num_classes, q=0.5):
        super(LQLoss, self).__init__()
        self.num_classes = num_classes
        self.q = q

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        loss_tmp = (1 - torch.sum(label_one_hot * pred, dim=1) ** self.q) / self.q

        return torch.mean(loss_tmp)


class SCELoss(torch.nn.Module):
    def __init__(self, num_classes, alpha=0.1, beta=1):
        super(SCELoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.num_classes = num_classes
        self.cross_entropy = torch.nn.CrossEntropyLoss()

    def forward(self, pred, labels):
        # CCE
        ce = self.cross_entropy(pred, labels)

        # RCE
        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)
        )
        label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0)
        rce = -1 * torch.sum(pred * torch.log(label_one_hot), dim=1)

        # Loss
        loss = self.alpha * ce + self.beta * rce.mean()
        return loss



@contextlib.contextmanager
def _disable_tracking_bn_stats(model):
    def switch_attr(m):
        if hasattr(m, "track_running_stats"):
            m.track_running_stats ^= True

    model.apply(switch_attr)
    yield
    model.apply(switch_attr)


def _l2_normalize(d):
    d_reshaped = d.view(d.shape[0], -1, *(1 for _ in range(d.dim() - 2)))
    d /= torch.norm(d_reshaped, dim=1, keepdim=True) + 1e-8
    return d


class VATLoss(nn.Module):
    def __init__(self, xi=10, eps=8, ip=1, alpha=1):
        """VAT loss
        :param xi: hyperparameter of VAT (default: 10.0)
        :param eps: hyperparameter of VAT (default: 1.0)
        :param ip: iteration times of computing adv noise (default: 1)
        """
        super(VATLoss, self).__init__()
        self.xi = xi
        self.eps = eps
        self.ip = ip
        self.alpha = alpha

    def forward(self, model, x, labels):
        with torch.no_grad():
            pred = F.softmax(model(x), dim=1)

        # prepare random unit tensor
        d = torch.rand(x.shape).sub(0.5).to(x.device)
        d = _l2_normalize(d)

        with _disable_tracking_bn_stats(model):
            # calc adversarial direction
            for _ in range(self.ip):
                d.requires_grad_()
                pred_hat = model(x + self.xi * d)
                logp_hat = F.log_softmax(pred_hat, dim=1)
                adv_distance = F.kl_div(logp_hat, pred, reduction="batchmean")
                adv_distance.backward()
                d = _l2_normalize(d.grad)
                model.zero_grad()

            # calc LDS
            r_adv = d * self.eps
            pred_hat = model(x + r_adv)
            logp_hat = F.log_softmax(pred_hat, dim=1)
            lds = F.kl_div(logp_hat, pred, reduction="batchmean")

        classification_loss = F.cross_entropy(model(x), labels)
        loss = classification_loss + self.alpha * lds
        return loss


class LabelSmoothingCrossEntropyLoss(nn.Module):
    def __init__(self, num_classes, smoothing=0.15, dim=-1):
        super(LabelSmoothingCrossEntropyLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.cls = num_classes
        self.dim = dim

    def forward(self, pred, target):
        pred = pred.log_softmax(dim=self.dim)
        with torch.no_grad():
            # true_dist = pred.data.clone()
            true_dist = torch.zeros_like(pred)
            true_dist.fill_(self.smoothing / (self.cls - 1))
            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))


class TransferLoss(nn.Module):

    def __init__(self, reg, leave_last=2, decay=False):
        super().__init__()
        self.reg = reg
        self.decay = decay
        self.leave_last = leave_last

    def forward(self, params1, params2):
        loss = 0
        reg = self.reg
        if isinstance(params1, nn.Module):
            params1 = list(params1.parameters())
        if isinstance(params2, nn.Module):
            params2 = list(params2.parameters())
        n_params = min(len(params1), len(params2))
        for i, (p1, p2) in enumerate(zip(params1, params2)):
            if i == n_params - self.leave_last:
                break
            loss += torch.pow(p1 - p2, 2).sum() * reg
            if self.decay:
                reg -= self.reg / n_params
        return loss


class PretrainedTransferLoss(TransferLoss):

    def __init__(self, model, reg, leave_last=2, decay=False):
        super().__init__(reg=reg, decay=decay)
        self.weight = [p.detach().clone() for p in model.parameters()]

    def cuda(self):
        self.weight = [w.cuda() for w in self.weight]

    def forward(self, model1, model2=None):
        return super().forward(model1, self.weight)


class SupConLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf

        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss

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


@torch.no_grad()
def mixup_data(X, y, alpha=0.2):
    """Returns mixed inputs, pairs of targets, and lambda
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    def mix(x):
        return lam * x + (1 - lam) * x.flip(dims=(0,))

    if isinstance(X, list):
        mixed_X = [mix(x) for x in X]
    else:
        mixed_X = mix(X)
     
    y_a, y_b = y, y.flip(dims=(0,))
    return mixed_X, y_a, y_b, lam


def mixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)



class FocalLoss(nn.Module):
    def __init__(self, num_classes, alpha=1, gamma=2):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.num_classes = num_classes

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        weight = torch.pow(-pred + 1.0, self.gamma)
        focal = -self.alpha * weight * torch.log(pred)
        loss_tmp = torch.sum(label_one_hot * focal, dim=1)

        return torch.mean(loss_tmp)


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

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        loss_tmp = 1 - torch.sum(label_one_hot * pred, dim=1)

        return torch.mean(loss_tmp)


class LQLoss(nn.Module):
    def __init__(self, num_classes, q=0.5):
        super(LQLoss, self).__init__()
        self.num_classes = num_classes
        self.q = q

    def forward(self, pred, labels):
        label_one_hot = (
            torch.nn.functional.one_hot(labels, self.num_classes)
            .float()
            .to(pred.device)
        )
        pred = F.softmax(pred, dim=1) + 1e-8
        loss_tmp = (1 - torch.sum(label_one_hot * pred, dim=1) ** self.q) / self.q

        return torch.mean(loss_tmp)


class SCELoss(torch.nn.Module):
    def __init__(self, num_classes, alpha=0.1, beta=1):
        super(SCELoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.num_classes = num_classes
        self.cross_entropy = torch.nn.CrossEntropyLoss()

    def forward(self, pred, labels):
        # CCE
        ce = self.cross_entropy(pred, labels)

        # RCE
        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)
        )
        label_one_hot = torch.clamp(label_one_hot, min=1e-4, max=1.0)
        rce = -1 * torch.sum(pred * torch.log(label_one_hot), dim=1)

        # Loss
        loss = self.alpha * ce + self.beta * rce.mean()
        return loss



@contextlib.contextmanager
def _disable_tracking_bn_stats(model):
    def switch_attr(m):
        if hasattr(m, "track_running_stats"):
            m.track_running_stats ^= True

    model.apply(switch_attr)
    yield
    model.apply(switch_attr)


def _l2_normalize(d):
    d_reshaped = d.view(d.shape[0], -1, *(1 for _ in range(d.dim() - 2)))
    d /= torch.norm(d_reshaped, dim=1, keepdim=True) + 1e-8
    return d


class VATLoss(nn.Module):
    def __init__(self, xi=10, eps=8, ip=1, alpha=1):
        """VAT loss
        :param xi: hyperparameter of VAT (default: 10.0)
        :param eps: hyperparameter of VAT (default: 1.0)
        :param ip: iteration times of computing adv noise (default: 1)
        """
        super(VATLoss, self).__init__()
        self.xi = xi
        self.eps = eps
        self.ip = ip
        self.alpha = alpha

    def forward(self, model, x, labels):
        with torch.no_grad():
            pred = F.softmax(model(x), dim=1)

        # prepare random unit tensor
        d = torch.rand(x.shape).sub(0.5).to(x.device)
        d = _l2_normalize(d)

        with _disable_tracking_bn_stats(model):
            # calc adversarial direction
            for _ in range(self.ip):
                d.requires_grad_()
                pred_hat = model(x + self.xi * d)
                logp_hat = F.log_softmax(pred_hat, dim=1)
                adv_distance = F.kl_div(logp_hat, pred, reduction="batchmean")
                adv_distance.backward()
                d = _l2_normalize(d.grad)
                model.zero_grad()

            # calc LDS
            r_adv = d * self.eps
            pred_hat = model(x + r_adv)
            logp_hat = F.log_softmax(pred_hat, dim=1)
            lds = F.kl_div(logp_hat, pred, reduction="batchmean")

        classification_loss = F.cross_entropy(model(x), labels)
        loss = classification_loss + self.alpha * lds
        return loss


class LabelSmoothingCrossEntropyLoss(nn.Module):
    def __init__(self, num_classes, smoothing=0.15, dim=-1):
        super(LabelSmoothingCrossEntropyLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.cls = num_classes
        self.dim = dim

    def forward(self, pred, target):
        pred = pred.log_softmax(dim=self.dim)
        with torch.no_grad():
            # true_dist = pred.data.clone()
            true_dist = torch.zeros_like(pred)
            true_dist.fill_(self.smoothing / (self.cls - 1))
            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))


class TransferLoss(nn.Module):

    def __init__(self, reg, leave_last=2, decay=False):
        super().__init__()
        self.reg = reg
        self.decay = decay
        self.leave_last = leave_last

    def forward(self, params1, params2):
        loss = 0
        reg = self.reg
        if isinstance(params1, nn.Module):
            params1 = list(params1.parameters())
        if isinstance(params2, nn.Module):
            params2 = list(params2.parameters())
        n_params = min(len(params1), len(params2))
        for i, (p1, p2) in enumerate(zip(params1, params2)):
            if i == n_params - self.leave_last:
                break
            loss += torch.pow(p1 - p2, 2).sum() * reg
            if self.decay:
                reg -= self.reg / n_params
        return loss


class PretrainedTransferLoss(TransferLoss):

    def __init__(self, model, reg, leave_last=2, decay=False):
        super().__init__(reg=reg, decay=decay)
        self.weight = [p.detach().clone() for p in model.parameters()]

    def cuda(self):
        self.weight = [w.cuda() for w in self.weight]

    def forward(self, model1, model2=None):
        return super().forward(model1, self.weight)


class SupConLoss(nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf

        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss

Returns mixed inputs, pairs of targets, and lambda # CCE # RCE # Loss VAT loss         :param xi: hyperparameter of VAT (default: 10.0)         :param eps: hyperparameter of VAT (default: 1.0)         :param ip: iteration times of computing adv noise (default: 1) # prepare random unit tensor # calc adversarial direction # calc LDS # true_dist = pred.data.clone() Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.     It also supports the unsupervised contrastive loss in SimCLR Compute loss for model. If both `labels` and `mask` are None,         it degenerates to SimCLR unsupervised loss:         https://arxiv.org/pdf/2002.05709.pdf          Args:             features: hidden vector of shape [bsz, n_views, ...].             labels: ground truth of shape [bsz].             mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j                 has the same class as sample i. Can be asymmetric.         Returns:             A loss scalar. # compute logits # for numerical stability # tile mask # mask-out self-contrast cases # compute log_prob # compute mean of log-likelihood over positive # loss