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


BATCH_SIZE = 4096
NUM_CLASSES = 1000 # ImageNet-1K
SHAPE = (BATCH_SIZE, NUM_CLASSES)

EPSILON = 2.0 
REDUCTION = 'none'

class Poly1CrossEntropyLoss(nn.Module):
    """
    PolyLoss: A Polynomial Expansion Perspective of Classification Loss Functions (ICLR 2022)
    Equation: L = -log(Pt) + epsilon * (1 - Pt)
    """
    def __init__(self, epsilon=2.0, reduction='mean'):
        super(Poly1CrossEntropyLoss, self).__init__()
        self.epsilon = epsilon
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # labels: (N) long
        
        # 1. 计算 Pt (Target Probabilities)
        probs = F.softmax(logits, dim=-1)
        
        pt = probs.gather(1, labels.unsqueeze(1)).squeeze(1)
        
        # 2. 计算 Cross Entropy: -log(Pt)
        ce_loss = F.cross_entropy(logits, labels, reduction='none')
        
        # 3. Poly-1 Term: eps * (1 - Pt)
        poly1 = self.epsilon * (1 - pt)
        
        # 4. Final Loss
        loss = ce_loss + poly1
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, epsilon=2.0, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = Poly1CrossEntropyLoss(epsilon=epsilon, reduction=reduction)
    
    def forward(self, logits, labels):
        return self.loss_fn(logits, labels)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    labels = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), labels.contiguous()]

def get_init_inputs():
    return [EPSILON, REDUCTION]