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')
