import torch
import torch.nn as nn
from .reliability_loss import APLoss


class MultiPixelAPLoss(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.sampler = sampler
        self.base = 0.25
        self.dec_base = 0.20

    def loss_from_ap(self, ap, rel, noise_ap, noise_rel):
        dec_ap = torch.clamp(ap - noise_ap, min=0, max=1)
        return (1 - ap * noise_rel - (1 - noise_rel) * self.base), (
            1.0 - dec_ap * (1 - noise_rel) - noise_rel * self.dec_base
        )

    def forward(
        self,
        feat0,
        feat1,
        noise_feat0,
        noise_feat1,
        conf0,
        conf1,
        noise_conf0,
        noise_conf1,
        pos0,
        pos1,
        B,
        H,
        W,
        N=1500,
    ):
        # subsample things
        scores, noise_scores, gt, msk, qconf, noise_qconf = self.sampler(
            feat0,
            feat1,
            noise_feat0,
            noise_feat1,
            conf0,
            conf1,
            noise_conf0,
            noise_conf1,
            pos0,
            pos1,
            B,
            H,
            W,
            N=1500,
        )

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

        pixel_loss = self.loss_from_ap(ap, qconf, noise_ap, noise_qconf)

        loss = pixel_loss[0][msk].mean(), pixel_loss[1][msk].mean()
        return loss