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

BATCH_SIZE = 128
NUM_ANCHORS = 64
NUM_CLASSES = 64
SHAPE = (BATCH_SIZE, NUM_ANCHORS * NUM_CLASSES)

GAMMA = 2.0
ALPHA = 0.25
EPSILON = 1.0

class Poly1FocalLoss(nn.Module):
    """
    Optimized Poly1FocalLoss Implementation.
    L = L_fl + epsilon * (1-Pt)^(gamma+1)
    """
    def __init__(self, gamma=2.0, alpha=0.25, epsilon=1.0, reduction='mean'):
        super(Poly1FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.epsilon = epsilon
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # 1. 计算 BCE Loss (即 -log(Pt))
        ce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')

        # 2. 计算预测概率 p
        p = torch.sigmoid(logits)

        # 3. 计算 (1 - Pt)
        p_t_complement = (p - targets).abs()

        # 4. 计算 Alpha_t
        alpha_t = torch.where(targets == 1, self.alpha, 1.0 - self.alpha)

        # 5. 计算 Poly1 Focal Loss
        modulating_factor = p_t_complement.pow(self.gamma)
        
        loss = modulating_factor * (alpha_t * ce_loss + self.epsilon * p_t_complement)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, gamma=2.0, alpha=0.25, epsilon=1.0, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = Poly1FocalLoss(gamma=gamma, alpha=alpha, epsilon=epsilon, 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, ALPHA, EPSILON, 'none']