from pathlib import Path
import argparse

from .utils import create_query_list_with_intrinsics, scale_sfm_images, evaluate
from ... import extract_features, match_features, pairs_from_covisibility
from ... import triangulation, localize_sfm, pairs_from_retrieval, logger

SCENES = ["KingsCollege", "OldHospital", "ShopFacade", "StMarysChurch", "GreatCourt"]


def run_scene(images, gt_dir, outputs, results, num_covis, num_loc):
    ref_sfm_sift = gt_dir / "model_train"
    test_list = gt_dir / "list_query.txt"

    outputs.mkdir(exist_ok=True, parents=True)
    ref_sfm = outputs / "sfm_superpoint+superglue"
    ref_sfm_scaled = outputs / "sfm_sift_scaled"
    query_list = outputs / "query_list_with_intrinsics.txt"
    sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt"
    loc_pairs = outputs / f"pairs-query-netvlad{num_loc}.txt"

    feature_conf = {
        "output": "feats-superpoint-n4096-r1024",
        "model": {
            "name": "superpoint",
            "nms_radius": 3,
            "max_keypoints": 4096,
        },
        "preprocessing": {
            "grayscale": True,
            "resize_max": 1024,
        },
    }
    matcher_conf = match_features.confs["superglue"]
    retrieval_conf = extract_features.confs["netvlad"]

    create_query_list_with_intrinsics(
        gt_dir / "empty_all", query_list, test_list, ext=".txt", image_dir=images
    )
    with open(test_list, "r") as f:
        query_seqs = {q.split("/")[0] for q in f.read().rstrip().split("\n")}

    global_descriptors = extract_features.main(retrieval_conf, images, outputs)
    pairs_from_retrieval.main(
        global_descriptors,
        loc_pairs,
        num_loc,
        db_model=ref_sfm_sift,
        query_prefix=query_seqs,
    )

    features = extract_features.main(feature_conf, images, outputs, as_half=True)
    pairs_from_covisibility.main(ref_sfm_sift, sfm_pairs, num_matched=num_covis)
    sfm_matches = match_features.main(
        matcher_conf, sfm_pairs, feature_conf["output"], outputs
    )

    scale_sfm_images(ref_sfm_sift, ref_sfm_scaled, images)
    triangulation.main(
        ref_sfm, ref_sfm_scaled, images, sfm_pairs, features, sfm_matches
    )

    loc_matches = match_features.main(
        matcher_conf, loc_pairs, feature_conf["output"], outputs
    )

    localize_sfm.main(
        ref_sfm,
        query_list,
        loc_pairs,
        features,
        loc_matches,
        results,
        covisibility_clustering=False,
        prepend_camera_name=True,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--scenes", default=SCENES, choices=SCENES, nargs="+")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--dataset",
        type=Path,
        default="datasets/cambridge",
        help="Path to the dataset, default: %(default)s",
    )
    parser.add_argument(
        "--outputs",
        type=Path,
        default="outputs/cambridge",
        help="Path to the output directory, default: %(default)s",
    )
    parser.add_argument(
        "--num_covis",
        type=int,
        default=20,
        help="Number of image pairs for SfM, default: %(default)s",
    )
    parser.add_argument(
        "--num_loc",
        type=int,
        default=10,
        help="Number of image pairs for loc, default: %(default)s",
    )
    args = parser.parse_args()

    gt_dirs = args.dataset / "CambridgeLandmarks_Colmap_Retriangulated_1024px"

    all_results = {}
    for scene in args.scenes:
        logger.info(f'Working on scene "{scene}".')
        results = args.outputs / scene / "results.txt"
        if args.overwrite or not results.exists():
            run_scene(
                args.dataset / scene,
                gt_dirs / scene,
                args.outputs / scene,
                results,
                args.num_covis,
                args.num_loc,
            )
        all_results[scene] = results

    for scene in args.scenes:
        logger.info(f'Evaluate scene "{scene}".')
        evaluate(
            gt_dirs / scene / "empty_all",
            all_results[scene],
            gt_dirs / scene / "list_query.txt",
            ext=".txt",
        )