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


class Model(nn.Module):
    def __init__(self, alpha=0.5, smooth=1e-6):
        super().__init__()
        self.alpha = alpha
        self.smooth = smooth

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

        probs = logits.sigmoid()
        probs = probs.flatten(1)
        targets_f = targets_f.flatten(1)

        intersection = (probs * targets_f).sum(dim=1)
        union = probs.sum(dim=1) + targets_f.sum(dim=1)

        angle = 1.0 - (2.0 * intersection + self.smooth) / (union + self.smooth)
        angular_loss = (angle * (1.0 - angle)).mean()

        bce_loss = F.binary_cross_entropy_with_logits(logits, targets_f, reduction='mean')

        return self.alpha * angular_loss + (1.0 - self.alpha) * bce_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, 1e-6]