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

# 多标签分类场景
BATCH_SIZE = 4096
NUM_CLASSES = 80 # COCO classes
SHAPE = (BATCH_SIZE, NUM_CLASSES)

GAMMA_POS = 0.0
GAMMA_NEG = 4.0
MARGIN = 0.05
EPS = 1e-8
REDUCTION = 'none'

class AsymmetricLoss(nn.Module):
    def __init__(self, gamma_neg=4.0, gamma_pos=1.0, clip=0.05, eps=1e-8, reduction='mean'):
        super(AsymmetricLoss, self).__init__()
        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip = clip
        self.eps = eps
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N, C) binary 0/1
        p = torch.sigmoid(logits)
        
        pos_loss = (1 - p).pow(self.gamma_pos) * torch.log(p + self.eps)
        pos_loss = -pos_loss
        
        p_m = (p - self.clip).clamp(min=0.0)
        neg_loss = p_m.pow(self.gamma_neg) * torch.log(1 - p_m + self.eps)
        neg_loss = -neg_loss
        
        loss = targets * pos_loss + (1 - targets) * neg_loss
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, gamma_neg=4.0, gamma_pos=1.0, clip=0.05, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = AsymmetricLoss(gamma_neg=gamma_neg, gamma_pos=gamma_pos, clip=clip, 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, 2, SHAPE, dtype=torch.float32)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    return [GAMMA_NEG, GAMMA_POS, MARGIN, REDUCTION]