import argparse
from pathlib import Path
import numpy as np
import scipy.spatial

from . import logger
from .utils.read_write_model import read_images_binary
from .pairs_from_retrieval import pairs_from_score_matrix

DEFAULT_ROT_THRESH = 30  # in degrees


def get_pairwise_distances(images):
    ids = np.array(list(images.keys()))
    Rs = []
    ts = []
    for id_ in ids:
        image = images[id_]
        R = image.qvec2rotmat()
        t = image.tvec
        Rs.append(R)
        ts.append(t)
    Rs = np.stack(Rs, 0)
    ts = np.stack(ts, 0)

    # Invert the poses from world-to-camera to camera-to-world.
    Rs = Rs.transpose(0, 2, 1)
    ts = -(Rs @ ts[:, :, None])[:, :, 0]

    dist = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(ts))

    # Instead of computing the angle between two camera orientations,
    # we compute the angle between the principal axes, as two images rotated
    # around their principal axis still observe the same scene.
    axes = Rs[:, :, -1]
    dots = np.einsum('mi,ni->mn', axes, axes, optimize=True)
    dR = np.rad2deg(np.arccos(np.clip(dots, -1., 1.)))

    return ids, dist, dR


def main(model, output, num_matched, rotation_threshold=DEFAULT_ROT_THRESH):
    logger.info('Reading the COLMAP model...')
    images = read_images_binary(model / 'images.bin')

    logger.info(
        f'Obtaining pairwise distances between {len(images)} images...')
    ids, dist, dR = get_pairwise_distances(images)
    scores = -dist

    invalid = (dR >= rotation_threshold)
    np.fill_diagonal(invalid, True)
    pairs = pairs_from_score_matrix(scores, invalid, num_matched)
    pairs = [(images[ids[i]].name, images[ids[j]].name) for i, j in pairs]

    logger.info(f'Found {len(pairs)} pairs.')
    with open(output, 'w') as f:
        f.write('\n'.join(' '.join(p) for p in pairs))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True, type=Path)
    parser.add_argument('--output', required=True, type=Path)
    parser.add_argument('--num_matched', required=True, type=int)
    parser.add_argument('--rotation_threshold',
                        default=DEFAULT_ROT_THRESH, type=float)
    args = parser.parse_args()
    main(**args.__dict__)
