# Author: Tomas Hodan (hodantom@cmp.felk.cvut.cz)
# Center for Machine Perception, Czech Technical University in Prague

"""Matching of estimated poses to the ground-truth poses."""

import numpy as np


def match_poses(errs, error_ths, max_ests_count=0, gt_valid_mask=None):
  """Matches the estimated poses to the ground-truth poses.

  The estimated poses are greedily matched to the ground truth poses in the
  order of decreasing score of the estimates. An estimated pose is matched to a
  ground-truth pose if the error w.r.t. the ground-truth pose is below the
  specified threshold. Each estimated pose is matched to up to one ground-truth
  pose and each ground-truth pose is matched to up to one estimated pose.

  :param errs: List of dictionaries, where each dictionary holds the following
    info about one pose estimate:
    - 'est_id': ID of the pose estimate.
    - 'score': Confidence score of the pose estimate.
    - 'errors': Dictionary mapping ground-truth ID's to errors of the pose
        estimate w.r.t. the ground-truth poses.
  :param error_ths: Thresholds of correctness. The pose error can be given
    by more than one element (e.g. translational + rotational error), in which
    case there is one threshold for each element.
  :param max_ests_count: Top k pose estimates to consider (0 = all).
  :param gt_valid_mask: Mask of ground-truth poses which can be considered.
  :return: List of dictionaries, where each dictionary holds info for one pose
    estimate (the estimates are ordered as in errs) about the matching
    ground-truth pose:
    - 'est_id': ID of the pose estimate.
    - 'gt_id': ID of the matched ground-truth pose (-1 means there is no
        matching ground-truth pose).
    - 'score': Confidence score of the pose estimate.
    - 'error': Error of the pose estimate w.r.t. the matched ground-truth pose.
    - 'error_norm': Error normalized by the threshold value.
  """
  # Sort the estimated poses by decreasing confidence score.
  errs_sorted = sorted(errs, key=lambda e: e['score'], reverse=True)

  # Keep only the required number of poses with the highest confidence score.
  # 0 = all pose estimates are considered.
  if max_ests_count > 0:
    errs_sorted = errs_sorted[:max_ests_count]

  # Number of values defining the error (e.g. 1 for "ADD", 2 for "5deg 5cm").
  error_num_elems = len(list(error_ths))

  # Greedily match the estimated poses to the ground truth poses in the order of
  # decreasing score of the estimates.
  matches = []
  gt_matched = []
  for e in errs_sorted:

    best_gt_id = -1
    best_error = list(error_ths)
    for gt_id, error in e['errors'].items():

      # If the mask of valid GT poses is not provided, consider all valid.
      is_valid = not gt_valid_mask or gt_valid_mask[gt_id]

      # Only valid GT poses that have not been matched yet are considered.
      if is_valid and gt_id not in gt_matched:

        # The current pose estimate is considered the best so far if all error
        # elements are the lowest so far.
        if np.all([error[i] < best_error[i] for i in range(error_num_elems)]):
          best_gt_id = gt_id
          best_error = error

    if best_gt_id >= 0:

      # Mark the GT pose as matched.
      gt_matched.append(best_gt_id)

      # Error normalized by the threshold.
      best_errors_normed = [best_error[i] / float(error_ths[i])
                            for i in range(error_num_elems)]

      # Save info about the match.
      matches.append({
        'est_id': e['est_id'],
        'gt_id': best_gt_id,
        'score': e['score'],
        'error': best_error,
        'error_norm': best_errors_normed
      })

  return matches


def match_poses_scene(scene_id, scene_gt, scene_gt_valid, scene_errs,
                      correct_th, n_top):
  """Matches the estimated poses to the ground-truth poses in one scene.

  :param scene_id: Scene ID.
  :param scene_gt: Dictionary mapping image ID's to lists of dictionaries with:
    - 'obj_id': Object ID of the ground-truth pose.
  :param scene_gt_valid: Dictionary mapping image ID's to lists of boolean
    values indicating which ground-truth poses should be considered.
  :param scene_errs: List of dictionaries with:
    - 'im_id': Image ID.
    - 'obj_id': Object ID.
    - 'est_id': ID of the pose estimate.
    - 'score': Confidence score of the pose estimate.
    - 'errors': Dictionary mapping ground-truth ID's to errors of the pose
        estimate w.r.t. the ground-truth poses.
  :param error_obj_threshs: Dictionary mapping object ID's to values of the
    threshold of correctness.
  :param n_top: Top N pose estimates (with the highest score) to be evaluated
    for each object class in each image.
  :return:
  """
  # Organize the errors by image ID and object ID (for faster query).
  scene_errs_org = {}
  for e in scene_errs:
    scene_errs_org.setdefault(
      e['im_id'], {}).setdefault(e['obj_id'], []).append(e)

  # Matching of poses in individual images.
  scene_matches = []
  for im_id, im_gts in scene_gt.items():
    im_matches = []

    for gt_id, gt in enumerate(im_gts):
      im_matches.append({
        'scene_id': scene_id,
        'im_id': im_id,
        'obj_id': gt['obj_id'],
        'gt_id': gt_id,
        'est_id': -1,
        'score': -1,
        'error': -1,
        'error_norm': -1,
        'valid': scene_gt_valid[im_id][gt_id],
      })

    # Treat estimates of each object separately.
    im_obj_ids = set([gt['obj_id'] for gt in im_gts])
    for obj_id in im_obj_ids:
      if im_id in scene_errs_org.keys()\
            and obj_id in scene_errs_org[im_id].keys():

        # Greedily match the estimated poses to the ground truth poses.
        errs_im_obj = scene_errs_org[im_id][obj_id]
        ms = match_poses(
          errs_im_obj, correct_th, n_top, scene_gt_valid[im_id])

        # Update info about the matched GT poses.
        for m in ms:
          g = im_matches[m['gt_id']]
          g['est_id'] = m['est_id']
          g['score'] = m['score']
          g['error'] = m['error']
          g['error_norm'] = m['error_norm']

    scene_matches += im_matches

  return scene_matches
