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

# 人脸识别场景：Batch Size 较小，但分类数 (Classes/Identities) 极大
BATCH_SIZE = 512
NUM_CLASSES = 10000 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

# CosFace 超参数
SCALE_S = 64.0
MARGIN_M = 0.35

class CosFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of CosFace Loss.
    """
    def __init__(self, s=64.0, m=0.35, reduction='mean'):
        super(CosFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C) - Normalized Features @ Normalized Weights
        # label: (N)
        
        # 1. 创建 One-hot 掩码 
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # 2. Apply Margin: cos_theta - m (only for target class)
        logits = cosine - one_hot * self.m
        
        # 3. Scale: s * logits
        logits = logits * self.s
        
        # 4. CrossEntropy
        loss = F.cross_entropy(logits, label, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, s=64.0, m=0.35, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = CosFaceLoss(s=s, m=m, reduction=reduction)
    
    def forward(self, cosine, label):
        return self.loss_fn(cosine, label)

def get_inputs():
    # 模拟归一化后的 Cosine Similarity (-1 ~ 1)
    cosine = torch.randn(SHAPE, dtype=torch.float32)
    # 归一化到 [-1, 1] 模拟真实余弦值
    cosine = torch.clamp(cosine, -1.0, 1.0)
    
    label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [cosine.contiguous(), label.contiguous()]

def get_init_inputs():
    return [SCALE_S, MARGIN_M, 'none']