######################################
#######ORIGINAL IMPLEMENTATION########
######################################
# FROM https://github.com/kunhe/FastAP-metric-learning/blob/master/pytorch/FastAP_loss.py
# This code is copied directly from the official implementation
# so that we can make sure our implementation returns the same result.
# It's copied under the MIT license.
import torch
from torch.autograd import Variable


def softBinning(D, mid, Delta):
    y = 1 - torch.abs(D - mid) / Delta
    return torch.max(torch.tensor([0], dtype=D.dtype).to(D.device), y)


def dSoftBinning(D, mid, Delta):
    side1 = (D > (mid - Delta)).type(D.dtype)
    side2 = (D <= mid).type(D.dtype)
    ind1 = side1 * side2  # .type(torch.uint8)

    side1 = (D > mid).type(D.dtype)
    side2 = (D <= (mid + Delta)).type(D.dtype)
    ind2 = side1 * side2  # .type(torch.uint8)

    return (ind1 - ind2) / Delta


######################################
#######ORIGINAL IMPLEMENTATION########
######################################
# FROM https://github.com/kunhe/FastAP-metric-learning/blob/master/pytorch/FastAP_loss.py
# This code is copied directly from the official implementation
# so that we can make sure our implementation returns the same result.
# It's copied under the MIT license.
class OriginalImplementationFastAP(torch.autograd.Function):
    """
    FastAP - autograd function definition

    This class implements the FastAP loss from the following paper:
    "Deep Metric Learning to Rank",
    F. Cakir, K. He, X. Xia, B. Kulis, S. Sclaroff. CVPR 2019

    NOTE:
        Given a input batch, FastAP does not sample triplets from it as it's not
        a triplet-based method. Therefore, FastAP does not take a Sampler as input.
        Rather, we specify how the input batch is selected.
    """

    @staticmethod
    def forward(ctx, input, target, num_bins):
        """
        Args:
            input:     torch.Tensor(N x embed_dim), embedding matrix
            target:    torch.Tensor(N x 1), class labels
            num_bins:  int, number of bins in distance histogram
        """
        N = target.size()[0]
        assert input.size()[0] == N, "Batch size donesn't match!"

        # 1. get affinity matrix
        Y = target.unsqueeze(1)
        Aff = 2 * (Y == Y.t()).type(input.dtype) - 1
        Aff.masked_fill_(
            torch.eye(N, N).bool().to(input.device), 0
        )  # set diagonal to 0

        I_pos = (Aff > 0).type(input.dtype).to(input.device)
        I_neg = (Aff < 0).type(input.dtype).to(input.device)
        N_pos = torch.sum(I_pos, 1)

        # 2. compute distances from embeddings
        # squared Euclidean distance with range [0,4]
        dist2 = 2 - 2 * torch.mm(input, input.t())

        # 3. estimate discrete histograms
        Delta = torch.tensor(4.0 / num_bins).to(input.device)
        Z = torch.linspace(0.0, 4.0, steps=num_bins + 1).to(input.device)
        L = Z.size()[0]
        h_pos = torch.zeros((N, L), dtype=input.dtype).to(input.device)
        h_neg = torch.zeros((N, L), dtype=input.dtype).to(input.device)
        for l in range(L):
            pulse = softBinning(dist2, Z[l], Delta)
            h_pos[:, l] = torch.sum(pulse * I_pos, 1)
            h_neg[:, l] = torch.sum(pulse * I_neg, 1)

        H_pos = torch.cumsum(h_pos, 1)
        h = h_pos + h_neg
        H = torch.cumsum(h, 1)

        # 4. compate FastAP
        FastAP = h_pos * H_pos / H
        FastAP[torch.isnan(FastAP) | torch.isinf(FastAP)] = 0
        FastAP = torch.sum(FastAP, 1) / N_pos
        FastAP = FastAP[~torch.isnan(FastAP)]
        loss = 1 - torch.mean(FastAP)

        # 6. save for backward
        ctx.save_for_backward(input, target)
        ctx.Z = Z
        ctx.Delta = Delta
        ctx.dist2 = dist2
        ctx.I_pos = I_pos
        ctx.I_neg = I_neg
        ctx.h_pos = h_pos
        ctx.h_neg = h_neg
        ctx.H_pos = H_pos
        ctx.N_pos = N_pos
        ctx.h = h
        ctx.H = H
        ctx.L = torch.tensor(L)

        return loss

    @staticmethod
    def backward(ctx, grad_output):
        input, target = ctx.saved_tensors

        Z = Variable(ctx.Z, requires_grad=False)
        Delta = Variable(ctx.Delta, requires_grad=False)
        dist2 = Variable(ctx.dist2, requires_grad=False)
        I_pos = Variable(ctx.I_pos, requires_grad=False)
        I_neg = Variable(ctx.I_neg, requires_grad=False)
        h = Variable(ctx.h, requires_grad=False)  # noqa: F841
        H = Variable(ctx.H, requires_grad=False)
        h_pos = Variable(ctx.h_pos, requires_grad=False)
        h_neg = Variable(ctx.h_neg, requires_grad=False)  # noqa: F841
        H_pos = Variable(ctx.H_pos, requires_grad=False)
        N_pos = Variable(ctx.N_pos, requires_grad=False)

        L = Z.size()[0]
        H2 = torch.pow(H, 2)
        H_neg = H - H_pos

        # 1. d(FastAP)/d(h+)
        LTM1 = torch.tril(torch.ones(L, L), -1)  # lower traingular matrix
        tmp1 = h_pos * H_neg / H2
        tmp1[torch.isnan(tmp1)] = 0

        d_AP_h_pos = (H_pos * H + h_pos * H_neg) / H2
        d_AP_h_pos = d_AP_h_pos + torch.mm(tmp1, LTM1.cuda())
        d_AP_h_pos = d_AP_h_pos / N_pos.repeat(L, 1).t()
        d_AP_h_pos[torch.isnan(d_AP_h_pos) | torch.isinf(d_AP_h_pos)] = 0

        # 2. d(FastAP)/d(h-)
        LTM0 = torch.tril(torch.ones(L, L), 0)  # lower triangular matrix
        tmp2 = -h_pos * H_pos / H2
        tmp2[torch.isnan(tmp2)] = 0

        d_AP_h_neg = torch.mm(tmp2, LTM0.cuda())
        d_AP_h_neg = d_AP_h_neg / N_pos.repeat(L, 1).t()
        d_AP_h_neg[torch.isnan(d_AP_h_neg) | torch.isinf(d_AP_h_neg)] = 0

        # 3. d(FastAP)/d(embedding)
        d_AP_x = 0
        for l in range(L):
            dpulse = dSoftBinning(dist2, Z[l], Delta)
            dpulse[torch.isnan(dpulse) | torch.isinf(dpulse)] = 0
            ddp = dpulse * I_pos
            ddn = dpulse * I_neg

            alpha_p = torch.diag(d_AP_h_pos[:, l])  # N*N
            alpha_n = torch.diag(d_AP_h_neg[:, l])
            Ap = torch.mm(ddp, alpha_p) + torch.mm(alpha_p, ddp)
            An = torch.mm(ddn, alpha_n) + torch.mm(alpha_n, ddn)

            # accumulate gradient
            d_AP_x = d_AP_x - torch.mm(input.t(), (Ap + An))

        grad_input = -d_AP_x
        return grad_input.t(), None, None


