# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/10_noise.ipynb (unless otherwise specified).

__all__ = ['suspect', 'suspect', 'suspect']

# Internal Cell
import numpy as np
import pandas as pd
from plum import dispatch
from loguru import logger

# Internal Cell
def __rank_suspicious(labels: np.ndarray, unique_labels: np.ndarray=None, *, probas: np.ndarray) -> np.ndarray:
    """ The internal method to rank the suspicious labels given negative probas from a binary classifier.

    Params:
        labels: Numpy Array of shape n x 1, label 0 for negative, 1 for positive cases.
        probas: Numpy Array of shape n x 1, probabilites for positive cases

    Return:
        ranks: Numpy Array of shape n x 1, the magnitude of suspicious labels, values between (-1, 1), the larger the more suspicious.
    """
    logger.debug(f"Numpy arrays labels and probas must have same length {len(labels)} vs {len(probas)}")
    if unique_labels is None:
        unique_labels = np.unique(labels)
    logger.debug(f"Unique labels {unique_labels}")
    if len(labels) != len(probas) or len(probas.shape) != len(unique_labels):
        logger.debug("Trying to reshape probas")
        probas = probas.reshape((len(labels),-1))
    logger.debug(f"Shape of labels and probas: {labels.shape} vs {probas.shape}")
    assert len(labels) == len(probas)
    assert probas.shape[1] == 1 # positive probabilieis from a binary classifier
    assert set(unique_labels).issubset([0, 1]) # binary classification
    assert (probas >= 0).all() and (probas<=1).all() # probas should be between 0 and 1

    # TBD: temporary implementation
    ranks = np.abs(labels - probas[:, 0])
    logger.debug(f"ranks {ranks}")
    assert len(labels) == len(ranks)

    return ranks

# Cell
@dispatch
def suspect(labels: np.ndarray, *, probas: np.ndarray) -> np.ndarray:
    """ Rank the  suspicious using Numpy API
    Params:
        labels: Numpy Array of shape n x 1
        probas: Numpy Array of shape n x m (n samples, m labels)

    Return:
        ranks: Numpy Array of shape n x 1 (values between (-1, 1), the larger the more suspicious)
    """
    ranks = __rank_suspicious(labels, probas=probas)
    return ranks

# Cell
@dispatch
def suspect(labels: pd.Series, *, probas: pd.DataFrame) -> pd.Series:
    """ Rank the suspicious using Pandas API
    """
    logger.debug(f"Pandas series labels and dataframe probas must have same length {len(labels)} vs {len(probas)}")
    ranks = __rank_suspicious(labels=labels.to_numpy(), probas=probas.iloc[:, [0]].to_numpy())
    ranks = pd.Series(ranks, index=labels.index)
    return ranks

# Cell
@dispatch
def suspect(labels: list, *, probas: list) -> list:
    """ Rank the  suspicious using normal Python lists
    Params:
        labels: a 1 dimension Python list of len n (n samples)
        probas: a 2 dimension Python list shape n x m (m labels)

    Return:
        ranks: a 1 dimension Python list of len n (values between (-1, 1), the larger the more suspicious)
    """
    logger.debug("Normal Python lists")
    ranks = __rank_suspicious(np.array(labels), probas=np.array(probas))
    # ranks = ranks.tolist() # DO NOT do this. Keep the numerical accuracy
    return ranks