from pathlib import Path
import argparse

from .utils import create_reference_sfm
from .create_gt_sfm import correct_sfm_with_gt_depth
from ..Cambridge.utils import create_query_list_with_intrinsics, evaluate
from ... import extract_features, match_features, pairs_from_covisibility
from ... import triangulation, localize_sfm, logger

SCENES = ["chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs"]


def run_scene(
    images,
    gt_dir,
    retrieval,
    outputs,
    results,
    num_covis,
    use_dense_depth,
    depth_dir=None,
):
    outputs.mkdir(exist_ok=True, parents=True)
    ref_sfm_sift = outputs / "sfm_sift"
    ref_sfm = outputs / "sfm_superpoint+superglue"
    query_list = outputs / "query_list_with_intrinsics.txt"

    feature_conf = {
        "output": "feats-superpoint-n4096-r1024",
        "model": {
            "name": "superpoint",
            "nms_radius": 3,
            "max_keypoints": 4096,
        },
        "preprocessing": {
            "globs": ["*.color.png"],
            "grayscale": True,
            "resize_max": 1024,
        },
    }
    matcher_conf = match_features.confs["superglue"]
    matcher_conf["model"]["sinkhorn_iterations"] = 5

    test_list = gt_dir / "list_test.txt"
    create_reference_sfm(gt_dir, ref_sfm_sift, test_list)
    create_query_list_with_intrinsics(gt_dir, query_list, test_list)

    features = extract_features.main(
        feature_conf, images, outputs, as_half=True
    )

    sfm_pairs = outputs / f"pairs-db-covis{num_covis}.txt"
    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
    )
    if not (use_dense_depth and ref_sfm.exists()):
        triangulation.main(
            ref_sfm, ref_sfm_sift, images, sfm_pairs, features, sfm_matches
        )
    if use_dense_depth:
        assert depth_dir is not None
        ref_sfm_fix = outputs / "sfm_superpoint+superglue+depth"
        correct_sfm_with_gt_depth(ref_sfm, depth_dir, ref_sfm_fix)
        ref_sfm = ref_sfm_fix

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

    localize_sfm.main(
        ref_sfm,
        query_list,
        retrieval,
        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/7scenes",
        help="Path to the dataset, default: %(default)s",
    )
    parser.add_argument(
        "--outputs",
        type=Path,
        default="outputs/7scenes",
        help="Path to the output directory, default: %(default)s",
    )
    parser.add_argument("--use_dense_depth", action="store_true")
    parser.add_argument(
        "--num_covis",
        type=int,
        default=30,
        help="Number of image pairs for SfM, default: %(default)s",
    )
    args = parser.parse_args()

    gt_dirs = args.dataset / "7scenes_sfm_triangulated/{scene}/triangulated"
    retrieval_dirs = args.dataset / "7scenes_densevlad_retrieval_top_10"

    all_results = {}
    for scene in args.scenes:
        logger.info(f'Working on scene "{scene}".')
        results = (
            args.outputs
            / scene
            / "results_{}.txt".format(
                "dense" if args.use_dense_depth else "sparse"
            )
        )
        if args.overwrite or not results.exists():
            run_scene(
                args.dataset / scene,
                Path(str(gt_dirs).format(scene=scene)),
                retrieval_dirs / f"{scene}_top10.txt",
                args.outputs / scene,
                results,
                args.num_covis,
                args.use_dense_depth,
                depth_dir=args.dataset / f"depth/7scenes_{scene}/train/depth",
            )
        all_results[scene] = results

    for scene in args.scenes:
        logger.info(f'Evaluate scene "{scene}".')
        gt_dir = Path(str(gt_dirs).format(scene=scene))
        evaluate(gt_dir, all_results[scene], gt_dir / "list_test.txt")