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

N, C, H, W = 32, 1, 64, 64


class FocalLoss(nn.Module):
    def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
        super().__init__()
        self.reduction = reduction
        self.alpha = float(alpha)
        self.gamma = float(gamma)
        self.epsilon = 1e-8

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

        bce_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')

        p = torch.sigmoid(input)

        p_t = p * target + (1 - p) * (1 - target)

        modulating_factor = (1.0 - p_t).pow(self.gamma)

        alpha_t = self.alpha * target + (1 - self.alpha) * (1 - target)

        loss = alpha_t * modulating_factor * bce_loss

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


class Model(nn.Module):
    def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
        super().__init__()
        self.op = FocalLoss(reduction, alpha, gamma)

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        if isinstance(input, (list, tuple)) and len(input) > 0:
            input = input[0]
            target = target[0] if len(target) > 0 else target

        return self.op(input, target)


def get_inputs():
    input = torch.randn(N, C, H, W, dtype=torch.float32)
    target = torch.randint(0, 2, (N, C, H, W), dtype=torch.float32)
    return [input, target]


def get_init_inputs():
    return ['mean', 0.25, 2.0]
