# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Helper classes for computing performance metrics."""

import collections
from typing import Dict, List, Optional, Tuple

import lingvo.compat as tf
from lingvo.core import hyperparams
from lingvo.core import plot
from lingvo.core import py_utils
from lingvo.core import scorers
import numpy as np
try:
  import sklearn.metrics  # pylint: disable=g-import-not-at-top
  HAS_SKLEARN = True
except ImportError:
  HAS_SKLEARN = False
try:
  import scipy.stats  # pylint: disable=g-import-not-at-top
  HAS_SCIPY_STATS = True
except ImportError:
  HAS_SCIPY_STATS = False


def CreateScalarSummary(name, simple_value):
  return tf.Summary(
      value=[tf.Summary.Value(tag=name, simple_value=simple_value)])


class BaseMetric:
  """Base class for aggregating statistics to compute a performance metric."""

  def Update(self, *args, **kwargs):
    """Updates this metric (e.g. accumulates statistics) from the arguments."""
    pass

  @property
  def value(self):
    """Current value of this metric."""
    return None

  def Summary(self, name):
    """Converts the current state of this metric to a `tf.Summary`.

    Args:
      name: A string to use as the summary value tag.

    Returns:
      A `tf.Summary` proto.
    """
    return CreateScalarSummary(name, self.value)


class ConfigurableMetric(BaseMetric):
  """A Metric class with configurable params."""

  @classmethod
  def Params(cls):
    p = hyperparams.InstantiableParams(cls)
    return p

  def __init__(self, params):
    self.params = params


class AverageMetric(BaseMetric):
  """Class to compute a weighted (arithmetic) average value metric."""

  def __init__(self):
    self._total_value = 0.0
    self._total_weight = 0.0

  def Update(self, value, weight=1.0):
    if weight < 0.0:
      raise ValueError('weight must be non-negative.  Got: %f' % weight)
    self._total_value += value * weight
    self._total_weight += weight

  # We may want both a getter and a setter method for total_value and
  # total_weight, respectively.
  def GetTotalValue(self):
    return self._total_value

  def SetTotalValue(self, val):
    self._total_value = val

  total_value = property(GetTotalValue, SetTotalValue)

  def GetTotalWeight(self):
    return self._total_weight

  def SetTotalWeight(self, val):
    self._total_weight = val

  total_weight = property(GetTotalWeight, SetTotalWeight)

  @property
  def value(self):
    return (self._total_value /
            self._total_weight if self._total_weight > 0 else 0)


class UniqueAverageMetric(AverageMetric):
  """Computes average metric on keyed results to ensure unique counts.

  This implementation's Update field takes in a `key` argument that allows
  tracking of per-key values.  If a previously seen key is provided, this
  implementation checks that the previous value is the same as the current
  value for that key to help flag issues with non-determinism, and ensures
  that the final value reflects an average over unique keys only.
  """

  def __init__(self, mismatch_is_error=True):
    super().__init__()
    self._stored_values = {}
    self._mismatch_is_error = mismatch_is_error
    self._keys_with_different_values = set()

  def Update(self, key, value, weight=1.0):
    prev_value = self._stored_values.get(key, None)

    if prev_value is not None:
      # Check that value is equal to curr_value, and raise
      # an error if not equal.
      #
      # We may want to make the tolerances configurable via ConfigurableMetric.
      if not np.isclose(prev_value, value):
        self._keys_with_different_values.add(key)
      return

    self._stored_values[key] = value
    return super().Update(value, weight)

  def Summary(self, name):
    """Converts the current state of this metric to a `tf.Summary`.

    Args:
      name: A string to use as the summary value tag.

    Returns:
      A `tf.Summary` proto.
    """
    summary = tf.Summary(value=[
        tf.Summary.Value(tag=name, simple_value=self.value),
        tf.Summary.Value(
            tag=name + '/total_count', simple_value=len(self._stored_values)),
        tf.Summary.Value(
            tag=name + '/total_value', simple_value=self._total_value),
        tf.Summary.Value(
            tag=name + '/total_weight', simple_value=self._total_weight),
    ])
    return summary

  @property
  def value(self):
    if self._keys_with_different_values:
      msg = ('UniqueAverageMetric received duplicate keys with different '
             f'values: {self._keys_with_different_values}')
      if self._mismatch_is_error:
        # Log and raise an error, since current infrastructure can
        # swallow exceptions.
        tf.logging.error(msg)
        raise ValueError(msg)
      else:
        tf.logging.warning(msg)

    return super().value


