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)

# SphereFace 超参数
SCALE_S = 64.0
MARGIN_M = 4 

class SphereFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of SphereFace (A-Softmax).
    psi(theta) = (-1)^k * cos(m*theta) - 2k
    """
    def __init__(self, s=64.0, m=4, reduction='mean'):
        super(SphereFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction
        self.pi = math.pi

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C) should be clamped to (-1, 1)
        # label: (N)
        
        # 1. Gather target cosines: cos(theta_yi)
        target_cosine = cosine.gather(1, label.view(-1, 1)).squeeze(1)
        
        # 2. Compute theta: acos(x)
        # clamp to avoid nan
        input_theta = torch.acos(torch.clamp(target_cosine, -1.0 + 1e-7, 1.0 - 1e-7))
        
        # 3. Compute k
        # theta in [k*pi/m, (k+1)*pi/m]
        # k = floor(theta * m / pi)
        k = (input_theta * self.m / self.pi).floor()
        
        # 4. Compute psi(theta)
        # (-1)^k * cos(m*theta) - 2k
        minus_one_pow_k = torch.pow(-1, k)
        cos_m_theta = torch.cos(input_theta * self.m)
        psi = minus_one_pow_k * cos_m_theta - 2 * k
        
        # 5. Scatter back to logits
        # In SphereFace, logits = s * cosine for j!=yi, and s * psi for j==yi
        
        # Create a copy or use scatter
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # logits = (1 - one_hot) * cosine + one_hot * psi
        logits = cosine * (1.0 - one_hot) + psi.unsqueeze(1) * one_hot
        
        # 6. Scale
        logits = logits * self.s
        
        # 7. 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=4, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = SphereFaceLoss(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']