import sys

sys.path.append(".")
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
import cv2
import argparse
from RDD.RDD_helper import RDD_helper
from RDD.RDD import build
import matplotlib.pyplot as plt
import matplotlib
import os
from numpy.linalg import inv
from benchmarks.utils import (
    pose_auc,
    compute_symmetrical_epipolar_errors,
    compute_pose_error,
    compute_relative_pose,
    estimate_pose,
    dynamic_alpha,
)


def read_scannet_pose(path):
    """Read ScanNet's Camera2World pose and transform it to World2Camera.

    Returns:
        pose_w2c (np.ndarray): (4, 4)
    """
    cam2world = np.loadtxt(path, delimiter=" ")
    world2cam = inv(cam2world)
    return world2cam


def make_matching_figure(img0, img1, mkpts0, mkpts1, color, kpts0=None, kpts1=None, text=[], dpi=75, path=None):
    # draw image pair
    assert mkpts0.shape[0] == mkpts1.shape[0], f"mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}"
    fig, axes = plt.subplots(1, 2, figsize=(10, 6), dpi=dpi)
    axes[0].imshow(img0, cmap="gray")
    axes[1].imshow(img1, cmap="gray")
    for i in range(2):  # clear all frames
        axes[i].get_yaxis().set_ticks([])
        axes[i].get_xaxis().set_ticks([])
        for spine in axes[i].spines.values():
            spine.set_visible(False)
    plt.tight_layout(pad=1)

    if kpts0 is not None:
        assert kpts1 is not None
        axes[0].scatter(kpts0[:, 0], kpts0[:, 1], c="w", s=2)
        axes[1].scatter(kpts1[:, 0], kpts1[:, 1], c="w", s=2)

    # draw matches
    if mkpts0.shape[0] != 0 and mkpts1.shape[0] != 0:
        fig.canvas.draw()
        transFigure = fig.transFigure.inverted()
        fkpts0 = transFigure.transform(axes[0].transData.transform(mkpts0))
        fkpts1 = transFigure.transform(axes[1].transData.transform(mkpts1))
        fig.lines = [
            matplotlib.lines.Line2D(
                (fkpts0[i, 0], fkpts1[i, 0]),
                (fkpts0[i, 1], fkpts1[i, 1]),
                transform=fig.transFigure,
                c=color[i],
                linewidth=1,
            )
            for i in range(len(mkpts0))
        ]

        axes[0].scatter(mkpts0[:, 0], mkpts0[:, 1], c=color, s=4)
        axes[1].scatter(mkpts1[:, 0], mkpts1[:, 1], c=color, s=4)

    # put txts
    txt_color = "k" if img0[:100, :200].mean() > 200 else "w"
    fig.text(
        0.01, 0.99, "\n".join(text), transform=fig.axes[0].transAxes, fontsize=15, va="top", ha="left", color=txt_color
    )

    # save or return figure
    if path:
        plt.savefig(str(path), bbox_inches="tight", pad_inches=0)
        plt.close()
    else:
        return fig


def error_colormap(err, thr, alpha=1.0):
    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
    x = 1 - np.clip(err / (thr * 2), 0, 1)
    return np.clip(np.stack([2 - x * 2, x * 2, np.zeros_like(x), np.ones_like(x) * alpha], -1), 0, 1)


def _make_evaluation_figure(img0, img1, kpts0, kpts1, epi_errs, e_t, e_R, alpha="dynamic", path=None):
    conf_thr = 1e-4

    img0 = np.array(img0)
    img1 = np.array(img1)

    kpts0 = kpts0
    kpts1 = kpts1

    epi_errs = epi_errs.cpu().numpy()
    correct_mask = epi_errs < conf_thr
    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
    n_correct = np.sum(correct_mask)

    # recall might be larger than 1, since the calculation of conf_matrix_gt
    # uses groundtruth depths and camera poses, but epipolar distance is used here.

    # matching info
    if alpha == "dynamic":
        alpha = dynamic_alpha(len(correct_mask))
    color = error_colormap(epi_errs, conf_thr, alpha=alpha)

    text = [
        f"#Matches {len(kpts0)}",
        f"Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}",
        f"e_t: {e_t:.2f} | e_R: {e_R:.2f}",
    ]

    # make the figure
    figure = make_matching_figure(img0, img1, kpts0, kpts1, color, text=text, path=path)
    return figure