class F1Metric(BaseMetric):
  """Class to compute F1 metrics."""

  def __init__(self):
    self._true_pos = 0.0
    self._false_pos = 0.0
    self._false_neg = 0.0

  def UpdateTruePositive(self, count=1.0):
    self._true_pos += count

  def UpdateFalsePositive(self, count=1.0):
    self._false_pos += count

  def UpdateFalseNegative(self, count=1.0):
    self._false_neg += count

  @property
  def value(self):
    if (self._true_pos + self._false_pos) > 0:
      precision = self._true_pos / (self._true_pos + self._false_pos)
    else:
      precision = 0.0
    if (self._true_pos + self._false_neg) > 0:
      recall = self._true_pos / (self._true_pos + self._false_neg)
    else:
      recall = 0.0
    if (precision + recall) > 0:
      return 2.0 * precision * recall / (precision + recall)
    else:
      return 0.0


class MCCMetric(F1Metric):
  """Class to compute Matthews correlation coefficient metric."""

  def __init__(self):
    super().__init__()
    self._true_neg = 0.0

  def UpdateTrueNegative(self, count=1.0):
    self._true_neg += count

  @property
  def value(self):
    tp = self._true_pos
    tn = self._true_neg
    fp = self._false_pos
    fn = self._false_neg

    denominator = (tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)
    if denominator > 0.0:
      return (tp * tn - fp * fn) / denominator**0.5
    else:
      return 0.0


class CorpusBleuMetric(BaseMetric):
  """Metric class to compute the corpus-level BLEU score."""

  def __init__(self, **kwargs):
    self._scorer = scorers.BleuScorer(**kwargs)

  def Update(self, ref_str, hyp_str):
    self._scorer.AddSentence(ref_str, hyp_str)

  @property
  def unsegmenter(self):
    return self._scorer.unsegmenter

  @property
  def value(self):
    return self._scorer.ComputeOverallScore()


class TpuEvalMetrics:
  """Manages computation of metrics during TPU execution.

  TPU execution runs a training loop on device. To get eval metrics out of this,
  metric values and weights must be carried through the loop. This requires
  passing initial values to the loop setup, updated the values during the loop,
  and doing a final aggregation after the loop. This class wraps the metrics
  dictionary so that the needed ops can be built at the right time as the
  training loop is built.

  Note that because the model is not constructed until the loop body function is
  called, the initial values must be known statically. This is done currently by
  hardcoding a limit on the number of metrics and casting each metric and value
  count to float32, regardless of the number of actual metrics the model
  produced.

  Note that this implementation computes the metrics over all replicas, for the
  last step of the loop only (could be changed to average over all loop steps
  instead).
  """

  def __init__(self, max_metrics: int = 256):
    self._metrics: Optional[Dict[str, Tuple[tf.Tensor, tf.Tensor]]] = None
    self._max_metrics: int = max_metrics

    # Loop-carried values alternate value and weight; all values are scalars.
    self._initial_values = (2 * self._max_metrics) * [
        tf.constant(0, tf.float32)
    ]

  @property
  def initial_values(self):
    return self._initial_values

  @property
  def metrics(self):
    return self._metrics

  def PackStepMetricsForAccumulation(self, metric_dict, step_args):
    """Sets the metrics to evaluate and the per-step output tensors.

    Args:
      metric_dict: dict of (name -> (tensor of values, tensor of weights))
      step_args: the loop-carried tensors which hold accumulated metric values.
        These share the same structure of alternating value and weight scalars
        as the initial values and the output of this function.

    Returns:
      The tensors to return from the training loop body.  For entries that are
      for metrics in self._metrics, returns the value computed within the loop
      (the step_args value passed in); for all others, the value will never be
      used at the end and so the step_args value is passed through (which has
      the effect of passing the initial values through every iteration of the
      loop).
    """
    num_metrics = len(metric_dict)
    assert num_metrics <= self._max_metrics, (
        'Increase max_metrics to >= %d' % num_metrics
    )
    # Persisting to reuse structure in Finalize
    self._metrics = metric_dict

    # self._metrics contains a map of (metric_value, metric_weight).
    # We convert it into [metric_value * metric_weight, metric_weight] to make
    # it easier to aggregate metric values across steps and TPU replicas.
    ret = []
    for _, (value, weight) in sorted(metric_dict.items()):
      assert value.shape.is_fully_defined(), '%s' % value
      assert weight.shape.is_fully_defined(), '%s' % weight
      weight = tf.cast(weight, tf.float32)
      value = tf.cast(value, tf.float32) * weight
      ret += [value, weight]
    # Each metric has two tensors: value and weight.
    assert len(ret) == 2 * num_metrics
    ret += list(step_args)[len(ret) :]
    return ret

  def FinalizeMetrics(
      self, loop_carried_metrics: List[tf.Tensor]
  ) -> List[tf.Tensor]:
    """Compute final average of the metrics, given loop_carried_metrics tensors.

    To be called outside the training loop body , but still in the scope of
    tpu.batch_parallel.

    Args:
      loop_carried_metrics: Result of the training loop.

    Returns:
      The tensors of the final avg values and total weights.
    """
    # Each metric has two tensors in the loop carrying result.
    metrics = [
        tf.tpu.cross_replica_sum(x)
        for x in loop_carried_metrics[: (2 * len(self.metrics))]
    ]

    return tf.nest.flatten(
        [
            py_utils.WeightedAvg(tf.math.divide_no_nan(value, weight), weight)
            for value, weight in py_utils.Chunked(metrics)
        ]
    )

  def PackMetricsValues(self, values) -> None:
    """Packs numpy values into self._metrics, a dict of metrics."""
    for k, v in zip(sorted(self._metrics.keys()), py_utils.Chunked(values)):
      self._metrics[k] = v

  @classmethod
  def ToAverageMetric(cls, value, weight=1.0) -> AverageMetric:
    avg_metric = AverageMetric()
    avg_metric.total_weight = weight
    avg_metric.total_value = weight * value
    return avg_metric

  def ToAverageMetrics(self) -> Dict[str, AverageMetric]:
    """Wrap the final metric values into AverageMetric objects.

    Returns:
      A dict that maps metric names to AverageMetric objects.
    """
    return {
        name: self.ToAverageMetric(value, weight)
        for name, (value, weight) in self._metrics.items()
    }


