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

BATCH_SIZE = 2048
NUM_CLASSES = 2048
SHAPE = (BATCH_SIZE, NUM_CLASSES)

Q_VALUE = 0.7
REDUCTION = 'mean'

class GCELoss(nn.Module):
    """
    Generalized Cross Entropy Loss
    L = (1 - Pt^q) / q
    """
    def __init__(self, q=0.7, reduction='mean'):
        super(GCELoss, self).__init__()
        self.q = q
        self.reduction = reduction
        self.epsilon = 1e-7

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        probs = F.softmax(logits, dim=1)
        probs_t = probs.gather(1, targets.unsqueeze(1)).squeeze(1)
        
        probs_t = probs_t.clamp(min=self.epsilon, max=1.0)
        loss = (1.0 - probs_t.pow(self.q)) / self.q
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, q=0.7, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = GCELoss(q=q, reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

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

def get_init_inputs():
    return [Q_VALUE, REDUCTION]