import cv2
import warnings
from pathlib import Path
from hloc import logger
from hloc import matchers, extractors, logger
from hloc import match_dense, match_features, extract_features
from hloc.utils.viz import add_text, plot_keypoints
from common.utils import (
    load_config,
    get_model,
    get_feature_model,
    ransac_zoo,
    get_matcher_zoo,
    filter_matches,
    device,
    ROOT,
)
from common.viz import (
    fig2im,
    plot_images,
    display_matches,
    plot_color_line_matches,
)
import time
import matplotlib.pyplot as plt

warnings.simplefilter("ignore")


def test_modules(config: dict):
    img_path1 = ROOT / "datasets/sacre_coeur/mapping/02928139_3448003521.jpg"
    img_path2 = ROOT / "datasets/sacre_coeur/mapping/17295357_9106075285.jpg"
    image0 = cv2.imread(str(img_path1))
    image1 = cv2.imread(str(img_path2))
    keypoint_threshold = 0.0
    extract_max_keypoints = 2000
    match_threshold = 0.2
    log_path = ROOT / "experiments"
    log_path.mkdir(exist_ok=True, parents=True)

    matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])
    for k, v in matcher_zoo_restored.items():
        if image0 is None or image1 is None:
            logger.error("Error: No images found! Please upload two images.")
        # init output
        output_keypoints = None
        output_matches_raw = None
        output_matches_ransac = None
        match_conf = v["matcher"]

        # update match config
        match_conf["model"]["match_threshold"] = match_threshold
        match_conf["model"]["max_keypoints"] = extract_max_keypoints
        matcher = get_model(match_conf)
        t1 = time.time()
        if v["dense"]:
            pred = match_dense.match_images(
                matcher,
                image0,
                image1,
                match_conf["preprocessing"],
                device=device,
            )
            del matcher
            extract_conf = None
            last_fixed = "{}".format(match_conf["model"]["name"])
        else:
            extract_conf = v["feature"]

            # update extract config
            extract_conf["model"]["max_keypoints"] = extract_max_keypoints
            extract_conf["model"]["keypoint_threshold"] = keypoint_threshold
            extractor = get_feature_model(extract_conf)
            pred0 = extract_features.extract(
                extractor, image0, extract_conf["preprocessing"]
            )
            pred1 = extract_features.extract(
                extractor, image1, extract_conf["preprocessing"]
            )
            pred = match_features.match_images(matcher, pred0, pred1)
            del extractor
            last_fixed = "{}_{}".format(
                extract_conf["model"]["name"], match_conf["model"]["name"]
            )

        # keypoints on images
        logger.info(f"Match features done using: {time.time()-t1:.3f}s")
        t1 = time.time()
        texts = [
            f"image pairs: {img_path1.name} & {img_path2.name}",
            "",
        ]
        titles = [
            "Image 0 - Keypoints",
            "Image 1 - Keypoints",
        ]
        output_keypoints = plot_images([image0, image1], titles=titles, dpi=300)
        if "keypoints0" in pred.keys() and "keypoints1" in pred.keys():
            plot_keypoints([pred["keypoints0"], pred["keypoints1"]])
            text = (
                f"# keypoints0: {len(pred['keypoints0'])} \n"
                + f"# keypoints1: {len(pred['keypoints1'])}"
            )
            add_text(0, text, fs=15)
        output_keypoints = fig2im(output_keypoints)

        # plot images with raw matches
        titles = [
            "Image 0 - Raw matched keypoints",
            "Image 1 - Raw matched keypoints",
        ]
        output_matches_raw, num_matches_raw = display_matches(
            pred, titles=titles
        )
        logger.info(f"Plot keypoints done using: {time.time()-t1:.3f}s")
        t1 = time.time()

        filter_matches(
            pred,
            ransac_method=config["defaults"]["ransac_method"],
            ransac_reproj_threshold=config["defaults"][
                "ransac_reproj_threshold"
            ],
            ransac_confidence=config["defaults"]["ransac_confidence"],
            ransac_max_iter=config["defaults"]["ransac_max_iter"],
        )
        # plot images with ransac matches
        titles = [
            "Image 0 - Ransac matched keypoints",
            "Image 1 - Ransac matched keypoints",
        ]
        output_matches_ransac, num_matches_ransac = display_matches(
            pred, titles=titles
        )
        logger.info(f"RANSAC matches done using: {time.time()-t1:.3f}s")

        img_keypoints_path = log_path / f"img_keypoints_{last_fixed}.png"
        img_matches_raw_path = log_path / f"img_matches_raw_{last_fixed}.png"
        img_matches_ransac_path = (
            log_path / f"img_matches_ransac_{last_fixed}.png"
        )
        cv2.imwrite(str(img_keypoints_path), output_keypoints)
        cv2.imwrite(str(img_matches_raw_path), output_matches_raw)
        cv2.imwrite(str(img_matches_ransac_path), output_matches_ransac)

        plt.close("all")


if __name__ == "__main__":
    import argparse

    config = load_config(ROOT / "common/config.yaml")
    test_modules(config)