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


class Model(nn.Module):
    def __init__(self, alpha_combo=0.5, gamma_focal=2.0, smooth=1.0):
        super().__init__()
        self.alpha_combo = alpha_combo
        self.gamma_focal = gamma_focal
        self.smooth = smooth

    def _dice_loss(self, inputs, targets) -> torch.Tensor:
        inputs = inputs.sigmoid()
        inputs = inputs.flatten(1)
        targets = targets.flatten(1)

        intersection = (inputs * targets).sum(dim=1)
        dice = (2.0 * intersection + self.smooth) / (inputs.sum(dim=1) + targets.sum(dim=1) + self.smooth)

        return 1.0 - dice.mean()

    def _focal_loss(self, inputs, targets) -> torch.Tensor:
        BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)

        focal_loss = ((1.0 - pt) ** self.gamma_focal) * BCE_loss
        return focal_loss.mean()

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        targets_f = targets.float()

        dice_loss = self._dice_loss(logits, targets_f)

        focal_loss = self._focal_loss(logits, targets_f)

        return self.alpha_combo * dice_loss + (1.0 - self.alpha_combo) * focal_loss


batch_size = 512
feature_dim = 128


def get_inputs():
    logits = torch.randn(batch_size, feature_dim, dtype=torch.float64)
    targets = torch.randint(0, 2, (batch_size, feature_dim), dtype=torch.float64)
    return [logits, targets]


def get_init_inputs():
    return [0.5, 2.0, 1.0]