from pathlib import Path
import argparse

from ... import extract_features, match_features, localize_sfm, logger
from .utils import get_timestamps, delete_unused_images
from .utils import generate_query_lists, generate_localization_pairs
from .utils import prepare_submission, evaluate_submission

relocalization_files = {
    'training': 'RelocalizationFilesTrain//relocalizationFile_recording_2020-03-24_17-36-22.txt',
    'validation': 'RelocalizationFilesVal/relocalizationFile_recording_2020-03-03_12-03-23.txt',
    'test0': 'RelocalizationFilesTest/relocalizationFile_recording_2020-03-24_17-45-31_*.txt',
    'test1': 'RelocalizationFilesTest/relocalizationFile_recording_2020-04-23_19-37-00_*.txt',
}

parser = argparse.ArgumentParser()
parser.add_argument('--sequence', type=str, required=True,
                    choices=['training', 'validation', 'test0', 'test1'],
                    help='Sequence to be relocalized.')
parser.add_argument('--dataset', type=Path, default='datasets/4Seasons',
                    help='Path to the dataset, default: %(default)s')
parser.add_argument('--outputs', type=Path, default='outputs/4Seasons',
                    help='Path to the output directory, default: %(default)s')
args = parser.parse_args()
sequence = args.sequence

data_dir = args.dataset
ref_dir = data_dir / 'reference'
assert ref_dir.exists(), f'{ref_dir} does not exist'
seq_dir = data_dir / sequence
assert seq_dir.exists(), f'{seq_dir} does not exist'
seq_images = seq_dir / 'undistorted_images'
reloc = ref_dir / relocalization_files[sequence]

output_dir = args.outputs
output_dir.mkdir(exist_ok=True, parents=True)
query_list = output_dir / f'{sequence}_queries_with_intrinsics.txt'
ref_pairs = output_dir / 'pairs-db-dist20.txt'
ref_sfm = output_dir / 'sfm_superpoint+superglue'
results_path = output_dir / f'localization_{sequence}_hloc+superglue.txt'
submission_dir = output_dir / 'submission_hloc+superglue'

num_loc_pairs = 10
loc_pairs = output_dir / f'pairs-query-{sequence}-dist{num_loc_pairs}.txt'

fconf = extract_features.confs['superpoint_max']
mconf = match_features.confs['superglue']

# Not all query images that are used for the evaluation
# To save time in feature extraction, we delete unsused images.
timestamps = get_timestamps(reloc, 1)
delete_unused_images(seq_images, timestamps)

# Generate a list of query images with their intrinsics.
generate_query_lists(timestamps, seq_dir, query_list)

# Generate the localization pairs from the given reference frames.
generate_localization_pairs(
    sequence, reloc, num_loc_pairs, ref_pairs, loc_pairs)

# Extract, match, amd localize.
ffile = extract_features.main(fconf, seq_images, output_dir)
mfile = match_features.main(mconf, loc_pairs, fconf['output'], output_dir)
localize_sfm.main(
    ref_sfm, query_list, loc_pairs, ffile, mfile, results_path)

# Convert the absolute poses to relative poses with the reference frames.
submission_dir.mkdir(exist_ok=True)
prepare_submission(results_path, reloc, ref_dir / 'poses.txt', submission_dir)

# If not a test sequence: evaluation the localization accuracy
if 'test' not in sequence:
    logger.info('Evaluating the relocalization submission...')
    evaluate_submission(submission_dir, reloc)
