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

class Model(nn.Module):
    """
    Focal Loss with Label Smoothing implementation for binary classification.
    Combines label smoothing with focal loss for better generalization.
    Focal Loss = -α * (1-pt)^γ * log(pt_smoothed)
    where pt = p if target=1, else (1-p), p = sigmoid(logit)
    """
    def __init__(self, alpha=0.25, gamma=2.0, smoothing=0.1, reduction='mean'):
        super(Model, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.smoothing = smoothing
        self.reduction = reduction
    
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        Compute Focal Loss with Label Smoothing.
        
        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 with label smoothing
        """
        # 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)
        
        # Apply label smoothing to targets
        # For binary classification: smooth_target = (1-smoothing)*target + smoothing/2
        smoothed_targets = (1.0 - self.smoothing) * targets + self.smoothing / 2.0
        
        # Compute probabilities with sigmoid
        probs = torch.sigmoid(inputs)
        
        # Compute pt based on original targets (not smoothed)
        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 with smoothed targets
        bce = F.binary_cross_entropy_with_logits(inputs, smoothed_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 [0.25, 2.0, 0.1]  # alpha, gamma, smoothing
