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

class Model(nn.Module):
    """
    Focal Loss implementation for binary classification.
    Focal Loss = -α * (1-pt)^γ * log(pt)
    where pt = p if target=1, else (1-p)
    """
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
        super(Model, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
    
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        Compute Focal Loss between inputs and targets.
        
        Args:
            inputs (torch.Tensor): Predicted logits of shape (batch_size, num_classes)
            targets (torch.Tensor): Ground truth labels of shape (batch_size,)
        
        Returns:
            torch.Tensor: Computed focal loss
        """
        # 确保输入输出类型完全一致
        inputs = inputs.to(torch.float32)
        targets = targets.to(torch.float32)
        
        # 确保targets的形状与inputs匹配
        if inputs.dim() == 2 and inputs.size(1) == 1:
            targets = targets.unsqueeze(1)  # 将targets从[batch_size]变为[batch_size, 1]
        
        # Convert logits to probabilities
        probs = torch.sigmoid(inputs)
        
        # Compute pt
        pt = torch.where(targets == 1, probs, 1 - probs)
        
        # Compute focal weight
        focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
        
        # Compute binary cross entropy
        bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        
        # Apply focal weight
        focal_loss = focal_weight * bce
        
        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

batch_size = 32
num_classes = 1  # Binary classification

def get_inputs():
    # Generate random logits - 明确指定float32
    inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
    # Generate random binary targets (0 or 1) - 明确指定float32
    targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
    return [inputs, targets]

def get_init_inputs():
    return []  # No special initialization inputs needed