import cv2
import os
from tqdm import tqdm
import torch
import numpy as np
import sys
import poselib

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

import argparse
import datetime

parser = argparse.ArgumentParser(description="HPatch dataset evaluation script")
parser.add_argument("--name", type=str, default="LiftFeat", help="experiment name")
parser.add_argument("--gpu", type=str, default="0", help="GPU ID")
args = parser.parse_args()

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

top_k = None
n_i = 52
n_v = 56

DATASET_ROOT = os.path.join(os.path.dirname(__file__), "../data/HPatch")

from evaluation.eval_utils import *
from models.liftfeat_wrapper import LiftFeat


poselib_config = {"ransac_th": 3.0, "options": {}}


class PoseLibHomographyEstimator:
    def __init__(self, conf):
        self.conf = conf

    def estimate(self, mkpts0, mkpts1):
        M, info = poselib.estimate_homography(
            mkpts0,
            mkpts1,
            {
                "max_reproj_error": self.conf["ransac_th"],
                **self.conf["options"],
            },
        )
        success = M is not None
        if not success:
            M = np.eye(3, dtype=np.float32)
            inl = np.zeros(mkpts0.shape[0], dtype=np.bool_)
        else:
            inl = info["inliers"]

        estimation = {
            "success": success,
            "M_0to1": M,
            "inliers": inl,
        }

        return estimation


estimator = PoseLibHomographyEstimator(poselib_config)


def poselib_homography_estimate(mkpts0, mkpts1):
    data = estimator.estimate(mkpts0, mkpts1)
    return data


def generate_standard_image(img, target_size=(1920, 1080)):
    sh, sw = img.shape[0], img.shape[1]
    rh, rw = float(target_size[1]) / float(sh), float(target_size[0]) / float(sw)
    ratio = min(rh, rw)
    nh, nw = int(ratio * sh), int(ratio * sw)
    ph, pw = target_size[1] - nh, target_size[0] - nw
    nimg = cv2.resize(img, (nw, nh))
    nimg = cv2.copyMakeBorder(nimg, 0, ph, 0, pw, cv2.BORDER_CONSTANT, value=(0, 0, 0))

    return nimg, ratio, ph, pw


def benchmark_features(match_fn):
    lim = [1, 9]
    rng = np.arange(lim[0], lim[1] + 1)

    seq_names = sorted(os.listdir(DATASET_ROOT))

    n_feats = []
    n_matches = []
    seq_type = []
    i_err = {thr: 0 for thr in rng}
    v_err = {thr: 0 for thr in rng}

    i_err_homo = {thr: 0 for thr in rng}
    v_err_homo = {thr: 0 for thr in rng}

    for seq_idx, seq_name in tqdm(enumerate(seq_names), total=len(seq_names)):
        # load reference image
        ref_img = cv2.imread(os.path.join(DATASET_ROOT, seq_name, "1.ppm"))
        ref_img_shape = ref_img.shape

        # load query images
        for im_idx in range(2, 7):
            # read ground-truth homography
            homography = np.loadtxt(os.path.join(DATASET_ROOT, seq_name, "H_1_" + str(im_idx)))
            query_img = cv2.imread(os.path.join(DATASET_ROOT, seq_name, f"{im_idx}.ppm"))

            mkpts_a, mkpts_b = match_fn(ref_img, query_img)

            pos_a = mkpts_a
            pos_a_h = np.concatenate([pos_a, np.ones([pos_a.shape[0], 1])], axis=1)
            pos_b_proj_h = np.transpose(np.dot(homography, np.transpose(pos_a_h)))
            pos_b_proj = pos_b_proj_h[:, :2] / pos_b_proj_h[:, 2:]

            pos_b = mkpts_b

            dist = np.sqrt(np.sum((pos_b - pos_b_proj) ** 2, axis=1))

            n_matches.append(pos_a.shape[0])
            seq_type.append(seq_name[0])

            if dist.shape[0] == 0:
                dist = np.array([float("inf")])

            for thr in rng:
                if seq_name[0] == "i":
                    i_err[thr] += np.mean(dist <= thr)
                else:
                    v_err[thr] += np.mean(dist <= thr)

            # estimate homography
            gt_homo = homography
            pred_homo, _ = cv2.findHomography(mkpts_a, mkpts_b, cv2.USAC_MAGSAC)
            if pred_homo is None:
                homo_dist = np.array([float("inf")])
            else:
                corners = np.array(
                    [
                        [0, 0],
                        [ref_img_shape[1] - 1, 0],
                        [0, ref_img_shape[0] - 1],
                        [ref_img_shape[1] - 1, ref_img_shape[0] - 1],
                    ]
                )
                real_warped_corners = homo_trans(corners, gt_homo)
                warped_corners = homo_trans(corners, pred_homo)
                homo_dist = np.mean(np.linalg.norm(real_warped_corners - warped_corners, axis=1))

            for thr in rng:
                if seq_name[0] == "i":
                    i_err_homo[thr] += np.mean(homo_dist <= thr)
                else:
                    v_err_homo[thr] += np.mean(homo_dist <= thr)

    seq_type = np.array(seq_type)
    n_feats = np.array(n_feats)
    n_matches = np.array(n_matches)

    return i_err, v_err, i_err_homo, v_err_homo, [seq_type, n_feats, n_matches]


if __name__ == "__main__":
    errors = {}

    weights = os.path.join(os.path.dirname(__file__), "../weights/LiftFeat.pth")
    liftfeat = LiftFeat(weight=weights)

    errors = benchmark_features(liftfeat.match_liftfeat)

    i_err, v_err, i_err_hom, v_err_hom, _ = errors

    cur_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    print(f"\n==={cur_time}==={args.name}===")
    print(f"MHA@3 MHA@5 MHA@7")
    for thr in [3, 5, 7]:
        ill_err_hom = i_err_hom[thr] / (n_i * 5)
        view_err_hom = v_err_hom[thr] / (n_v * 5)
        print(f"{ill_err_hom * 100:.2f}%-{view_err_hom * 100:.2f}%")
