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

class SupConLoss(nn.Module):
    def __init__(self, temperature=0.07, loss_weight=1.0):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.loss_weight = loss_weight

    def forward(self, features, labels):
        """ 
        features: [bsz, n_views, feat_dim]
        labels: [bsz]
        """
        device = features.device
        bsz = features.shape[0]
        labels = labels.contiguous().view(-1, 1)  # [bsz, 1]
        mask = torch.eq(labels, labels.T).float().to(device)  # [bsz, bsz]

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)  # [bsz * n_views, dim]

        anchor_feature = contrast_feature
        anchor_count = contrast_count

        # cosine similarity matrix
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature
        )

        # log prob
        logits_mask = torch.ones_like(anchor_dot_contrast) - torch.eye(bsz * anchor_count, device=device)
        mask = mask.repeat(anchor_count, contrast_count)
        logits_mask = logits_mask * mask

        exp_logits = torch.exp(anchor_dot_contrast) * logits_mask
        log_prob = anchor_dot_contrast - torch.log(exp_logits.sum(1, keepdim=True) + 1e-12)

        mean_log_prob_pos = (log_prob * mask).sum(1) / mask.sum(1)

        loss = -mean_log_prob_pos
        loss = loss.view(anchor_count, bsz).mean()

        return loss * self.loss_weight
