import torch
import torch.nn.functional as F

from torch import nn
import cv2
import numpy as np
from copy import deepcopy
from RDD.dataset.megadepth.utils import warp


def plot_keypoints(image, kpts, radius=2, color=(255, 0, 0)):
    image = image.cpu().detach().numpy() if isinstance(image, torch.Tensor) else image
    kpts = kpts.cpu().detach().numpy() if isinstance(kpts, torch.Tensor) else kpts

    if image.dtype is not np.dtype("uint8"):
        image = image * 255
        image = image.astype(np.uint8)

    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    out = np.ascontiguousarray(deepcopy(image))
    kpts = np.round(kpts).astype(int)

    for kpt in kpts:
        y0, x0 = kpt
        cv2.drawMarker(out, (x0, y0), color, cv2.MARKER_CROSS, radius)

        # cv2.circle(out, (x0, y0), radius, color, -1, lineType=cv2.LINE_4)
    return out


class DetectorLoss(nn.Module):
    def __init__(self, temperature=0.1, scores_th=0.1, peaky_weight=0.5, reprojection_weight=1, scoremap_weight=0.5):
        super().__init__()
        self.temperature = temperature
        self.scores_th = scores_th
        self.peaky_weight = peaky_weight
        self.reprojection_weight = reprojection_weight
        self.scoremap_weight = scoremap_weight

        self.PeakyLoss = PeakyLoss(scores_th=scores_th)
        self.ReprojectionLocLoss = ReprojectionLocLoss(scores_th=scores_th)
        self.ScoreMapRepLoss = ScoreMapRepLoss(temperature=temperature)

    def forward(self, correspondences, pred0_with_rand, pred1_with_rand):
        loss_peaky0 = self.PeakyLoss(pred0_with_rand)
        loss_peaky1 = self.PeakyLoss(pred1_with_rand)
        loss_peaky = (loss_peaky0 + loss_peaky1) / 2.0

        loss_reprojection = self.ReprojectionLocLoss(pred0_with_rand, pred1_with_rand, correspondences)

        loss_score_map_rp = self.ScoreMapRepLoss(pred0_with_rand, pred1_with_rand, correspondences)

        loss_kp = (
            loss_peaky * self.peaky_weight
            + loss_reprojection * self.reprojection_weight
            + loss_score_map_rp * self.scoremap_weight
        )

        return loss_kp


class PeakyLoss(object):
    """PeakyLoss to avoid an uniform score map"""

    def __init__(self, scores_th: float = 0.1):
        super().__init__()
        self.scores_th = scores_th

    def __call__(self, pred):
        b, c, h, w = pred["scores_map"].shape
        loss_mean = 0
        CNT = 0

        for idx in range(b):
            n_original = len(pred["score_dispersity"][idx])
            scores_kpts = pred["scores"][idx][:n_original]
            valid = scores_kpts > self.scores_th
            loss_peaky = pred["score_dispersity"][idx][valid]

            loss_mean = loss_mean + loss_peaky.sum()
            CNT = CNT + len(loss_peaky)

        loss_mean = loss_mean / CNT if CNT != 0 else pred["scores_map"].new_tensor(0)
        assert not torch.isnan(loss_mean)
        return loss_mean


