import pickle
import h5py
import numpy as np
import torch
from dkm.utils import *
from PIL import Image
from tqdm import tqdm


class Yfcc100mBenchmark:
    def __init__(self, data_root="data/yfcc100m_test") -> None:
        self.scenes = [
            "buckingham_palace",
            "notre_dame_front_facade",
            "reichstag",
            "sacre_coeur",
        ]
        self.data_root = data_root

    def benchmark(self, model, r=2):
        model.train(False)
        with torch.no_grad():
            data_root = self.data_root
            meta_info = open(
                f"{data_root}/yfcc_test_pairs_with_gt.txt", "r"
            ).readlines()
            tot_e_t, tot_e_R, tot_e_pose = [], [], []
            for scene_ind in range(len(self.scenes)):
                scene = self.scenes[scene_ind]
                pairs = np.array(
                    pickle.load(
                        open(f"{data_root}/pairs/{scene}-te-1000-pairs.pkl", "rb")
                    )
                )
                scene_dir = f"{data_root}/yfcc100m/{scene}/test/"
                calibs = open(scene_dir + "calibration.txt", "r").read().split("\n")
                images = open(scene_dir + "images.txt", "r").read().split("\n")
                pair_inds = np.random.choice(
                    range(len(pairs)), size=len(pairs), replace=False
                )
                for pairind in tqdm(pair_inds):
                    idx1, idx2 = pairs[pairind]
                    params = meta_info[1000 * scene_ind + pairind].split()
                    rot1, rot2 = int(params[2]), int(params[3])
                    calib1 = h5py.File(scene_dir + calibs[idx1], "r")
                    K1, R1, t1, _, _ = get_pose(calib1)
                    calib2 = h5py.File(scene_dir + calibs[idx2], "r")
                    K2, R2, t2, _, _ = get_pose(calib2)

                    R, t = compute_relative_pose(R1, t1, R2, t2)
                    im1 = images[idx1]
                    im2 = images[idx2]
                    im1 = Image.open(scene_dir + im1).rotate(rot1 * 90, expand=True)
                    w1, h1 = im1.size
                    im2 = Image.open(scene_dir + im2).rotate(rot2 * 90, expand=True)
                    w2, h2 = im2.size
                    K1 = rotate_intrinsic(K1, rot1)
                    K2 = rotate_intrinsic(K2, rot2)

                    dense_matches, dense_certainty = model.match(im1, im2)
                    dense_certainty = dense_certainty ** (1 / r)
                    sparse_matches, sparse_confidence = model.sample(
                        dense_matches, dense_certainty, 10000
                    )
                    scale1 = 480 / min(w1, h1)
                    scale2 = 480 / min(w2, h2)
                    w1, h1 = scale1 * w1, scale1 * h1
                    w2, h2 = scale2 * w2, scale2 * h2
                    K1 = K1 * scale1
                    K2 = K2 * scale2

                    kpts1 = sparse_matches[:, :2]
                    kpts1 = np.stack(
                        (w1 * kpts1[:, 0] / 2, h1 * kpts1[:, 1] / 2), axis=-1
                    )
                    kpts2 = sparse_matches[:, 2:]
                    kpts2 = np.stack(
                        (w2 * kpts2[:, 0] / 2, h2 * kpts2[:, 1] / 2), axis=-1
                    )
                    try:
                        threshold = 1.0
                        norm_threshold = threshold / (
                            np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2]))
                        )
                        R_est, t_est, mask = estimate_pose(
                            kpts1,
                            kpts2,
                            K1[:2, :2],
                            K2[:2, :2],
                            norm_threshold,
                            conf=0.9999999,
                        )
                        T1_to_2 = np.concatenate((R_est, t_est), axis=-1)  #
                        e_t, e_R = compute_pose_error(T1_to_2, R, t)
                        e_pose = max(e_t, e_R)
                    except:
                        e_t, e_R = 90, 90
                        e_pose = max(e_t, e_R)
                    tot_e_t.append(e_t)
                    tot_e_R.append(e_R)
                    tot_e_pose.append(e_pose)
            tot_e_pose = np.array(tot_e_pose)
            thresholds = [5, 10, 20]
            auc = pose_auc(tot_e_pose, thresholds)
            acc_5 = (tot_e_pose < 5).mean()
            acc_10 = (tot_e_pose < 10).mean()
            acc_15 = (tot_e_pose < 15).mean()
            acc_20 = (tot_e_pose < 20).mean()
            map_5 = acc_5
            map_10 = np.mean([acc_5, acc_10])
            map_20 = np.mean([acc_5, acc_10, acc_15, acc_20])
            return {
                "auc_5": auc[0],
                "auc_10": auc[1],
                "auc_20": auc[2],
                "map_5": map_5,
                "map_10": map_10,
                "map_20": map_20,
            }