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

class Model(nn.Module):
    """
    Focal Loss with Sigmoid implementation for binary classification.
    Fully fused version that computes sigmoid and focal loss in a single pass.
    Focal Loss = -α * (1-pt)^γ * log(pt)
    where pt = p if target=1, else (1-p), p = sigmoid(logit)
    """
    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 with fused sigmoid computation.
        
        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
        """
        # Ensure input types are consistent
        inputs = inputs.to(torch.float32)
        targets = targets.to(torch.float32)
        
        # Handle shape matching
        if inputs.dim() == 2 and inputs.size(1) == 1:
            inputs = inputs.squeeze(1)
        
        # Fully fused computation using highly optimized built-in functions
        # Compute sigmoid with numerical stability
        sigmoid_inputs = torch.sigmoid(inputs)
        
        # Compute pt based on targets
        pt = torch.where(targets == 1, sigmoid_inputs, 1 - sigmoid_inputs)
        
        # Compute focal weight
        focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
        
        # Compute binary cross entropy with logits (more numerically stable)
        bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        
        # Apply focal weight
        focal_loss = focal_weight * bce
        
        # Apply reduction
        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 with explicit float32
    inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
    # Generate random binary targets (0 or 1) with explicit float32
    targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
    return [inputs, targets]

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