# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/20-alg.ipynb.

# %% auto 0
__all__ = ['bias_proba2label', 'suspect']

# %% ../nbs/20-alg.ipynb 2
import numpy as np
import pandas as pd
from loguru import logger

# %% ../nbs/20-alg.ipynb 21
def bias_proba2label(
    *,
    probas: np.ndarray,
    labels: np.ndarray,
) -> np.ndarray:
    """Lọc xác suất mô hình tính toán cho lớp nhãn đã gán với từng mẫu dữ liệu.`"""
    confidence = np.array([probas[i, j] for i, j in enumerate(labels)])
    bias = 1 - confidence
    logger.debug(bias)
    return bias

# %% ../nbs/20-alg.ipynb 25
def avg_confidence(
    *,
    probas: np.ndarray,
    labels: np.ndarray,
):
    """Compute the average model confidence for samples in each class. If there is no sample in some specific class, take np.nan."""
    m_class = probas.shape[1]  # number of classes
    thress = np.zeros(m_class)
    for j in range(m_class):
        logger.debug(f"j {j}")
        i = labels == j
        logger.debug(f"i {i}")
        if i.any():
            p_ij = probas[i, j]
            logger.debug(f"p_ij {p_ij}")
            thress[j] = p_ij.mean()
        else:
            thress[j] = np.nan
    logger.debug(thress)
    return thress

# %% ../nbs/20-alg.ipynb 29
def find_likeliests(
    *,
    thress: np.ndarray,  # thresholds for each class
    probas: np.ndarray,  # computed probabily matrix
):
    """
    Dự đoán nhãn đáng tin nhất cho từng mẫu dữ liệu.
    Trường hợp không dự đoán được thì gán -1 vào hàng tương ứng.
    """
    likelies = probas > thress
    ps_filtered = np.where(likelies, probas, -1)
    most_ll = ps_filtered.argmax(axis=1)
    most_ll = np.where((likelies).any(axis=1), most_ll, -1)
    logger.debug(most_ll)
    return most_ll

# %% ../nbs/20-alg.ipynb 38
from scipy.sparse import coo_matrix


def est_noise_mat(
    *,
    likeliests: np.ndarray,  # likeliest labels
    labels: np.ndarray,  # observed labels
    m_class: int,  # number of classes
    normalize: bool = True,
    cleanlab_compat: bool = False  # CleanLab compatible. See https://github.com/cleanlab/cleanlab/blob/3fb4133a5c4f132ecdcce779580dc3aadbe52144/cleanlab/count.py#L295-L296
) -> coo_matrix:
    z = np.column_stack((labels[likeliests >= 0], likeliests[likeliests >= 0]))
    ij, cnt = np.unique(z, axis=0, return_counts=True)
    counts_mat = coo_matrix((cnt, (ij[:, 0], ij[:, 1])), shape=(m_class, m_class))
    if cleanlab_compat:
        counts_mat.setdiag(counts_mat.diagonal().clip(min=1))  # to avoid division by 0
    logger.debug("counts_mat\n{}", counts_mat.toarray())
    o_cnt = np.array(
        [
            {k: v for k, v in zip(*np.unique(labels, return_counts=True))}.get(i, 0)
            for i in range(m_class)
        ]
    )
    l_cnt = np.array(
        [
            {k: v for k, v in zip(*np.unique(likeliests, return_counts=True))}.get(i, 0)
            for i in range(m_class)
        ]
    )
    l_cnt = np.maximum(l_cnt, 1)  # to avoid division by 0
    logger.debug("o_cnt\n{}", o_cnt)
    logger.debug("l_cnt\n{}", l_cnt)
    o_l_rate = o_cnt / l_cnt
    logger.debug("o_l_rate\n{}", o_l_rate)
    res = counts_mat.multiply(o_l_rate)
    if normalize:
        res = res / res.sum()
    logger.debug("res\n{}", res.toarray())
    return res

# %% ../nbs/20-alg.ipynb 40
def est_noise_cnt(
    *,
    likeliests: np.ndarray,  # likeliest labels
    labels: np.ndarray,  # observed labels
    m_class: int,  # number of classes
) -> int:
    """Estimate number of erronous labels"""
    # noise_cnt = sum(likeliests != labels) # fast and lazy
    noise_mat = est_noise_mat(
        labels=labels, likeliests=likeliests, m_class=m_class, normalize=True
    )
    noise_rate = 1 - sum(noise_mat.diagonal())
    logger.debug("noise_rate {}", noise_rate)
    noise_cnt = int(noise_rate * len(labels))
    logger.debug("noise_cnt {}", noise_cnt)
    return noise_cnt

# %% ../nbs/20-alg.ipynb 44
def rank_suspicious(
    *,
    probas: np.ndarray,
    likeliests: np.ndarray,
    labels: np.ndarray,
):
    from scipy.sparse import lil_matrix

    n, m_class = probas.shape
    e = lil_matrix((n, 1))
    ids = (likeliests != labels) & (likeliests != -1)
    # logger.debug(ids)
    # TODO: make the loop below faster
    for k in np.arange(n)[np.where(ids)]:
        e[k, 0] = probas[k, likeliests[k]] - probas[k, labels[k]]
    # logger.debug(e)
    coo = e.tocoo()
    id = coo.row
    err = coo.data
    # logger.debug(id)
    # logger.debug(err)
    err_df = pd.DataFrame(dict(id=id, err=err)).set_index("id")
    logger.debug(err_df)
    return err_df

# %% ../nbs/20-alg.ipynb 48
def suspect(
    probas: np.ndarray,  # probabilities
    # *, # turn of keyword arguments due to nbdev show_doc bug
    labels: np.ndarray,  # observed labels
) -> np.ndarray:
    "The internal method to rank the suspicious labels given probas from a classifier."
    # logger.debug(f"Shape of labels and probas: {labels.shape} vs {probas.shape}")
    assert len(labels) == len(probas)
    assert (probas >= 0).all() and (
        probas <= 1
    ).all()  # probas should be between 0 and 1

    bias = bias_proba2label(probas=probas, labels=labels)
    thress = avg_confidence(probas=probas, labels=labels)
    likeliests = find_likeliests(thress=thress, probas=probas)
    noise_cnt = est_noise_cnt(
        likeliests=likeliests, labels=labels, m_class=probas.shape[1]
    )
    ids = np.argsort(-bias)[:noise_cnt]
    err_df = (
        pd.DataFrame(
            dict(
                # id=ids,
                # err=bias[ids]
                err=bias,
                suspected=False,
            )
        )
        # .set_index("id")
    )
    err_df.loc[ids, "suspected"] = True
    return err_df
