"""Evaluation metrics."""

import numpy as np
from sklearn.metrics import matthews_corrcoef  # noqa
from sklearn.metrics import recall_score  # noqa
from sklearn.metrics import cohen_kappa_score
from sklearn.metrics import r2_score  # noqa
from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import precision_score  # noqa
from sklearn.metrics import precision_recall_curve
from sklearn.metrics import auc
from sklearn.metrics import jaccard_score
from sklearn.metrics import f1_score
from sklearn.metrics import roc_auc_score  # noqa
from sklearn.metrics import accuracy_score  # noqa
from sklearn.metrics import balanced_accuracy_score  # noqa
from scipy.stats import pearsonr

# kappa_score is an alias for `sklearn.metrics.cohen_kappa_score`
kappa_score = cohen_kappa_score


def pearson_r2_score(y: np.ndarray, y_pred: np.ndarray) -> float:
  """Computes Pearson R^2 (square of Pearson correlation).

  Parameters
  ----------
  y: np.ndarray
    ground truth array
  y_pred: np.ndarray
    predicted array

  Returns
  -------
  float
    The Pearson-R^2 score.
  """
  return pearsonr(y, y_pred)[0]**2


def jaccard_index(y: np.ndarray, y_pred: np.ndarray) -> float:
  """Computes Jaccard Index which is the Intersection Over Union metric
  which is commonly used in image segmentation tasks.

  DEPRECATED: WILL BE REMOVED IN A FUTURE VERSION OF DEEEPCHEM. USE `jaccard_score` instead.

  Parameters
  ----------
  y: np.ndarray
    ground truth array
  y_pred: np.ndarray
    predicted array

  Returns
  -------
  score: float
    The jaccard index. A number between 0 and 1.
  """
  return jaccard_score(y, y_pred)


def pixel_error(y: np.ndarray, y_pred: np.ndarray) -> float:
  """An error metric in case y, y_pred are images.

  Defined as 1 - the maximal F-score of pixel similarity, or squared
  Euclidean distance between the original and the result labels.

  Parameters
  ----------
  y: np.ndarray
    ground truth array
  y_pred: np.ndarray
    predicted array

  Returns
  -------
  score: float
    The pixel-error. A number between 0 and 1.
  """
  return 1 - f1_score(y, y_pred)


def prc_auc_score(y: np.ndarray, y_pred: np.ndarray) -> float:
  """Compute area under precision-recall curve

  Parameters
  ----------
  y: np.ndarray
    A numpy array of shape `(N, n_classes)` or `(N,)` with true labels
  y_pred: np.ndarray
    Of shape `(N, n_classes)` with class probabilities.

  Returns
  -------
  float
    The area under the precision-recall curve. A number between 0 and 1.
  """
  precision, recall, _ = precision_recall_curve(y[:, 1], y_pred[:, 1])
  return auc(recall, precision)


def rms_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
  """Computes RMS error."""
  return np.sqrt(mean_squared_error(y_true, y_pred))


def mae_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
  """Computes MAE."""
  return mean_absolute_error(y_true, y_pred)


def bedroc_score(y_true: np.ndarray, y_pred: np.ndarray, alpha: float = 20.0):
  """Compute BEDROC metric.

  BEDROC metric implemented according to Truchon and Bayley that modifies
  the ROC score by allowing for a factor of early recognition.
  Please confirm details from [1]_.

  Parameters
  ----------
  y_true: np.ndarray
    Binary class labels. 1 for positive class, 0 otherwise
  y_pred: np.ndarray
    Predicted labels
  alpha: float, default 20.0
    Early recognition parameter

  Returns
  -------
  float
    Value in [0, 1] that indicates the degree of early recognition

  Notes
  -----
  This function requires RDKit to be installed.

  References
  ----------
  .. [1] Truchon et al. "Evaluating virtual screening methods: good and bad metrics
     for the “early recognition” problem." Journal of chemical information and modeling
     47.2 (2007): 488-508.
  """
  try:
    from rdkit.ML.Scoring.Scoring import CalcBEDROC
  except ModuleNotFoundError:
    raise ValueError("This function requires RDKit to be installed.")

  # validation
  assert len(y_true) == len(y_pred), 'Number of examples do not match'
  assert np.array_equal(
      np.unique(y_true).astype(int),
      [0, 1]), ('Class labels must be binary: %s' % np.unique(y_true))

  yt = np.asarray(y_true)
  yp = np.asarray(y_pred)

  yt = yt.flatten()
  yp = yp[:, 1].flatten()  # Index 1 because one_hot predictions

  scores = list(zip(yt, yp))
  scores = sorted(scores, key=lambda pair: pair[1], reverse=True)

  return CalcBEDROC(scores, 0, alpha)

