# Copyright 2020 Valentin Gabeur
# Copyright 2020 Samuel Albanie, Yang Liu and Arsha Nagrani
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module for computing performance metrics.

Code based on the implementation of "Collaborative Experts":
https://github.com/albanie/collaborative-experts
"""

import ipdb
import numpy as np
import scipy.stats


def t2v_metrics(sims, query_masks=None):
  """Compute retrieval metrics from a similiarity matrix.

  Args:
    sims: N x M matrix of similarities between embeddings, where x_{i,j} =
      <text_embd[i], vid_embed[j]>
    query_masks: mask any missing queries from the dataset (two videos in MSRVTT
      only have 19, rather than 20 captions)

  Returns:
    (dict[str:float]): retrieval metrics
  """
  assert sims.ndim == 2, "expected a matrix"
  nq, nv = sims.shape
  dists = -sims
  sorted_dists = np.sort(dists, axis=1)

  # if False:
  #   import sys
  #   import matplotlib
  #   from pathlib import Path
  #   matplotlib.use("Agg")
  #   import matplotlib.pyplot as plt
  #   sys.path.insert(0, str(Path.home() / "coding/src/zsvision/python"))
  #   from zsvision.zs_iterm import zs_dispFig  # NOQA
  #   plt.matshow(dists)
  #   zs_dispFig()
  #   # import ipdb; ipdb.set_trace()

  # The indices are computed such that they slice out the ground truth distances
  # from the psuedo-rectangular dist matrix
  qu = nq // nv  # Nb of queries per video
  gt_idx = [[
      np.ravel_multi_index([ii, jj], (nq, nv))
      for ii in range(jj * qu, (jj + 1) * qu)
  ]
            for jj in range(nv)]
  gt_idx = np.array(gt_idx)
  gt_dists = dists.reshape(-1)[gt_idx.reshape(-1)]
  gt_dists = gt_dists[:, np.newaxis]
  rows, cols = np.where(
      (sorted_dists - gt_dists) == 0)  # find column position of GT

  # --------------------------------
  # NOTE: Breaking ties
  # --------------------------------
  # We sometimes need to break ties (in general, these should occur extremely
  # rarely, but there are pathological cases when they can distort the scores,
  # such as when the similarity matrix is all zeros). Previous implementations
  # (e.g. the t2i evaluation function used here:
  # https://github.com/niluthpol/multimodal_vtt/blob/master/evaluation.py
  # and here:
  # https://github.com/linxd5/VSE_Pytorch/blob/master/evaluation.py#L87)
  # generally break ties "optimistically". However, if the similarity matrix is
  # constant this can evaluate to a perfect ranking. A principled option is to
  # average over all possible partial orderings implied by the ties. See this
  # paper for a discussion:
  #    McSherry, Frank, and Marc Najork,
  #    "Computing information retrieval performance measures efficiently in the
  #    presence of tied scores."
  #    European conference on information retrieval. Springer, Berlin,
  #    Heidelberg, 2008.
  # http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.145.8892&rep=rep1&type=pdf

  # break_ties = "optimistically"
  break_ties = "averaging"

  if rows.size > nq:
    assert np.unique(rows).size == nq, "issue in metric evaluation"
    if break_ties == "optimistically":
      _, idx = np.unique(rows, return_index=True)
      cols = cols[idx]
    elif break_ties == "averaging":
      # fast implementation, based on this code:
      # https://stackoverflow.com/a/49239335
      locs = np.argwhere((sorted_dists - gt_dists) == 0)

      # Find the split indices
      steps = np.diff(locs[:, 0])
      splits = np.nonzero(steps)[0] + 1
      splits = np.insert(splits, 0, 0)

      # Compute the result columns
      summed_cols = np.add.reduceat(locs[:, 1], splits)
      counts = np.diff(np.append(splits, locs.shape[0]))
      avg_cols = summed_cols / counts
      # if False:
      #   print("Running slower code to verify rank averaging across ties")
      #   # slow, but more interpretable version, used for testing
      #   avg_cols_slow = [
      #       np.mean(cols[rows == idx]) for idx in range(nq)
      #   ]
      #   assert np.array_equal(avg_cols,
      #                         avg_cols_slow), "slow vs fast difference"
      #   print("passed num check")
      cols = avg_cols

  msg = "expected ranks to match queries ({} vs {}) "
  if cols.size != nq:
    ipdb.set_trace()
  assert cols.size == nq, msg

  # if False:
  #   # overload mask to check that we can recover the scores for single-query
  #   # retrieval
  #   print("DEBUGGING MODE")
  #   query_masks = np.zeros_like(query_masks)
  #   query_masks[:, 0] = 1  # recover single query score

  if query_masks is not None:
    # remove invalid queries
    assert query_masks.size == nq, "invalid query mask shape"
    cols = cols[query_masks.reshape(-1).astype(np.bool)]
    assert cols.size == query_masks.sum(), "masking was not applied correctly"
    # update number of queries to account for those that were missing
    nq = query_masks.sum()

  # if False:
  #   # sanity check against old logic for square matrices
  #   gt_dists_old = np.diag(dists)
  #   gt_dists_old = gt_dists_old[:, np.newaxis]
  #   _, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
  #   assert np.array_equal(cols_old, cols), "new metric doesn't match"

  return cols2metrics(cols, nq)


def v2t_metrics(sims, query_masks=None):
  """Compute retrieval metrics from a similiarity matrix.

  Args:
    sims: N x M matrix of similarities between embeddings, where
      x_{i,j} = <text_embd[i], vid_embed[j]>
    query_masks: mask any missing captions from the dataset

  Returns:
    (dict[str:float]): retrieval metrics

  NOTES: We find the closest "GT caption" in the style of VSE, which
  corresponds
  to finding the rank of the closest relevant caption in embedding space:
  github.com/ryankiros/visual-semantic-embedding/blob/master/evaluation.py#L52-L56
  """
  # switch axes of text and video
  sims = sims.T

  # if False:
  #   # experiment with toy example
  #   sims = np.ones((3, 3))
  #   sims[0, 0] = 2
  #   sims[1, 1:2] = 2
  #   sims[2, :] = 2
  #   query_masks = None

  assert sims.ndim == 2, "expected a matrix"
  num_queries, num_caps = sims.shape
  dists = -sims
  caps_per_video = num_caps // num_queries
  break_ties = "averaging"

  missing_val = 1E8
  query_ranks = []
  for ii in range(num_queries):
    row_dists = dists[ii, :]
    if query_masks is not None:
      # Set missing queries to have a distance of infinity.  A missing query
      # refers to a query position `n` for a video that had less than `n`
      # captions (for example, a few MSRVTT videos only have 19 queries)
      row_dists[np.logical_not(query_masks.reshape(-1))] = missing_val

    # NOTE: Using distance subtraction to perform the ranking is easier to make
    # deterministic than using argsort, which suffers from the issue of defining
    # "stability" for equal distances.  Example of distance subtraction code:
    # github.com/antoine77340/Mixture-of-Embedding-Experts/blob/master/train.py
    sorted_dists = np.sort(row_dists)

    min_rank = np.inf
    for jj in range(ii * caps_per_video, (ii + 1) * caps_per_video):
      if row_dists[jj] == missing_val:
        # skip rankings of missing captions
        continue
      ranks = np.where((sorted_dists - row_dists[jj]) == 0)[0]
      if break_ties == "optimistically":
        rank = ranks[0]
      elif break_ties == "averaging":
        # NOTE: If there is more than one caption per video, its possible for
        # the method to do "worse than chance" in the degenerate case when all
        # similarities are tied.
        rank = ranks.mean()
      if rank < min_rank:
        min_rank = rank
    query_ranks.append(min_rank)
  query_ranks = np.array(query_ranks)

  #   # sanity check against old version of code
  #   if False:
  #     sorted_dists = np.sort(dists, axis=1)
  #     gt_dists_old = np.diag(dists)
  #     gt_dists_old = gt_dists_old[:, np.newaxis]
  #     rows_old, cols_old = np.where((sorted_dists - gt_dists_old) == 0)
  #     if rows_old.size > num_queries:
  #       _, idx = np.unique(rows_old, return_index=True)
  #       cols_old = cols_old[idx]
  #     num_diffs = (1 - (cols_old == query_ranks)).sum()
  #     msg = f"new metric doesn't match in {num_diffs} places"
  #     assert np.array_equal(cols_old, query_ranks), msg

  #     # visualise the distance matrix
  #     import sys
  #     import matplotlib
  #     matplotlib.use("Agg")
  #     import matplotlib.pyplot as plt
  #     sys.path.insert(0, str(Path.home() / "coding/src/zsvision/python"))
  #     from zsvision.zs_iterm import zs_dispFig  # NOQA
  #     plt.matshow(dists)
  #     zs_dispFig()

  return cols2metrics(query_ranks, num_queries)


def cols2metrics(cols, num_queries):
  """Compute the metrics."""
  metrics = {}
  metrics["R1"] = 100 * float(np.sum(cols == 0)) / num_queries
  metrics["R5"] = 100 * float(np.sum(cols < 5)) / num_queries
  metrics["R10"] = 100 * float(np.sum(cols < 10)) / num_queries
  metrics["R50"] = 100 * float(np.sum(cols < 50)) / num_queries
  metrics["MedR"] = np.median(cols) + 1
  metrics["MeanR"] = np.mean(cols) + 1
  stats = [metrics[x] for x in ("R1", "R5", "R10")]
  metrics["geometric_mean_R1-R5-R10"] = scipy.stats.mstats.gmean(stats)
  metrics["cols"] = [int(i) for i in list(cols)]
  return metrics
