import argparse
from argparse import Namespace
import os
from pathlib import Path
import yaml
import immatch
from immatch.utils.localize_sfm_helper import *


def generate_covis_db_pairs(args):
    from immatch.utils.colmap.data_parsing import covis_pairs_from_nvm

    nvm_path = Path(args.dataset_dir) / "3D-models/all-merged/all.nvm"
    covis_pairs_from_nvm(nvm_path, args.pair_dir, topk=args.topk)


def eval_robotcar(args):
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # Initialize Model
    config_file = f"configs/{args.config}.yml"
    with open(config_file, "r") as f:
        model_conf = yaml.load(f, Loader=yaml.FullLoader)[args.benchmark_name]
        class_name = model_conf["class"]
        print(f"Method:{class_name} Conf: {model_conf}")
    model = immatch.__dict__[class_name](model_conf)
    matcher = lambda im1, im2: model.match_pairs(im1, im2)

    # Merge args
    args = Namespace(**vars(args), **model_conf)
    args.model_tag = model.name
    args.conf_tag = f"im{model.imsize}"
    if args.prefix:
        args.conf_tag += f".{args.prefix}"

    # Load db & query pairs
    pair_list = load_pairs(args)
    if not pair_list:
        print("No pairs for localizations! Please turn on --generate_covis_db_pairs!")
        return

    # Initialize experiment paths
    args.dataset_dir = Path(args.dataset_dir)
    args.im_dir = args.dataset_dir / "images"
    args = init_paths(args)
    print(args)

    # Match pairs
    if not args.skip_match:
        match_pairs_exporth5(pair_list, matcher, args.im_dir, args.output_dir)

    # Extract keypoints
    process_matches_and_keypoints_exporth5(
        pair_list,
        args.output_dir,
        args.result_dir,
        qt_psize=args.qt_psize,
        qt_dthres=args.qt_dthres,
        sc_thres=args.sc_thres,
    )

    # Localization
    init_empty_sfm(args)
    reconstruct_database_pairs(args)
    localize_queries(args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Localize RobotCar Seasons")
    parser.add_argument("--gpu", "-gpu", type=str, default="0")
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--prefix", type=str, default=None)
    parser.add_argument("--colmap", type=str, required=True)
    parser.add_argument("--skip_match", action="store_true")
    parser.add_argument("--dataset_dir", type=str, default="data/datasets/RobotCar")
    parser.add_argument("--pair_dir", type=str, default="data/pairs")
    parser.add_argument("--benchmark_name", type=str, default="robotcar")

    # Turn on to only generate covis db pairs
    parser.add_argument("--generate_covis_db_pairs", action="store_true")
    parser.add_argument("--topk", type=int, default=20)

    args = parser.parse_args()

    # Generate covis pairs
    if args.generate_covis_db_pairs:
        generate_covis_db_pairs(args)
    else:
        eval_robotcar(args)