class ScanNetPoseMNNBenchmark:
    def __init__(self, data_root="./scannet_test_1500") -> None:
        self.scenes = np.load(f"{data_root}/test.npz", allow_pickle=True)["name"]

        self.data_root = data_root
        self.intrinsics = np.load(f"{data_root}/intrinsics.npz", allow_pickle=True)

    def benchmark(
        self,
        model_helper,
        model_name=None,
        scale_intrinsics=False,
        calibrated=True,
        plot_every_iter=1,
        plot=False,
        method="sparse",
    ):

        with torch.no_grad():
            data_root = self.data_root
            tot_e_t, tot_e_R, tot_e_pose = [], [], []
            thresholds = [5, 10, 20]

            for pairind in tqdm(range(len(self.scenes))):

                scene_name, scene_sub_name, stem_name_0, stem_name_1 = self.scenes[pairind]

                seq_name = f"scene{scene_name:04d}_{scene_sub_name:02d}"

                K0 = K1 = self.intrinsics[seq_name].copy()

                T0 = read_scannet_pose(f"{data_root}/scene{scene_name:04d}_{scene_sub_name:02d}/pose/{stem_name_0}.txt")
                T1 = read_scannet_pose(f"{data_root}/scene{scene_name:04d}_{scene_sub_name:02d}/pose/{stem_name_1}.txt")
                R0, t0 = T0[:3, :3], T0[:3, 3]
                R1, t1 = T1[:3, :3], T1[:3, 3]
                R, t = compute_relative_pose(R0, t0, R1, t1)
                T0_to_1 = np.concatenate((R, t[:, None]), axis=-1)

                im_A_path = f"{data_root}/scene{scene_name:04d}_{scene_sub_name:02d}/color/{stem_name_0}.jpg"
                im_B_path = f"{data_root}/scene{scene_name:04d}_{scene_sub_name:02d}/color/{stem_name_1}.jpg"

                im_A = cv2.imread(im_A_path)
                im_B = cv2.imread(im_B_path)
                im_A = cv2.resize(im_A, (640, 480))
                im_B = cv2.resize(im_B, (640, 480))

                if method == "dense":
                    kpts0, kpts1, conf = model_helper.match_dense(im_A, im_B, thr=0.01, resize=None)
                elif method == "lightglue":
                    kpts0, kpts1, conf = model_helper.match_lg(im_A, im_B, thr=0.01, resize=None)
                elif method == "sparse":
                    kpts0, kpts1, conf = model_helper.match(im_A, im_B, thr=0.01, resize=None)
                else:
                    kpts0, kpts1, conf = model_helper.match_3rd_party(im_A, im_B, thr=0.01, resize=None, model=method)

                if scale_intrinsics:

                    scale0 = 840 / max(w0, h0)
                    scale1 = 840 / max(w1, h1)
                    w0, h0 = scale0 * w0, scale0 * h0
                    w1, h1 = scale1 * w1, scale1 * h1
                    K0, K1 = K0.copy(), K1.copy()
                    K0[:2] = K0[:2] * scale0
                    K1[:2] = K1[:2] * scale1

                threshold = 0.5
                if calibrated:
                    norm_threshold = threshold / (np.mean(np.abs(K0[:2, :2])) + np.mean(np.abs(K1[:2, :2])))
                    ret = estimate_pose(
                        kpts0,
                        kpts1,
                        K0,
                        K1,
                        norm_threshold,
                        conf=0.99999,
                    )
                if ret is not None:
                    R_est, t_est, mask = ret
                    T0_to_1_est = np.concatenate((R_est, t_est), axis=-1)  #
                    T0_to_1 = np.concatenate((R, t[:, None]), axis=-1)
                    e_t, e_R = compute_pose_error(T0_to_1_est, R, t)

                    epi_errs = compute_symmetrical_epipolar_errors(T0_to_1, kpts0, kpts1, K0, K1)
                    if pairind % plot_every_iter == 0 and plot:

                        if not os.path.exists(f"outputs/scannet/{model_name}_{method}"):
                            os.mkdir(f"outputs/scannet/{model_name}_{method}")
                        name = f"outputs/scannet/{model_name}_{method}/{scene_name}_{pairind}.png"
                        _make_evaluation_figure(im_A, im_B, kpts0, kpts1, epi_errs, e_t, e_R, path=name)
                    e_pose = max(e_t, e_R)
                else:
                    e_t, e_R = np.inf, np.inf  # or any large sentinel value to indicate failure
                    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)
            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])
            print(f"{model_name} auc: {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,
            }


def parse_arguments():
    parser = argparse.ArgumentParser(description="Testing script.")

    parser.add_argument(
        "--data_root", type=str, default="./data/scannet_test_1500", help="Path to the ScanNet dataset."
    )

    parser.add_argument("--weights", type=str, default="./weights/RDD-v2.pth", help="Path to the model checkpoint.")

    parser.add_argument("--plot", action="store_true", help="Whether to plot the results.")

    parser.add_argument("--method", type=str, default="sparse", help="Method for matching.")

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_arguments()
    if not os.path.exists("outputs"):
        os.mkdir("outputs")

    if not os.path.exists(f"outputs/scannet"):
        os.mkdir(f"outputs/scannet")

    model = build(weights=args.weights)
    benchmark = ScanNetPoseMNNBenchmark(data_root=args.data_root)
    model.eval()
    model_helper = RDD_helper(model)
    with torch.no_grad():
        method = args.method
        out = benchmark.benchmark(model_helper, model_name="RDD", plot_every_iter=1, plot=args.plot, method=method)
        with open(f"outputs/scannet/RDD_{method}.txt", "w") as f:
            f.write(str(out))
