import torch
import torch.nn as nn
import numpy as np


class APLoss(nn.Module):
    """differentiable AP loss, through quantization.

    Input: (N, M)   values in [min, max]
    label: (N, M)   values in {0, 1}

    Returns: list of query AP (for each n in {1..N})
             Note: typically, you want to minimize 1 - mean(AP)
    """

    def __init__(self, nq=25, min=0, max=1, euc=False):
        nn.Module.__init__(self)
        assert isinstance(nq, int) and 2 <= nq <= 100
        self.nq = nq
        self.min = min
        self.max = max
        self.euc = euc
        gap = max - min
        assert gap > 0

        # init quantizer = non-learnable (fixed) convolution
        self.quantizer = q = nn.Conv1d(1, 2 * nq, kernel_size=1, bias=True)
        a = (nq - 1) / gap
        # 1st half = lines passing to (min+x,1) and (min+x+1/a,0) with x = {nq-1..0}*gap/(nq-1)
        q.weight.data[:nq] = -a
        q.bias.data[:nq] = torch.from_numpy(
            a * min + np.arange(nq, 0, -1)
        )  # b = 1 + a*(min+x)
        # 2nd half = lines passing to (min+x,1) and (min+x-1/a,0) with x = {nq-1..0}*gap/(nq-1)
        q.weight.data[nq:] = a
        q.bias.data[nq:] = torch.from_numpy(
            np.arange(2 - nq, 2, 1) - a * min
        )  # b = 1 - a*(min+x)
        # first and last one are special: just horizontal straight line
        q.weight.data[0] = q.weight.data[-1] = 0
        q.bias.data[0] = q.bias.data[-1] = 1

    def compute_AP(self, x, label):
        N, M = x.shape
        # print(x.shape, label.shape)
        if self.euc:  # euclidean distance in same range than similarities
            x = 1 - torch.sqrt(2.001 - 2 * x)

        # quantize all predictions
        q = self.quantizer(x.unsqueeze(1))
        q = torch.min(q[:, : self.nq], q[:, self.nq :]).clamp(
            min=0
        )  # N x Q x M [1600, 20, 1681]

        nbs = q.sum(dim=-1)  # number of samples  N x Q = c
        rec = (q * label.view(N, 1, M).float()).sum(
            dim=-1
        )  # nb of correct samples = c+ N x Q
        prec = rec.cumsum(dim=-1) / (1e-16 + nbs.cumsum(dim=-1))  # precision
        rec /= rec.sum(dim=-1).unsqueeze(1)  # norm in [0,1]

        ap = (prec * rec).sum(dim=-1)  # per-image AP
        return ap

    def forward(self, x, label):
        assert x.shape == label.shape  # N x M
        return self.compute_AP(x, label)


class PixelAPLoss(nn.Module):
    """Computes the pixel-wise AP loss:
    Given two images and ground-truth optical flow, computes the AP per pixel.

    feat1:  (B, C, H, W)   pixel-wise features extracted from img1
    feat2:  (B, C, H, W)   pixel-wise features extracted from img2
    aflow:  (B, 2, H, W)   absolute flow: aflow[...,y1,x1] = x2,y2
    """

    def __init__(self, sampler, nq=20):
        nn.Module.__init__(self)
        self.aploss = APLoss(nq, min=0, max=1, euc=False)
        self.name = "pixAP"
        self.sampler = sampler

    def loss_from_ap(self, ap, rel):
        return 1 - ap

    def forward(self, feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200):
        # subsample things
        scores, gt, msk, qconf = self.sampler(
            feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200
        )

        # compute pixel-wise AP
        n = qconf.numel()
        if n == 0:
            return 0
        scores, gt = scores.view(n, -1), gt.view(n, -1)
        ap = self.aploss(scores, gt).view(msk.shape)

        pixel_loss = self.loss_from_ap(ap, qconf)

        loss = pixel_loss[msk].mean()
        return loss


class ReliabilityLoss(PixelAPLoss):
    """same than PixelAPLoss, but also train a pixel-wise confidence
    that this pixel is going to have a good AP.
    """

    def __init__(self, sampler, base=0.5, **kw):
        PixelAPLoss.__init__(self, sampler, **kw)
        assert 0 <= base < 1
        self.base = base

    def loss_from_ap(self, ap, rel):
        return 1 - ap * rel - (1 - rel) * self.base