from PIL import Image
import numpy as np

import os

from tqdm import tqdm
from roma.utils import pose_auc
import cv2


class HpatchesHomogBenchmark:
    """Hpatches grid goes from [0,n-1] instead of [0.5,n-0.5]"""

    def __init__(self, dataset_path) -> None:
        seqs_dir = "hpatches-sequences-release"
        self.seqs_path = os.path.join(dataset_path, seqs_dir)
        self.seq_names = sorted(os.listdir(self.seqs_path))
        # Ignore seqs is same as LoFTR.
        self.ignore_seqs = set(
            [
                "i_contruction",
                "i_crownnight",
                "i_dc",
                "i_pencils",
                "i_whitebuilding",
                "v_artisans",
                "v_astronautis",
                "v_talent",
            ]
        )

    def convert_coordinates(self, im_A_coords, im_A_to_im_B, wq, hq, wsup, hsup):
        offset = 0.5  # Hpatches assumes that the center of the top-left pixel is at [0,0] (I think)
        im_A_coords = (
            np.stack(
                (
                    wq * (im_A_coords[..., 0] + 1) / 2,
                    hq * (im_A_coords[..., 1] + 1) / 2,
                ),
                axis=-1,
            )
            - offset
        )
        im_A_to_im_B = (
            np.stack(
                (
                    wsup * (im_A_to_im_B[..., 0] + 1) / 2,
                    hsup * (im_A_to_im_B[..., 1] + 1) / 2,
                ),
                axis=-1,
            )
            - offset
        )
        return im_A_coords, im_A_to_im_B

    def benchmark(self, model, model_name = None):
        n_matches = []
        homog_dists = []
        for seq_idx, seq_name in tqdm(
            enumerate(self.seq_names), total=len(self.seq_names)
        ):
            im_A_path = os.path.join(self.seqs_path, seq_name, "1.ppm")
            im_A = Image.open(im_A_path)
            w1, h1 = im_A.size
            for im_idx in range(2, 7):
                im_B_path = os.path.join(self.seqs_path, seq_name, f"{im_idx}.ppm")
                im_B = Image.open(im_B_path)
                w2, h2 = im_B.size
                H = np.loadtxt(
                    os.path.join(self.seqs_path, seq_name, "H_1_" + str(im_idx))
                )
                dense_matches, dense_certainty = model.match(
                    im_A_path, im_B_path
                )
                good_matches, _ = model.sample(dense_matches, dense_certainty, 5000)
                pos_a, pos_b = self.convert_coordinates(
                    good_matches[:, :2], good_matches[:, 2:], w1, h1, w2, h2
                )
                try:
                    H_pred, inliers = cv2.findHomography(
                        pos_a,
                        pos_b,
                        method = cv2.RANSAC,
                        confidence = 0.99999,
                        ransacReprojThreshold = 3 * min(w2, h2) / 480,
                    )
                except:
                    H_pred = None
                if H_pred is None:
                    H_pred = np.zeros((3, 3))
                    H_pred[2, 2] = 1.0
                corners = np.array(
                    [[0, 0, 1], [0, h1 - 1, 1], [w1 - 1, 0, 1], [w1 - 1, h1 - 1, 1]]
                )
                real_warped_corners = np.dot(corners, np.transpose(H))
                real_warped_corners = (
                    real_warped_corners[:, :2] / real_warped_corners[:, 2:]
                )
                warped_corners = np.dot(corners, np.transpose(H_pred))
                warped_corners = warped_corners[:, :2] / warped_corners[:, 2:]
                mean_dist = np.mean(
                    np.linalg.norm(real_warped_corners - warped_corners, axis=1)
                ) / (min(w2, h2) / 480.0)
                homog_dists.append(mean_dist)

        n_matches = np.array(n_matches)
        thresholds = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        auc = pose_auc(np.array(homog_dists), thresholds)
        return {
            "hpatches_homog_auc_3": auc[2],
            "hpatches_homog_auc_5": auc[4],
            "hpatches_homog_auc_10": auc[9],
        }