import numpy as np
import logging
from romav2.geometry import (
    compute_pose_error,
    pose_auc,
    estimate_pose_cv2_ransac,
    compute_relative_pose,
)
from PIL import Image
from tqdm import tqdm

logger = logging.getLogger(__name__)


class Mega1500:
    def __init__(self, data_root="data/megadepth") -> None:
        self.scene_names = [
            "0015_0.1_0.3.npz",
            "0015_0.3_0.5.npz",
            "0022_0.1_0.3.npz",
            "0022_0.3_0.5.npz",
            "0022_0.5_0.7.npz",
        ]
        self.scenes = [np.load(f"{data_root}/{scene}", allow_pickle=True) for scene in self.scene_names]
        self.data_root = data_root

    def benchmark(self, model, model_name=None):
        data_root = self.data_root
        tot_e_t, tot_e_R, tot_e_pose = [], [], []
        thresholds = [5, 10, 20]
        for scene_ind in range(len(self.scenes)):
            scene = self.scenes[scene_ind]
            pairs = scene["pair_infos"]
            intrinsics = scene["intrinsics"]
            poses = scene["poses"]
            im_paths = scene["image_paths"]
            pair_inds = range(len(pairs))
            for pairind in (pbar := tqdm(pair_inds, desc="Mega1500 eval")):
                idx1, idx2 = pairs[pairind][0]
                K1 = intrinsics[idx1].copy()
                T1 = poses[idx1].copy()
                R1, t1 = T1[:3, :3], T1[:3, 3]
                K2 = intrinsics[idx2].copy()
                T2 = poses[idx2].copy()
                R2, t2 = T2[:3, :3], T2[:3, 3]
                R, t = compute_relative_pose(R1, t1, R2, t2)
                im_A_path = f"{data_root}/{im_paths[idx1]}"
                im_B_path = f"{data_root}/{im_paths[idx2]}"
                preds = model.match(im_A_path, im_B_path)
                im_A = Image.open(im_A_path)
                w1, h1 = im_A.size
                im_B = Image.open(im_B_path)
                w2, h2 = im_B.size
                if (
                    True
                ):  # Note: we keep this true as it was used in DKM/RoMa papers. There is very little difference compared to setting to False.
                    scale1 = 1200 / max(w1, h1)
                    scale2 = 1200 / max(w2, h2)
                    w1, h1 = scale1 * w1, scale1 * h1
                    w2, h2 = scale2 * w2, scale2 * h2
                    K1, K2 = K1.copy(), K2.copy()
                    K1[:2] = K1[:2] * scale1
                    K2[:2] = K2[:2] * scale2
                for _ in range(5):
                    matches, _, _, _ = model.sample(preds, 5_000)
                    kpts1, kpts2 = model.to_pixel_coordinates(matches, h1, w1, h2, w2)
                    kpts1, kpts2 = kpts1.cpu().numpy(), kpts2.cpu().numpy()
                    shuffling = np.random.permutation(np.arange(len(kpts1)))
                    kpts1 = kpts1[shuffling]
                    kpts2 = kpts2[shuffling]
                    try:
                        threshold = 0.5
                        norm_threshold = threshold / (np.mean(np.abs(K1[:2, :2])) + np.mean(np.abs(K2[:2, :2])))
                        R_est, t_est, _ = estimate_pose_cv2_ransac(
                            kpts1,
                            kpts2,
                            K1,
                            K2,
                            norm_threshold,
                            conf=0.99999,
                        )
                        e_t, e_R = compute_pose_error(R_est, t_est[:, 0], R, t)
                        e_pose = max(e_t, e_R)
                    except Exception as e:
                        logger.debug(f"Pose estimation error: {e}")
                        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)
                pbar.set_postfix(auc=f"{[f'{a.item():.3f}' for a in pose_auc(tot_e_pose, thresholds)]}")

        tot_e_pose = np.array(tot_e_pose)
        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])
        logger.info("%s auc: %s", model_name, auc)
        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,
        }
