# Copyright 2020 Toyota Research Institute.  All rights reserved.
# Adapted from: https://github.com/rpautrat/SuperPoint/blob/master/superpoint/evaluations/descriptor_evaluation.py

import random
from glob import glob
from os import path as osp

import cv2
import numpy as np

from utils import warp_keypoints


def select_k_best(points, descriptors, k):
    """ Select the k most probable points (and strip their probability).
    points has shape (num_points, 3) where the last coordinate is the probability.

    Parameters
    ----------
    points: numpy.ndarray (N,3)
        Keypoint vector, consisting of (x,y,probability).
    descriptors: numpy.ndarray (N,256)
        Keypoint descriptors.
    k: int
        Number of keypoints to select, based on probability.
    Returns
    -------
    
    selected_points: numpy.ndarray (k,2)
        k most probable keypoints.
    selected_descriptors: numpy.ndarray (k,256)
        Descriptors corresponding to the k most probable keypoints.
    """
    sorted_prob = points[points[:, 2].argsort(), :2]
    sorted_desc = descriptors[points[:, 2].argsort(), :]
    start = min(k, points.shape[0])
    selected_points = sorted_prob[-start:, :]
    selected_descriptors = sorted_desc[-start:, :]
    return selected_points, selected_descriptors


def keep_shared_points(keypoints, descriptors, H, shape, keep_k_points=1000):
    """
    Compute a list of keypoints from the map, filter the list of points by keeping
    only the points that once mapped by H are still inside the shape of the map
    and keep at most 'keep_k_points' keypoints in the image.
    
    Parameters
    ----------
    keypoints: numpy.ndarray (N,3)
        Keypoint vector, consisting of (x,y,probability).
    descriptors: numpy.ndarray (N,256)
        Keypoint descriptors.
    H: numpy.ndarray (3,3)
        Homography.
    shape: tuple 
        Image shape.
    keep_k_points: int
        Number of keypoints to select, based on probability.

    Returns
    -------    
    selected_points: numpy.ndarray (k,2)
        k most probable keypoints.
    selected_descriptors: numpy.ndarray (k,256)
        Descriptors corresponding to the k most probable keypoints.
    """
    
    def keep_true_keypoints(points, descriptors, H, shape):
        """ Keep only the points whose warped coordinates by H are still inside shape. """
        warped_points = warp_keypoints(points[:, [1, 0]], H)
        warped_points[:, [0, 1]] = warped_points[:, [1, 0]]
        mask = (warped_points[:, 0] >= 0) & (warped_points[:, 0] < shape[0]) &\
               (warped_points[:, 1] >= 0) & (warped_points[:, 1] < shape[1])
        return points[mask, :], descriptors[mask, :]

    selected_keypoints, selected_descriptors = keep_true_keypoints(keypoints, descriptors, H, shape)
    selected_keypoints, selected_descriptors = select_k_best(selected_keypoints, selected_descriptors, keep_k_points)
    return selected_keypoints, selected_descriptors


def compute_matching_score(data, keep_k_points=1000):
    """
    Compute the matching score between two sets of keypoints with associated descriptors.
    
    Parameters
    ----------
    data: dict
        Input dictionary containing:
        image_shape: tuple (H,W)
            Original image shape.
        homography: numpy.ndarray (3,3)
            Ground truth homography.
        prob: numpy.ndarray (N,3)
            Keypoint vector, consisting of (x,y,probability).
        warped_prob: numpy.ndarray (N,3)
            Warped keypoint vector, consisting of (x,y,probability).
        desc: numpy.ndarray (N,256)
            Keypoint descriptors.
        warped_desc: numpy.ndarray (N,256)
            Warped keypoint descriptors.
    keep_k_points: int
        Number of keypoints to select, based on probability.

    Returns
    -------    
    ms: float
        Matching score.
    """
    shape = data['image_shape']
    real_H = data['homography']

    # Filter out predictions
    keypoints = data['prob'][:, :2].T
    keypoints = keypoints[::-1]
    prob = data['prob'][:, 2]
    keypoints = np.stack([keypoints[0], keypoints[1], prob], axis=-1)

    warped_keypoints = data['warped_prob'][:, :2].T
    warped_keypoints = warped_keypoints[::-1]
    warped_prob = data['warped_prob'][:, 2]
    warped_keypoints = np.stack([warped_keypoints[0], warped_keypoints[1], warped_prob], axis=-1)

    desc = data['desc']
    warped_desc = data['warped_desc']
    
    # Keeps all points for the next frame. The matching for caculating M.Score shouldnt use only in view points.
    keypoints,        desc        = select_k_best(keypoints,               desc, keep_k_points)
    warped_keypoints, warped_desc = select_k_best(warped_keypoints, warped_desc, keep_k_points)
    
    # Match the keypoints with the warped_keypoints with nearest neighbor search
    # This part needs to be done with crossCheck=False.
    # All the matched pairs need to be evaluated without any selection.
    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)

    matches = bf.match(desc, warped_desc)
    matches_idx = np.array([m.queryIdx for m in matches])
    m_keypoints = keypoints[matches_idx, :]
    matches_idx = np.array([m.trainIdx for m in matches])
    m_warped_keypoints = warped_keypoints[matches_idx, :]

    true_warped_keypoints = warp_keypoints(m_warped_keypoints[:, [1, 0]], np.linalg.inv(real_H))[:,::-1]
    vis_warped = np.all((true_warped_keypoints >= 0) & (true_warped_keypoints <= (np.array(shape)-1)), axis=-1)
    norm1 = np.linalg.norm(true_warped_keypoints - m_keypoints, axis=-1)

    correct1 = (norm1 < 3)
    count1 = np.sum(correct1 * vis_warped)
    score1 = count1 / np.maximum(np.sum(vis_warped), 1.0)

    matches = bf.match(warped_desc, desc)
    matches_idx = np.array([m.queryIdx for m in matches])
    m_warped_keypoints = warped_keypoints[matches_idx, :]
    matches_idx = np.array([m.trainIdx for m in matches])
    m_keypoints = keypoints[matches_idx, :]

    true_keypoints = warp_keypoints(m_keypoints[:, [1, 0]], real_H)[:,::-1]
    vis = np.all((true_keypoints >= 0) & (true_keypoints <= (np.array(shape)-1)), axis=-1)
    norm2 = np.linalg.norm(true_keypoints - m_warped_keypoints, axis=-1)

    correct2 = (norm2 < 3)
    count2 = np.sum(correct2 * vis)
    score2 = count2 / np.maximum(np.sum(vis), 1.0)

    ms = (score1 + score2) / 2

    return ms

