import torch


def _compute_distances(gradients, f, mode, **kwargs):
  """ Aksel rule.
  Args:
    gradients Non-empty list of gradients to aggregate
    f         Number of Byzantine gradients to tolerate
    mode      Operation mode, one of: 'mid', 'n-f'
    ...       Ignored keyword-arguments
  Returns:
    List of (gradient, distance to the median) sorted by increasing distance,
    Number of gradients to aggregate
  """
  n = len(gradients)
  # Prepare gradients
  g = torch.stack(gradients)
  # Compute median
  m = g.median(dim=0)[0]
  # Measure squared distances to median
  d = list((i, (x - m).pow_(2).sum().item()) for i, x in enumerate(gradients))
  # Average closest to median according to mode
  if mode == "mid":
    c = (n + 1) // 2
  elif mode == "n-f":
    c = n - f
  else:
    raise NotImplementedError
  d.sort(key=lambda x: x[1])
  return d, c