class ReprojectionLocLoss(object):
    """
    Reprojection location errors of keypoints to train repeatable detector.
    """

    def __init__(self, norm: int = 1, scores_th: float = 0.1):
        super().__init__()
        self.norm = norm
        self.scores_th = scores_th

    def __call__(self, pred0, pred1, correspondences):
        b, c, h, w = pred0["scores_map"].shape
        loss_mean = 0
        CNT = 0
        for idx in range(b):
            if correspondences[idx]["correspondence0"] is None:
                continue

            if self.norm == 2:
                dist = correspondences[idx]["dist"]
            elif self.norm == 1:
                dist = correspondences[idx]["dist_l1"]
            else:
                raise TypeError("No such norm in correspondence.")

            ids0_d = correspondences[idx]["ids0_d"]
            ids1_d = correspondences[idx]["ids1_d"]

            scores0 = correspondences[idx]["scores0"].detach()[ids0_d]
            scores1 = correspondences[idx]["scores1"].detach()[ids1_d]
            valid = (scores0 > self.scores_th) * (scores1 > self.scores_th)
            reprojection_errors = dist[ids0_d, ids1_d][valid]

            loss_mean = loss_mean + reprojection_errors.sum()
            CNT = CNT + len(reprojection_errors)

        loss_mean = loss_mean / CNT if CNT != 0 else correspondences[0]["dist"].new_tensor(0)

        assert not torch.isnan(loss_mean)
        return loss_mean


def local_similarity(descriptor_map, descriptors, kpts_wh, radius):
    """
    :param descriptor_map: CxHxW
    :param descriptors: NxC
    :param kpts_wh: Nx2 (W,H)
    :return:
    """
    _, h, w = descriptor_map.shape
    ksize = 2 * radius + 1

    descriptor_map_unflod = torch.nn.functional.unfold(
        descriptor_map.unsqueeze(0), kernel_size=(ksize, ksize), padding=(radius, radius)
    )
    descriptor_map_unflod = descriptor_map_unflod[0].t().reshape(h * w, -1, ksize * ksize)
    # find the correspondence patch
    kpts_wh_long = kpts_wh.detach().long()
    patch_ids = kpts_wh_long[:, 0] + kpts_wh_long[:, 1] * h
    desc_patches = descriptor_map_unflod[patch_ids].permute(0, 2, 1).detach()  # N_kpts x s*s x 128

    local_sim = torch.einsum("nsd,nd->ns", desc_patches, descriptors)
    local_sim_sort = torch.sort(local_sim, dim=1, descending=True).values
    local_sim_sort_mean = local_sim_sort[:, 4:].mean(dim=1)  # 4 is safe radius for bilinear interplation

    return local_sim_sort_mean