[((60, 9, 60, 33), 'sklearn.metrics.jaccard_score', 'jaccard_score', ({(60, 23, 60, 24): 'y', (60, 26, 60, 32): 'y_pred'}, {}), '(y, y_pred)', False, 'from sklearn.metrics import jaccard_score\n'), ((99, 25, 99, 70), 'sklearn.metrics.precision_recall_curve', 'precision_recall_curve', ({(99, 48, 99, 55): 'y[:, (1)]', (99, 57, 99, 69): 'y_pred[:, (1)]'}, {}), '(y[:, (1)], y_pred[:, (1)])', False, 'from sklearn.metrics import precision_recall_curve\n'), ((100, 9, 100, 31), 'sklearn.metrics.auc', 'auc', ({(100, 13, 100, 19): 'recall', (100, 21, 100, 30): 'precision'}, {}), '(recall, precision)', False, 'from sklearn.metrics import auc\n'), ((110, 9, 110, 44), 'sklearn.metrics.mean_absolute_error', 'mean_absolute_error', ({(110, 29, 110, 35): 'y_true', (110, 37, 110, 43): 'y_pred'}, {}), '(y_true, y_pred)', False, 'from sklearn.metrics import mean_absolute_error\n'), ((155, 7, 155, 25), 'numpy.asarray', 'np.asarray', ({(155, 18, 155, 24): 'y_true'}, {}), '(y_true)', True, 'import numpy as np\n'), ((156, 7, 156, 25), 'numpy.asarray', 'np.asarray', ({(156, 18, 156, 24): 'y_pred'}, {}), '(y_pred)', True, 'import numpy as np\n'), ((164, 9, 164, 37), 'rdkit.ML.Scoring.Scoring.CalcBEDROC', 'CalcBEDROC', ({(164, 20, 164, 26): 'scores', (164, 28, 164, 29): '(0)', (164, 31, 164, 36): 'alpha'}, {}), '(scores, 0, alpha)', False, 'from rdkit.ML.Scoring.Scoring import CalcBEDROC\n'), ((81, 13, 81, 32), 'sklearn.metrics.f1_score', 'f1_score', ({(81, 22, 81, 23): 'y', (81, 25, 81, 31): 'y_pred'}, {}), '(y, y_pred)', False, 'from sklearn.metrics import f1_score\n'), ((105, 17, 105, 51), 'sklearn.metrics.mean_squared_error', 'mean_squared_error', ({(105, 36, 105, 42): 'y_true', (105, 44, 105, 50): 'y_pred'}, {}), '(y_true, y_pred)', False, 'from sklearn.metrics import mean_squared_error\n'), ((153, 52, 153, 69), 'numpy.unique', 'np.unique', ({(153, 62, 153, 68): 'y_true'}, {}), '(y_true)', True, 'import numpy as np\n'), ((39, 9, 39, 28), 'scipy.stats.pearsonr', 'pearsonr', ({(39, 18, 39, 19): 'y', (39, 21, 39, 27): 'y_pred'}, {}), '(y, y_pred)', False, 'from scipy.stats import pearsonr\n'), ((152, 6, 152, 23), 'numpy.unique', 'np.unique', ({(152, 16, 152, 22): 'y_true'}, {}), '(y_true)', True, 'import numpy as np\n')]