import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

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

MAX_M = 0.5
SCALE_S = 30.0
REDUCTION = 'none'

class LDAMLoss(nn.Module):
    """
    LDAM Loss (NeurIPS 2019)
    """
    def __init__(self, cls_num_list, max_m=0.5, s=30.0, reduction='mean'):
        super(LDAMLoss, self).__init__()
        self.s = s
        self.reduction = reduction
        # 预计算 Margins
        # Formula: m_j = C / n_j^(1/4)
        m_list = 1.0 / np.sqrt(np.sqrt(cls_num_list))
        m_list = m_list * (max_m / np.max(m_list))
        m_list = torch.FloatTensor(m_list)
        self.register_buffer('class_margins', m_list)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N)
        
        batch_margins = self.class_margins[targets]

        
        logits_m = logits.clone()
        
        idx = targets.view(-1, 1)
        logits_m.scatter_add_(1, idx, -batch_margins.view(-1, 1))
        
        output = logits_m * self.s
        
        loss = F.cross_entropy(output, targets, 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, cls_num_list, max_m=0.5, s=30.0, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = LDAMLoss(cls_num_list, max_m, s, 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():
    # 模拟长尾分布的样本计数
    cls_num_list = np.array([10000 * (0.9 ** (i / 10.0)) for i in range(NUM_CLASSES)])
    cls_num_list = np.maximum(cls_num_list, 1) # 避免除0
    
    return [cls_num_list, MAX_M, SCALE_S, REDUCTION]