class ScoreMapRepLoss(object):
    """Scoremap repetability"""

    def __init__(self, temperature: float = 0.1):
        super().__init__()
        self.temperature = temperature
        self.radius = 2

    def __call__(self, pred0, pred1, correspondences):
        b, c, h, w = pred0["scores_map"].shape
        wh = pred0["keypoints"][0].new_tensor([[w - 1, h - 1]])
        loss_mean = 0
        CNT = 0

        for idx in range(b):
            if correspondences[idx]["correspondence0"] is None:
                continue

            scores_map0 = pred0["scores_map"][idx]
            scores_map1 = pred1["scores_map"][idx]
            kpts01 = correspondences[idx]["kpts01"]
            kpts10 = correspondences[idx]["kpts10"]  # valid warped keypoints

            # =====================
            scores_kpts10 = torch.nn.functional.grid_sample(
                scores_map0.unsqueeze(0), kpts10.view(1, 1, -1, 2), mode="bilinear", align_corners=True
            )[0, 0, 0, :]
            scores_kpts01 = torch.nn.functional.grid_sample(
                scores_map1.unsqueeze(0), kpts01.view(1, 1, -1, 2), mode="bilinear", align_corners=True
            )[0, 0, 0, :]

            s0 = scores_kpts01 * correspondences[idx]["scores0"]  # repeatability
            s1 = scores_kpts10 * correspondences[idx]["scores1"]  # repeatability

            # ===================== repetability
            similarity_map_01 = correspondences[idx]["similarity_map_01_valid"]
            similarity_map_10 = correspondences[idx]["similarity_map_10_valid"]

            pmf01 = ((similarity_map_01.detach() - 1) / self.temperature).exp()
            pmf10 = ((similarity_map_10.detach() - 1) / self.temperature).exp()
            kpts01 = kpts01.detach()
            kpts10 = kpts10.detach()

            pmf01_kpts = torch.nn.functional.grid_sample(
                pmf01.unsqueeze(0), kpts01.view(1, 1, -1, 2), mode="bilinear", align_corners=True
            )[0, :, 0, :]
            pmf10_kpts = torch.nn.functional.grid_sample(
                pmf10.unsqueeze(0), kpts10.view(1, 1, -1, 2), mode="bilinear", align_corners=True
            )[0, :, 0, :]
            repetability01 = torch.diag(pmf01_kpts)
            repetability10 = torch.diag(pmf10_kpts)

            # ===================== reliability
            # ids0, ids1 = correspondences[idx]['ids0'], correspondences[idx]['ids1']
            # descriptor_map0 = pred0['descriptor_map'][idx].detach()
            # descriptor_map1 = pred1['descriptor_map'][idx].detach()
            # descriptors0 = pred0['descriptors'][idx][ids0].detach()
            # descriptors1 = pred1['descriptors'][idx][ids1].detach()
            # kpts0 = pred0['keypoints'][idx][ids0].detach()
            # kpts1 = pred1['keypoints'][idx][ids1].detach()
            # kpts0_wh = (kpts0 / 2 + 0.5) * wh
            # kpts1_wh = (kpts1 / 2 + 0.5) * wh
            # ls0 = local_similarity(descriptor_map0, descriptors0, kpts0_wh, self.radius)
            # ls1 = local_similarity(descriptor_map1, descriptors1, kpts1_wh, self.radius)
            # reliability0 = 1 - ((ls0 - 1) / self.temperature).exp()
            # reliability1 = 1 - ((ls1 - 1) / self.temperature).exp()

            fs0 = repetability01  # * reliability0
            fs1 = repetability10  # * reliability1

            if s0.sum() != 0:
                loss01 = (1 - fs0) * s0 * len(s0) / s0.sum()
                loss_mean = loss_mean + loss01.sum()
                CNT = CNT + len(loss01)
            if s1.sum() != 0:
                loss10 = (1 - fs1) * s1 * len(s1) / s1.sum()
                loss_mean = loss_mean + loss10.sum()
                CNT = CNT + len(loss10)

        loss_mean = loss_mean / CNT if CNT != 0 else pred0["scores_map"].new_tensor(0)
        assert not torch.isnan(loss_mean)
        return loss_mean


# +++++++++++++++++++++++++++++++++++++++++++++++++++Taken from ALIKE+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++


def compute_keypoints_distance(kpts0, kpts1, p=2, debug=False):
    """
    Args:
        kpts0: torch.tensor [M,2]
        kpts1: torch.tensor [N,2]
        p: (int, float, inf, -inf, 'fro', 'nuc', optional): the order of norm

    Returns:
        dist, torch.tensor [N,M]
    """
    dist = kpts0[:, None, :] - kpts1[None, :, :]  # [M,N,2]
    dist = torch.norm(dist, p=p, dim=2)  # [M,N]
    return dist


def mutual_argmax(value, mask=None, as_tuple=True):
    """
    Args:
        value: MxN
        mask:  MxN

    Returns:

    """
    value = value - value.min()  # convert to non-negative tensor
    if mask is not None:
        value = value * mask

    max0 = value.max(dim=1, keepdim=True)  # the col index the max value in each row
    max1 = value.max(dim=0, keepdim=True)

    valid_max0 = value == max0[0]
    valid_max1 = value == max1[0]

    mutual = valid_max0 * valid_max1
    if mask is not None:
        mutual = mutual * mask

    return mutual.nonzero(as_tuple=as_tuple)


def mutual_argmin(value, mask=None):
    return mutual_argmax(-value, mask)