class TpuVariableMetrics:
  """A rewrite of TpuEvalMetrics that uses tf.Variables."""

  def __init__(self, max_metrics: int, strategy: tf.distribute.TPUStrategy):
    # TODO(b/263176991): Remove in a later refactor because we depend now on the
    #  user specifying the correct number of metrics in
    #  program.params.max_metrics.
    self._max_metrics = max_metrics
    self._metrics: Optional[Dict[str, Tuple[tf.Tensor, tf.Tensor]]] = None
    self._strategy = strategy
    self._vars = [self._InitializeVar(i) for i in range(2 * self._max_metrics)]

  @staticmethod
  def _InitializeVar(i):
    return tf.Variable(
        0.0,
        name=f'values_by_weights_{i//2}' if i % 2 == 0 else f'weights_{i//2}',
        trainable=False,
        dtype=tf.float32,
        synchronization=tf.VariableSynchronization.ON_READ,
        aggregation=tf.VariableAggregation.SUM,
    )

  @property
  def variables(self) -> List[tf.Variable]:
    return self._vars

  def ResetState(self) -> None:
    """Resets the tracked metric variables to 0; required on every k steps."""
    for v in self._vars:
      v.assign(0.0)

  def AccumulateStepMetrics(
      self, metric_dict: Dict[str, Tuple[tf.Tensor, tf.Tensor]]
  ) -> None:
    """Accumulates values in self._vars. Should be called via strategy.run."""
    num_metrics = len(metric_dict)
    assert num_metrics <= self._max_metrics, (
        'Increase max_metrics to >= %d' % num_metrics
    )
    for (_, (value, weight)), (vw_acc, w_acc) in zip(
        sorted(metric_dict.items()), py_utils.Chunked(self._vars)
    ):
      vw_acc.assign_add(
          tf.cast(weight, tf.float32) * tf.cast(value, tf.float32)
      )
      w_acc.assign_add(tf.cast(weight, tf.float32))

  def FinalizeMetricsWithStructure(
      self, structure: Dict[str, Tuple[tf.Tensor, tf.Tensor]]
  ) -> List[tf.Tensor]:
    """Compute final average of the metrics using the accumulator variables.

    To be called outside a strategy.run context, so that the tf.Variables'
    tf.VariableAggregation logic is triggered on this read (i.e. a tpu
    cross_replica_sum).

    Args:
      structure: The dictionary structure we wish to tf.nest.pack_sequence_as.

    Returns:
      The tensors of the final avg values and total weights.
    """

    return tf.nest.pack_sequence_as(
        structure,
        tf.nest.flatten(
            [
                [tf.math.divide_no_nan(vw, w), w]
                for vw, w in py_utils.Chunked(self._vars)
            ]
        ),
    )


