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


class Model(nn.Module):

    def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, smooth=1.0):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.smooth = smooth

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

        probs = torch.sigmoid(logits)

        probs = probs.flatten(1)
        targets_f_flat = targets_f.flatten(1)

        tp = (probs * targets_f_flat).sum(dim=1)  # True Positive
        fp = (probs * (1.0 - targets_f_flat)).sum(dim=1)  # False Positive
        fn = ((1.0 - probs) * targets_f_flat).sum(dim=1)  # False Negative

        tversky_index = (tp + self.smooth) / (tp + self.alpha * fn + self.beta * fp + self.smooth)

        focal_tversky_loss = torch.pow(1.0 - tversky_index, self.gamma)

        return focal_tversky_loss.mean()


batch_size = 128
feature_dim = 100


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.7, 0.3, 0.75, 1.0]