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

N, C, H, W = 32, 64, 56, 56
EPS = 1e-8


class CosineEmbeddingLossCustom(nn.Module):

    def __init__(self, margin=0.0, reduction='mean', eps=1e-8):
        super().__init__()
        self.margin = margin
        self.reduction = reduction
        self.eps = eps

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

        dot_product = torch.sum(x1 * x2, dim=1)
        norm_x1 = torch.norm(x1, p=2, dim=1)
        norm_x2 = torch.norm(x2, p=2, dim=1)

        cos_sim = dot_product / (norm_x1 * norm_x2 + self.eps)

        loss_pos = 1.0 - cos_sim
        loss_neg = F.relu(cos_sim - self.margin)

        loss = torch.where(target == 1, loss_pos, loss_neg)

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


class Model(nn.Module):
    def __init__(self, margin=0.5):
        super().__init__()
        self.op = CosineEmbeddingLossCustom(margin=margin, reduction='mean', eps=EPS)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        x1_flat = x1.view(x1.size(0), -1)
        x2_flat = x2.view(x2.size(0), -1)
        return self.op(x1_flat, x2_flat, target)


def get_inputs():
    x1 = torch.randn(N, C, H, W, dtype=torch.float32)
    x2 = torch.randn(N, C, H, W, dtype=torch.float32)

    target = torch.randint(0, 2, (N,), dtype=torch.float32)  # 0 or 1
    target = torch.where(target == 0, torch.tensor(-1.0), torch.tensor(1.0))

    return [x1, x2, target]


def get_init_inputs():
    return [0.5]