def compute_homography(data, keep_k_points=1000):
    """
    Compute the homography between 2 sets of Keypoints and descriptors inside data. 
    Use the homography to compute the correctness metrics (1,3,5).

    Parameters
    ----------
    data: dict
        Input dictionary containing:
        image_shape: tuple (H,W)
            Original image shape.
        homography: numpy.ndarray (3,3)
            Ground truth homography.
        prob: numpy.ndarray (N,3)
            Keypoint vector, consisting of (x,y,probability).
        warped_prob: numpy.ndarray (N,3)
            Warped keypoint vector, consisting of (x,y,probability).
        desc: numpy.ndarray (N,256)
            Keypoint descriptors.
        warped_desc: numpy.ndarray (N,256)
            Warped keypoint descriptors.
    keep_k_points: int
        Number of keypoints to select, based on probability.

    Returns
    -------    
    correctness1: float
        correctness1 metric.
    correctness3: float
        correctness3 metric.
    correctness5: float
        correctness5 metric.
    """
    shape = data['image_shape']
    real_H = data['homography']

    # Filter out predictions
    keypoints = data['prob'][:, :2].T
    keypoints = keypoints[::-1]
    prob = data['prob'][:, 2]
    keypoints = np.stack([keypoints[0], keypoints[1], prob], axis=-1)

    warped_keypoints = data['warped_prob'][:, :2].T
    warped_keypoints = warped_keypoints[::-1]
    warped_prob = data['warped_prob'][:, 2]
    warped_keypoints = np.stack([warped_keypoints[0], warped_keypoints[1], warped_prob], axis=-1)

    desc = data['desc']
    warped_desc = data['warped_desc']
    
    # Keeps only the points shared between the two views
    keypoints, desc = keep_shared_points(keypoints, desc, real_H, shape, keep_k_points)
    warped_keypoints, warped_desc = keep_shared_points(warped_keypoints, warped_desc, np.linalg.inv(real_H), shape,
                                                       keep_k_points)

    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    matches = bf.match(desc, warped_desc)
    matches_idx = np.array([m.queryIdx for m in matches])
    m_keypoints = keypoints[matches_idx, :]
    matches_idx = np.array([m.trainIdx for m in matches])
    m_warped_keypoints = warped_keypoints[matches_idx, :]

    # Estimate the homography between the matches using RANSAC
    H, _ = cv2.findHomography(m_keypoints[:, [1, 0]],
                              m_warped_keypoints[:, [1, 0]], cv2.RANSAC, 3, maxIters=5000)

    if H is None:
        return 0, 0, 0

    shape = shape[::-1]

    # Compute correctness
    corners = np.array([[0, 0, 1],
                        [0, shape[1] - 1, 1],
                        [shape[0] - 1, 0, 1],
                        [shape[0] - 1, shape[1] - 1, 1]])
    real_warped_corners = np.dot(corners, np.transpose(real_H))
    real_warped_corners = real_warped_corners[:, :2] / real_warped_corners[:, 2:]
    warped_corners = np.dot(corners, np.transpose(H))
    warped_corners = warped_corners[:, :2] / warped_corners[:, 2:]
    
    mean_dist = np.mean(np.linalg.norm(real_warped_corners - warped_corners, axis=1))
    correctness1 = float(mean_dist <= 1)
    correctness3 = float(mean_dist <= 3)
    correctness5 = float(mean_dist <= 5)

    return correctness1, correctness3, correctness5