######################################
#######ORIGINAL IMPLEMENTATION########
######################################
# FROM https://github.com/kunhe/FastAP-metric-learning/blob/master/pytorch/FastAP_loss.py
# This code is copied directly from the official implementation
# so that we can make sure our implementation returns the same result.
# It's copied under the MIT license.
class OriginalImplementationFastAPLoss(torch.nn.Module):
    """
    FastAP - loss layer definition

    This class implements the FastAP loss from the following paper:
    "Deep Metric Learning to Rank",
    F. Cakir, K. He, X. Xia, B. Kulis, S. Sclaroff. CVPR 2019
    """

    def __init__(self, num_bins=10):
        super(OriginalImplementationFastAPLoss, self).__init__()
        self.num_bins = num_bins

    def forward(self, batch, labels):
        return OriginalImplementationFastAP.apply(batch, labels, self.num_bins)


### Testing this library's implementation ###
import unittest

from pytorch_metric_learning.losses import FastAPLoss

from .. import TEST_DEVICE, TEST_DTYPES
from ..zzz_testing_utils.testing_utils import angle_to_coord


class TestFastAPLoss(unittest.TestCase):
    def test_fast_ap_loss(self):
        num_bins = 5
        loss_func = FastAPLoss(num_bins)
        original_loss_func = OriginalImplementationFastAPLoss(num_bins)
        ref_emb = torch.randn(32, 32)
        ref_labels = torch.randint(0, 10, (32,))

        for dtype in TEST_DTYPES:
            embedding_angles = torch.arange(0, 180)
            embeddings = torch.tensor(
                [angle_to_coord(a) for a in embedding_angles],
                requires_grad=True,
                dtype=dtype,
            ).to(
                TEST_DEVICE
            )  # 2D embeddings
            labels = torch.randint(low=0, high=10, size=(180,)).to(TEST_DEVICE)

            loss = loss_func(embeddings, labels)
            loss.backward()
            original_loss = original_loss_func(
                torch.nn.functional.normalize(embeddings), labels
            )
            rtol = 1e-2 if dtype == torch.float16 else 1e-5
            self.assertTrue(torch.isclose(loss, original_loss, rtol=rtol))

            # fastap doesn't support ref_emb
            self.assertRaises(
                ValueError,
                lambda: loss_func(
                    embeddings, labels, ref_emb=ref_emb, ref_labels=ref_labels
                ),
            )
