import torch
import numpy as np
import tqdm
from dkm.datasets import MegadepthBuilder
from dkm.utils import warp_kpts
from torch.utils.data import ConcatDataset


class MegadepthDenseBenchmark:
    def __init__(self, data_root="data/megadepth", h = 384, w = 512, num_samples = 2000, device=None) -> None:
        mega = MegadepthBuilder(data_root=data_root)
        self.dataset = ConcatDataset(
            mega.build_scenes(split="test_loftr", ht=h, wt=w)
        )  # fixed resolution of 384,512
        self.num_samples = num_samples
        if device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.device = device

    def geometric_dist(self, depth1, depth2, T_1to2, K1, K2, dense_matches):
        b, h1, w1, d = dense_matches.shape
        with torch.no_grad():
            x1 = dense_matches[..., :2].reshape(b, h1 * w1, 2)
            # x1 = torch.stack((2*x1[...,0]/w1-1,2*x1[...,1]/h1-1),dim=-1)
            mask, x2 = warp_kpts(
                x1.double(),
                depth1.double(),
                depth2.double(),
                T_1to2.double(),
                K1.double(),
                K2.double(),
            )
            x2 = torch.stack(
                (w1 * (x2[..., 0] + 1) / 2, h1 * (x2[..., 1] + 1) / 2), dim=-1
            )
            prob = mask.float().reshape(b, h1, w1)
        x2_hat = dense_matches[..., 2:]
        x2_hat = torch.stack(
            (w1 * (x2_hat[..., 0] + 1) / 2, h1 * (x2_hat[..., 1] + 1) / 2), dim=-1
        )
        gd = (x2_hat - x2.reshape(b, h1, w1, 2)).norm(dim=-1)
        gd = gd[prob == 1]
        pck_1 = (gd < 1.0).float().mean()
        pck_3 = (gd < 3.0).float().mean()
        pck_5 = (gd < 5.0).float().mean()
        gd = gd.mean()
        return gd, pck_1, pck_3, pck_5

    def benchmark(self, model, batch_size=8):
        model.train(False)
        with torch.no_grad():
            gd_tot = 0.0
            pck_1_tot = 0.0
            pck_3_tot = 0.0
            pck_5_tot = 0.0
            sampler = torch.utils.data.WeightedRandomSampler(
                torch.ones(len(self.dataset)), replacement=False, num_samples=self.num_samples
            )
            dataloader = torch.utils.data.DataLoader(
                self.dataset, batch_size=8, num_workers=batch_size, sampler=sampler
            )
            for data in tqdm.tqdm(dataloader):
                im1, im2, depth1, depth2, T_1to2, K1, K2 = (
                    data["query"],
                    data["support"],
                    data["query_depth"].to(self.device),
                    data["support_depth"].to(self.device),
                    data["T_1to2"].to(self.device),
                    data["K1"].to(self.device),
                    data["K2"].to(self.device),
                )
                matches, certainty = model.match(im1, im2, batched=True)
                gd, pck_1, pck_3, pck_5 = self.geometric_dist(
                    depth1, depth2, T_1to2, K1, K2, matches
                )
                gd_tot, pck_1_tot, pck_3_tot, pck_5_tot = (
                    gd_tot + gd,
                    pck_1_tot + pck_1,
                    pck_3_tot + pck_3,
                    pck_5_tot + pck_5,
                )
        return {
            "mega_pck_1": pck_1_tot.item() / len(dataloader),
            "mega_pck_3": pck_3_tot.item() / len(dataloader),
            "mega_pck_5": pck_5_tot.item() / len(dataloader),
        }