class AUCMetric(BaseMetric):
  """Class to compute the AUC score for binary classification."""

  def __init__(self, mode='roc', samples=-1):
    """Constructor of the class.

    Args:
      mode: Possible values: 'roc' or 'pr'.
      samples: The number of sample points to compute the AUC. If -1, include
        all points seen thus far.

    Raises:
      ImportError: If user has not installed sklearn, raise an ImportError.
    """
    if not HAS_SKLEARN:
      raise ImportError('AUCMetric depends on sklearn.')
    self._mode = mode
    self._samples = samples
    self._label = []
    self._prob = []
    self._weight = []
    if self._mode == 'roc':
      self._curve_fn = sklearn.metrics.roc_curve
      self._score_fn = sklearn.metrics.roc_auc_score
      self._plot_labels = ['False Positive Rate', 'True Positive Rate']
    elif self._mode == 'pr':
      self._curve_fn = sklearn.metrics.precision_recall_curve
      self._score_fn = sklearn.metrics.average_precision_score
      self._plot_labels = ['Recall', 'Precision']
    else:
      raise ValueError('mode in AUCMetric must be one of "roc" or "pr".')

  def Update(self, label, prob, weight=None):
    """Updates the metrics.

    Args:
      label: An array to specify the groundtruth binary labels. Values must be
        either 0 or 1.
      prob: An array to specify the prediction probabilities. Values must be
        within [0, 1.0].
      weight: An array to specify the sample weight for the auc computation.
    """
    self._label += label
    self._prob += prob
    if weight:
      self._weight += weight
    else:
      self._weight += [1 for _ in range(len(label))]

    if self._samples > 0:
      self._label = self._label[-self._samples:]
      self._prob = self._prob[-self._samples:]
      self._weight = self._weight[-self._samples:]

  @property
  def value(self):
    try:
      auc_val = self._score_fn(
          self._label, self._prob, sample_weight=self._weight)
      # For precision/recall, _score_fn returns nan if only one type of label is
      # present rather than throwing a ValueError.
      if np.isnan(auc_val):
        return 0.0
      else:
        return auc_val
    except ValueError as exception:
      # In case self._label still has just 1 type of label, e.g. all(labels==0).
      if 'Only one class present in y_true.' in str(exception):
        return 0.0
      else:
        raise

  def Summary(self, name):

    def _Setter(fig, axes):
      # 20 ticks betweein 0 and 1.
      ticks = np.arange(0, 1.05, 0.05)
      axes.grid(b=True)
      axes.set_xlabel(self._plot_labels[0])
      axes.set_xticks(ticks)
      axes.set_ylabel(self._plot_labels[1])
      axes.set_yticks(ticks)
      fig.tight_layout()

    xs, ys, _ = self._curve_fn(
        self._label, self._prob, sample_weight=self._weight)
    if self._mode == 'pr':
      # Swap because sklearn returns <'precision', 'recall'>.
      xs, ys = ys, xs
    ret = plot.Curve(name=name, figsize=(12, 12), xs=xs, ys=ys, setter=_Setter)
    ret.value.add(tag=name, simple_value=self.value)
    return ret

  def _PrecisionAtRecall(self, recall):
    # calculate precision@recall metric
    assert self._mode == 'pr'
    p, r, t = self._curve_fn(
        self._label, self._prob, sample_weight=self._weight)
    last_p = 0.0
    for pp, rr, _ in zip(p, r, t):
      if rr >= recall:
        last_p = pp
    return last_p

  def _RecallAtPrecision(self, precision):
    # calculate recall@precision metric
    assert self._mode == 'pr'
    p, r, t = self._curve_fn(
        self._label, self._prob, sample_weight=self._weight)
    for pp, rr, _ in zip(p, r, t):
      if pp >= precision:
        return rr
    return 0.0


class PrecisionAtRecall(AUCMetric):

  def __init__(self, recall_threshold, samples=-1):
    super().__init__(mode='pr', samples=samples)
    self._recall_threshold = recall_threshold

  @property
  def value(self):
    return self._PrecisionAtRecall(self._recall_threshold)


