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

# 定义常量
BATCH_SIZE = 256
FEATURE_DIM = 512
MARGIN = 0.25
GAMMA = 256


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.margin = MARGIN
        self.gamma = GAMMA

    def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # Circle Loss的 PyTorch 实现
        # 假设 features 已经是 L2 归一化的
        # (B, D) @ (D, B) -> (B, B)
        similarities = torch.matmul(features, features.t())

        # 创建正负样本对的掩码
        mask_positive = labels.unsqueeze(1) == labels.unsqueeze(0)
        mask_negative = labels.unsqueeze(1) != labels.unsqueeze(0)

        # 收集正样本对和负样本对的相似度
        # .masked_select() 会将张量展平
        sp = similarities[mask_positive]
        sn = similarities[mask_negative]

        # 计算 Circle Loss 的 logits
        # .detach() 用于停止梯度反向传播
        ap = torch.clamp_min(-sp.detach() + 1 + self.margin, min=0.)
        an = torch.clamp_min(sn.detach() + self.margin, min=0.)

        delta_p = 1 - self.margin
        delta_n = self.margin

        logit_p = -ap * (sp - delta_p) * self.gamma
        logit_n = an * (sn - delta_n) * self.gamma

        # 使用 logsumexp 和 softplus 计算最终的 loss
        # 这是 "unified" 版本的 loss
        loss = F.softplus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0))

        return loss


def get_inputs():
    # 特征需要 L2 归一化
    features = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
    labels = torch.randint(0, 10, (BATCH_SIZE,), dtype=torch.long)  # 假设有 10 个类别
    return [features, labels]


def get_init_inputs():
    return []