def compute_correspondence(model, pred0, pred1, batch, radius=2, rand=True, train_gt_th=5, debug=False):
    b, c, h, w = pred0["scores_map"].shape
    wh = pred0["scores_map"][0].new_tensor([[w - 1, h - 1]])

    pred0_with_rand = pred0
    pred1_with_rand = pred1
    pred0_with_rand["scores"] = []
    pred1_with_rand["scores"] = []
    pred0_with_rand["descriptors"] = []
    pred1_with_rand["descriptors"] = []
    pred0_with_rand["num_det"] = []
    pred1_with_rand["num_det"] = []

    kps, score_dispersity, scores = model.softdetect.detect_keypoints(pred0["scores_map"])

    pred0_with_rand["keypoints"] = kps
    pred0_with_rand["score_dispersity"] = score_dispersity

    kps, score_dispersity, scores = model.softdetect.detect_keypoints(pred1["scores_map"])
    pred1_with_rand["keypoints"] = kps
    pred1_with_rand["score_dispersity"] = score_dispersity

    correspondences = []
    for idx in range(b):
        # =========================== prepare keypoints
        kpts0, kpts1 = pred0["keypoints"][idx], pred1["keypoints"][idx]  # (x,y), shape: Nx2

        # additional random keypoints
        if rand:
            rand0 = torch.rand(len(kpts0), 2, device=kpts0.device) * 2 - 1  # -1~1
            rand1 = torch.rand(len(kpts1), 2, device=kpts1.device) * 2 - 1  # -1~1
            kpts0 = torch.cat([kpts0, rand0])
            kpts1 = torch.cat([kpts1, rand1])

            pred0_with_rand["keypoints"][idx] = kpts0
            pred1_with_rand["keypoints"][idx] = kpts1

        scores_map0 = pred0["scores_map"][idx]
        scores_map1 = pred1["scores_map"][idx]
        scores_kpts0 = torch.nn.functional.grid_sample(
            scores_map0.unsqueeze(0), kpts0.view(1, 1, -1, 2), mode="bilinear", align_corners=True
        ).squeeze()
        scores_kpts1 = torch.nn.functional.grid_sample(
            scores_map1.unsqueeze(0), kpts1.view(1, 1, -1, 2), mode="bilinear", align_corners=True
        ).squeeze()

        kpts0_wh_ = (kpts0 / 2 + 0.5) * wh  # N0x2, (w,h)
        kpts1_wh_ = (kpts1 / 2 + 0.5) * wh  # N1x2, (w,h)

        # ========================= nms
        dist = compute_keypoints_distance(kpts0_wh_.detach(), kpts0_wh_.detach())
        local_mask = dist < radius
        valid_cnt = torch.sum(local_mask, dim=1)
        indices_need_nms = torch.where(valid_cnt > 1)[0]

        for i in indices_need_nms:
            if valid_cnt[i] > 0:
                kpt_indices = torch.where(local_mask[i])[0]
                scs_max_idx = scores_kpts0[kpt_indices].argmax()

                tmp_mask = kpt_indices.new_ones(len(kpt_indices)).bool()
                tmp_mask[scs_max_idx] = False
                suppressed_indices = kpt_indices[tmp_mask]

                valid_cnt[suppressed_indices] = 0

        valid_mask = valid_cnt > 0
        kpts0_wh = kpts0_wh_[valid_mask]
        kpts0 = kpts0[valid_mask]
        scores_kpts0 = scores_kpts0[valid_mask]
        pred0_with_rand["keypoints"][idx] = kpts0

        valid_mask = valid_mask[: len(pred0_with_rand["score_dispersity"][idx])]
        pred0_with_rand["score_dispersity"][idx] = pred0_with_rand["score_dispersity"][idx][valid_mask]
        pred0_with_rand["num_det"].append(valid_mask.sum())

        dist = compute_keypoints_distance(kpts1_wh_.detach(), kpts1_wh_.detach())
        local_mask = dist < radius
        valid_cnt = torch.sum(local_mask, dim=1)
        indices_need_nms = torch.where(valid_cnt > 1)[0]
        for i in indices_need_nms:
            if valid_cnt[i] > 0:
                kpt_indices = torch.where(local_mask[i])[0]
                scs_max_idx = scores_kpts1[kpt_indices].argmax()

                tmp_mask = kpt_indices.new_ones(len(kpt_indices)).bool()
                tmp_mask[scs_max_idx] = False
                suppressed_indices = kpt_indices[tmp_mask]

                valid_cnt[suppressed_indices] = 0

        valid_mask = valid_cnt > 0
        kpts1_wh = kpts1_wh_[valid_mask]
        kpts1 = kpts1[valid_mask]
        scores_kpts1 = scores_kpts1[valid_mask]
        pred1_with_rand["keypoints"][idx] = kpts1

        valid_mask = valid_mask[: len(pred1_with_rand["score_dispersity"][idx])]
        pred1_with_rand["score_dispersity"][idx] = pred1_with_rand["score_dispersity"][idx][valid_mask]
        pred1_with_rand["num_det"].append(valid_mask.sum())

        # del dist, local_mask, valid_cnt, indices_need_nms, scs_max_idx, tmp_mask, suppressed_indices, valid_mask
        # torch.cuda.empty_cache()
        # ========================= nms

        pred0_with_rand["scores"].append(scores_kpts0)
        pred1_with_rand["scores"].append(scores_kpts1)

        descriptor_map0, descriptor_map1 = pred0["descriptor_map"][idx], pred1["descriptor_map"][idx]
        descriptor_map0 = F.normalize(descriptor_map0, dim=0)
        descriptor_map1 = F.normalize(descriptor_map1, dim=0)

        desc0 = torch.nn.functional.grid_sample(
            descriptor_map0.unsqueeze(0), kpts0.view(1, 1, -1, 2), mode="bilinear", align_corners=True
        )[0, :, 0, :].t()
        desc1 = torch.nn.functional.grid_sample(
            descriptor_map1.unsqueeze(0), kpts1.view(1, 1, -1, 2), mode="bilinear", align_corners=True
        )[0, :, 0, :].t()
        desc0 = F.normalize(desc0, dim=-1)
        desc1 = F.normalize(desc1, dim=-1)

        pred0_with_rand["descriptors"].append(desc0)
        pred1_with_rand["descriptors"].append(desc1)

        # =========================== prepare warp parameters
        warp01_params = {}
        for k, v in batch["warp01_params"].items():
            warp01_params[k] = v[idx]
        warp10_params = {}
        for k, v in batch["warp10_params"].items():
            warp10_params[k] = v[idx]

        # =========================== warp keypoints across images
        try:
            kpts0_wh, kpts01_wh, ids0, ids0_out = warp(kpts0_wh, warp01_params)
            kpts1_wh, kpts10_wh, ids1, ids1_out = warp(kpts1_wh, warp10_params)
        except:
            correspondences.append(
                {
                    "correspondence0": None,
                    "correspondence1": None,
                    "dist": kpts0_wh.new_tensor(0),
                }
            )
            continue

        if debug:
            from training.utils import save_image_in_actual_size

            image0 = batch["image0"][idx].cpu().detach().numpy().transpose(1, 2, 0)
            image1 = batch["image1"][idx].cpu().detach().numpy().transpose(1, 2, 0)

            p0 = kpts0_wh[:, [1, 0]].cpu().detach().numpy()
            img_kpts0 = plot_keypoints(image0, p0, radius=5, color=(255, 0, 0))
            # display_image_in_actual_size(img_kpts0)

            p1 = kpts1_wh[:, [1, 0]].cpu().detach().numpy()
            img_kpts1 = plot_keypoints(image1, p1, radius=5, color=(255, 0, 0))
            # display_image_in_actual_size(img_kpts1)

            p01 = kpts01_wh[:, [1, 0]].cpu().detach().numpy()
            img_kpts01 = plot_keypoints(img_kpts1, p01, radius=5, color=(0, 255, 0))
            save_image_in_actual_size(img_kpts01, name="kpts01.png")

            p10 = kpts10_wh[:, [1, 0]].cpu().detach().numpy()
            img_kpts10 = plot_keypoints(img_kpts0, p10, radius=5, color=(0, 255, 0))
            save_image_in_actual_size(img_kpts10, name="kpts10.png")

        # ============================= compute reprojection error
        dist01 = compute_keypoints_distance(kpts0_wh, kpts10_wh)
        dist10 = compute_keypoints_distance(kpts1_wh, kpts01_wh)

        dist_l2 = (dist01 + dist10.t()) / 2.0
        # find mutual correspondences by calculating the distance
        # between keypoints (I1) and warpped keypoints (I2->I1)
        mutual_min_indices = mutual_argmin(dist_l2)

        dist_mutual_min = dist_l2[mutual_min_indices]
        valid_dist_mutual_min = dist_mutual_min.detach() < train_gt_th

        ids0_d = mutual_min_indices[0][valid_dist_mutual_min]
        ids1_d = mutual_min_indices[1][valid_dist_mutual_min]

        correspondence0 = ids0[ids0_d]
        correspondence1 = ids1[ids1_d]

        # L1 distance
        dist01_l1 = compute_keypoints_distance(kpts0_wh, kpts10_wh, p=1)
        dist10_l1 = compute_keypoints_distance(kpts1_wh, kpts01_wh, p=1)

        dist_l1 = (dist01_l1 + dist10_l1.t()) / 2.0

        # =========================== compute cross image descriptor similarity_map
        similarity_map_01 = (desc0 @ descriptor_map1.reshape(h * w, 256).t()) * 20
        similarity_map_01 = similarity_map_01.softmax(dim=-2) * similarity_map_01.softmax(dim=-1)
        similarity_map_01 = similarity_map_01.reshape(desc0.shape[0], h, w)
        similarity_map_01 = torch.clamp(similarity_map_01, 1e-6, 1 - 1e-6)

        similarity_map_10 = (desc1 @ descriptor_map0.reshape(h * w, 256).t()) * 20
        similarity_map_10 = similarity_map_10.softmax(dim=-2) * similarity_map_10.softmax(dim=-1)
        similarity_map_10 = similarity_map_10.reshape(desc1.shape[0], h, w)
        similarity_map_10 = torch.clamp(similarity_map_10, 1e-6, 1 - 1e-6)

        similarity_map_01_valid = similarity_map_01[ids0]  # valid descriptors
        similarity_map_10_valid = similarity_map_10[ids1]

        kpts01 = 2 * kpts01_wh.detach() / wh - 1  # N0x2, (x,y), [-1,1]
        kpts10 = 2 * kpts10_wh.detach() / wh - 1  # N0x2, (x,y), [-1,1]

        correspondences.append(
            {
                "correspondence0": correspondence0,  # indices of matched kpts0 in all kpts
                "correspondence1": correspondence1,  # indices of matched kpts1 in all kpts
                "scores0": scores_kpts0[ids0],
                "scores1": scores_kpts1[ids1],
                "kpts01": kpts01,
                "kpts10": kpts10,  # warped valid kpts
                "ids0": ids0,
                "ids1": ids1,  # valid indices of kpts0 and kpts1
                "ids0_out": ids0_out,
                "ids1_out": ids1_out,
                "ids0_d": ids0_d,
                "ids1_d": ids1_d,  # match indices of valid kpts0 and kpts1
                "dist_l1": dist_l1,  # cross distance matrix of valid kpts using L1 norm
                "dist": dist_l2,  # cross distance matrix of valid kpts using L2 norm
                "similarity_map_01": similarity_map_01,  # all
                "similarity_map_10": similarity_map_10,  # all
                "similarity_map_01_valid": similarity_map_01_valid,  # valid
                "similarity_map_10_valid": similarity_map_10_valid,  # valid
            }
        )

    return correspondences, pred0_with_rand, pred1_with_rand


class EmptyTensorError(Exception):
    pass
