import torch
import torch.nn as nn

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

class SparsemaxLoss(nn.Module):
    """
    Sparsemax Loss (Martins & Astudillo, 2016)
    L = 0.5 * sum(z_j^2 - tau^2) + 0.5 - z_y
    """
    def __init__(self, reduction='mean'):
        super(SparsemaxLoss, self).__init__()
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N)
        
        # Sort (Descending) 
        z_sorted, _ = torch.sort(logits, dim=1, descending=True)
        
        z_cumsum = torch.cumsum(z_sorted, dim=1)
        
        k = torch.arange(1, logits.size(1) + 1, device=logits.device)
        
        support = (1 + k * z_sorted) > z_cumsum
        k_z = torch.sum(support, dim=1, keepdim=True) # (N, 1)
        
        zs_sum = torch.gather(z_cumsum, 1, k_z - 1)
        tau = (zs_sum - 1) / k_z
        
        mask = torch.arange(NUM_CLASSES, device=logits.device).unsqueeze(0) < k_z
        z_support = z_sorted * mask
        
        sum_sq_z = (z_support ** 2).sum(dim=1)
        sum_sq_tau = (tau.squeeze(1) ** 2) * k_z.squeeze(1).float()
        
        z_y = logits.gather(1, targets.unsqueeze(1)).squeeze(1)
        
        loss = 0.5 * (sum_sq_z - sum_sq_tau) + 0.5 - z_y
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = SparsemaxLoss(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 ['none']