from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F

class LovaszConLoss(nn.Module):
    
    def __init__(self, temperature=1.0, contrast_mode='all',
                 base_temperature=1.0, pos_reg_base=1.0, neg_reg_base=1.0, sim_mat=None, blend=1.0, stability=True):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature
        self.pos_reg_base = pos_reg_base
        self.neg_reg_base = neg_reg_base
        self.blend = blend
        self.stability = stability
        self.sim_mat = ((1-self.blend) * torch.eye(sim_mat.shape[0]).cuda() + self.blend * sim_mat.cuda()) if sim_mat is not None else None

    def forward(self, features, labels=None, mask=None, debug=True):
        
        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)

        """
        Feature normalization
        """
        features = features / torch.linalg.norm(features, dim=-1, keepdim=True)


        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.matmul(anchor_feature, contrast_feature.T)
        logits = anchor_dot_contrast


        # 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

       
        if self.sim_mat is None:
            logits = torch.div(logits, self.temperature)            

            logits_pos = logits * logits_mask * mask
            exp_logits_neg = torch.exp(logits) * logits_mask * (1-mask)

            loss_pos = -(logits_pos.sum(dim = 1))/(logits_mask * mask).sum(dim=1)
            loss_neg = torch.log(exp_logits_neg.sum(dim=1))

            loss = loss_pos.mean() + loss_neg.mean()

        else:
            logits_pos = torch.div(logits, self.temperature) * logits_mask * mask

            temp_labels = torch.cat([labels, labels], dim=0).flatten()
            similarities = self.sim_mat[temp_labels,:][:,temp_labels]
            
            # Same calculation as before, but incorporate the similarity matrix
            logits_neg = ((logits-similarities)/(1-similarities+1e-6))*logits_mask*(1-mask)
            logits_neg = torch.div(logits_neg, self.temperature)
            if self.stability:
                max_logits = torch.max(logits_neg, dim=1, keepdim=True)[0].detach()
                logits_neg = logits_neg - max_logits
                logits_pos = logits_pos - max_logits

            exp_logits_neg = torch.exp(logits_neg) * logits_mask * (1-mask)
            loss_pos = -(logits_pos.sum(dim = 1))/(logits_mask * mask).sum(dim=1)
            loss_neg = torch.log(exp_logits_neg.sum(dim=1))
            loss = loss_pos.mean() + loss_neg.mean()
 
        
        return loss