class RecallAtPrecision(AUCMetric):

  def __init__(self, precision_threshold, samples=-1):
    super().__init__(mode='pr', samples=samples)
    self._precision_threshold = precision_threshold

  @property
  def value(self):
    return self._RecallAtPrecision(self._precision_threshold)


class MultiClassAUCMetric(BaseMetric):
  """Class to compute mAP or mAUC for multiclass/multilabel classification.

  This metric is equivalent to mAP (mean average precision) or mAUC (multiclass
  AUC-ROC) for non-binary classification tasks. Another perspective of this
  measure defines it as the per-class average of the auc roc or auc pr curves.
  """

  def __init__(self, num_classes, mode='roc', samples=-1):
    """Construct a MultiClassAUCMetric instance.

    Args:
      num_classes: The number of classes.
      mode: Possible values: 'roc' or 'pr'.
      samples: The number of sample points to compute the AUC. If -1, include
        all points seen thus far.
    """
    self._num_classes = num_classes
    self._class_auc_metrics = []
    for _ in range(self._num_classes):
      self._class_auc_metrics.append(AUCMetric(mode, samples))

    # Track the classes that are present in the evaluation set.
    self._labels_seen_dict = {i: 0 for i in range(self._num_classes)}

  def Update(self, labels, probs, weights=None):
    """Updates the MultiClassAUCMetric.

    Args:
      labels: An list of arrays specifying the groundtruth binary labels for
        each class. Values must be either 0 or 1.
      probs: A list of arrays specifying the prediction probabilities for each
        class. Values must be within [0, 1.0].
      weights: An array to specify the sample weight for each example in the auc
        computation.
    """
    for i, auc_metric in enumerate(self._class_auc_metrics):
      sliced_labels = labels[i]
      sliced_probs = probs[i]
      auc_metric.Update(sliced_labels, sliced_probs, weights)

      for label in sliced_labels:
        self._labels_seen_dict[i] = max(label, self._labels_seen_dict[i])

  @property
  def value(self):
    auc_values = []
    for index, auc_metric in enumerate(self._class_auc_metrics):
      if self._labels_seen_dict[index]:
        auc_values.append(auc_metric.value)
    num_classes_present_in_eval = sum(self._labels_seen_dict.values())
    return np.sum(auc_values) / num_classes_present_in_eval


class CorrelationMetric(BaseMetric):
  """Class to compute correlation."""

  def __init__(self, mode='pearson', samples=-1):
    """Constructor of the class.

    Args:
      mode: Possible values: 'pearson', 'spearman', 'kendalltau'.
      samples: The number of sample points to compute the correlation. If -1,
        include all points seen thus far.

    Raises:
      ImportError: If user has not installed scipy.stats, raise an ImportError.
    """
    if not HAS_SCIPY_STATS:
      raise ImportError('CorrelationMetric depends on scipy.stats.')

    assert mode in ['pearson', 'spearman', 'kendalltau']
    self._mode = mode
    self._samples = samples
    self._target = []
    self._pred = []

  def Update(self, target, pred):
    """Updates the metrics.

    Args:
      target: An array to specify the groundtruth float target.
      pred: An array to specify the prediction.
    """
    self._target += target
    self._pred += pred

    if self._samples > 0:
      self._target = self._target[-self._samples:]
      self._pred = self._pred[-self._samples:]

  @property
  def value(self):
    # only use the correlation, p-value is ignored.
    if self._mode == 'pearson':
      return scipy.stats.pearsonr(self._target, self._pred)[0]
    elif self._mode == 'spearman':
      return scipy.stats.spearmanr(self._target, self._pred)[0]
    else:
      return scipy.stats.kendalltau(self._target, self._pred)[0]


