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

N, D = 32, 128


class TripletMarginWithDistanceLoss(nn.Module):

    def __init__(self, distance_function=None, margin=1.0, swap=False, reduction='mean'):
        super().__init__()
        self.distance_function = distance_function if distance_function is not None else nn.PairwiseDistance()
        self.margin = margin
        self.swap = swap
        self.reduction = reduction

    def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:

        d_ap = self.distance_function(anchor, positive)

        d_an = self.distance_function(anchor, negative)

        if self.swap:
            d_pn = self.distance_function(positive, negative)
            d_an = torch.min(d_an, d_pn)

        loss = torch.clamp(d_ap - d_an + self.margin, min=0.0)

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


class Model(nn.Module):
    def __init__(self, margin=1.0, swap=False):
        super().__init__()

        self.op = TripletMarginWithDistanceLoss(distance_function=nn.PairwiseDistance(), margin=margin, swap=swap)

    def forward(self, a: torch.Tensor, p: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
        return self.op(a, p, n)


def get_inputs():
    anchor = torch.randn(N, D, dtype=torch.float32)
    positive = torch.randn(N, D, dtype=torch.float32)
    negative = torch.randn(N, D, dtype=torch.float32)
    return [anchor, positive, negative]


def get_init_inputs():
    return [1.0, False]