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

BATCH_SIZE = 512
NUM_CLASSES = 10000 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

# ArcFace 超参数
SCALE_S = 64.0
MARGIN_M = 0.50

class ArcFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of ArcFace.
    """
    def __init__(self, s=64.0, m=0.50, reduction='mean'):
        super(ArcFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction
        
        self.cos_m = math.cos(m)
        self.sin_m = math.sin(m)
        # 阈值，防止 acos 出界
        self.th = math.cos(math.pi - m)
        self.mm = math.sin(math.pi - m) * m

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C)
        # label: (N)
        
        # 1. 提取 target 对应的 cosine 值
        pick_cosine = cosine.gather(1, label.view(-1, 1)).squeeze(1)
        
        # 2. 计算 theta: acos(x)
        theta = torch.acos(torch.clamp(pick_cosine, -1.0 + 1e-7, 1.0 - 1e-7))
        
        # 3. 加 margin: cos(theta + m)
        target_logit = torch.cos(theta + self.m)
        
        # 4. 替换回原矩阵 (这里通常非常慢)
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # output = (1-one_hot) * cosine + one_hot * target_logit
        logits = cosine * (1.0 - one_hot) + target_logit.unsqueeze(1) * one_hot
        
        # 5. Scale
        logits = logits * self.s
        
        # 6. 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.50, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = ArcFaceLoss(s=s, m=m, reduction=reduction)
    
    def forward(self, cosine, label):
        return self.loss_fn(cosine, label)

def get_inputs():
    cosine = torch.randn(SHAPE, dtype=torch.float32)
    cosine = torch.clamp(cosine, -0.99, 0.99)
    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']