class AverageKeyedCorrelationMetric(BaseMetric):
  """Class to compute correlation per key, then report average across all keys.
  """

  def __init__(self, mode='pearson', bypass_nan=True):
    """Constructor of the class.

    Args:
      mode: Possible values: 'pearson', 'spearman', 'kendalltau'.
      bypass_nan: for keys that has only 1 example, the metric will be NaN, turn
        on this flag to skip those keys. Depending on the scipy library, it may
        raise ValueException instead of produce NaN. In this case, the Exception
        will be suppressed and the key is skipped during calculation.

    Raises:
      ImportError: If user has not installed scipy.stats, raise an ImportError.
    """
    if not HAS_SCIPY_STATS:
      raise ImportError('CorrelationMetric depends on scipy.stats.')

    assert mode in ['pearson', 'spearman', 'kendalltau']
    self._mode = mode
    self._bypass_nan = bypass_nan
    self._target = collections.defaultdict(list)
    self._pred = collections.defaultdict(list)

  def Update(self, key, target, pred):
    """Updates the metrics.

    Args:
      key: The key this (target, pred) pair belongs to.
      target: An array to specify the groundtruth float target.
      pred: An array to specify the prediction.
    """
    self._target[key] += target
    self._pred[key] += pred

  @property
  def value(self):
    # only use the correlation, p-value is ignored.
    if self._mode == 'pearson':
      corr_f = scipy.stats.pearsonr
    elif self._mode == 'spearman':
      corr_f = scipy.stats.spearmanr
    else:
      corr_f = scipy.stats.kendalltau

    results = []
    for k in self._target:
      target = self._target[k]
      pred = self._pred[k]
      try:
        raw_corr = corr_f(target, pred)[0]
        if not self._bypass_nan or not np.isnan(raw_corr):
          results.append(raw_corr)
      except ValueError:
        continue

    if not self._bypass_nan or results:
      return np.mean(results)
    else:
      return 0.0


class SamplingMetric(ConfigurableMetric):
  """Sampling metric base class.

  Subclasses must implement _CreateSummary(); sampling will be handled
  by this base class.
  """

  @classmethod
  def Params(cls):
    p = super().Params()
    p.Define('num_samples', 8, 'The number of samples to store uniformly.')
    return p

  def __init__(self, params):
    super().__init__(params)
    p = self.params
    self._sampler = py_utils.UniformSampler(num_samples=p.num_samples)
    self._summary = None

  @property
  def samples(self):
    """Returns an iterable of sampled decoded outputs to compute Summaries."""
    return self._sampler.samples

  def Update(self, decoded_outputs):
    """Samples the input decoded_outputs NestedMap.

    Args:
      decoded_outputs: A `.NestedMap`.
    """
    self._sampler.Add(decoded_outputs)
    # Invalidate cache.
    self._summary = None

  def Summary(self, name):
    if self._summary is None:
      self._summary = self._CreateSummary(name)
      self._sampler = py_utils.UniformSampler(
          num_samples=self.params.num_samples)
    return self._summary

  def _CreateSummary(self, name):
    """Returns a tf.Summary for this metric."""
    raise NotImplementedError()


class GroupPairAUCMetric(AUCMetric):
  """Compute the AUC score for all pairs extracted from each group of items.

  For each group of items, the metric extracts all pairs with different
  target values. For each pair (i, j), the metric computes the binary
  classification AUC where the `label = 1 if target[i] > target[j] else 0` and
  `prob = sigmoid(logits[i] - logits[j])`.

  To prevent generating pairs across groups, an additional arg `group_ids` is
  required, which is a list of ints that specifies the group_id of each item.

  In addition, in order to achieve streaming computation, items from the same
  group need to form continuous chunks,
  e.g. group_ids = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2].

  In the case of [0, 0, 1, 1, 1, 0, 0, 2, 2, 2, 2], the second chunk of 0s will
  be treated as a separate 3rd group rather than part of the 1st group.
  """

  def UpdateRaw(self, group_ids, target, logits, weight=None):
    """Updates the metrics.

    Args:
      group_ids: An array to specify the group identity.
      target: An array to specify the groundtruth float values.
      logits: An array to specify the raw prediction logits.
      weight: An array to specify the sample weight for the auc computation.
    """

    assert self._samples <= 0

    sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))

    def _ProcessChunk(s, e):
      for i in range(s, e):
        for j in range(i + 1, e):
          if target[i] != target[j]:
            pair_label = 1 if target[i] > target[j] else 0
            pair_prob = sigmoid(logits[i] - logits[j])
            self._label.append(pair_label)
            self._prob.append(pair_prob)
            if weight:
              self._weight.append(min(1.0, weight[i] + weight[j]))
            else:
              self._weight.append(1.0)

    s, e = 0, 1
    while e <= len(target):
      # Find the end of a chunk
      if e == len(target) or group_ids[e] != group_ids[s]:
        # Process the current chunk [s:e]
        _ProcessChunk(s, e)

        # Start a new chunk by setting `s` to `e`
        s = e
      # Increment `e` by 1.
      e += 1
