import torch
import torch.nn as nn

class Model(nn.Module):
    """
    PyTorch基准实现：Dice Loss + BCEWithLogitsLoss
    """
    def __init__(self, dice_weight=1.0, bce_weight=1.0):
        super(Model, self).__init__()
        self.dice_weight = dice_weight
        self.bce_weight = bce_weight
        # BCEWithLogitsLoss内部融合了Sigmoid和BCE，更稳定
        self.bce_loss_fn = nn.BCEWithLogitsLoss(reduction='mean')

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Computes the combined Dice and BCE loss.

        Args:
            logits (torch.Tensor): Raw prediction tensor of shape [N, C, H, W].
            target (torch.Tensor): Target tensor of same shape as logits.

        Returns:
            torch.Tensor: Combined loss value (scalar).
        """
        # --- 第一步：计算Dice Loss ---
        pred = torch.sigmoid(logits)
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        
        intersection = (pred_flat * target_flat).sum()
        pred_sum = pred_flat.sum()
        target_sum = target_flat.sum()
        
        epsilon = 1e-6
        dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
        loss_dice = 1.0 - dice_score
        
        # --- 第二步：计算BCEWithLogitsLoss ---
        # PyTorch的BCEWithLogitsLoss已经非常高效且稳定
        loss_bce = self.bce_loss_fn(logits, target)
        
        # --- 第三步：组合损失 ---
        loss_total = self.dice_weight * loss_dice + self.bce_weight * loss_bce
        
        return loss_total

batch_size = 32
height, width = 256, 256
channels = 1

def get_inputs():
    logits = torch.randn(batch_size, channels, height, width)
    target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
    return [logits, target]

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