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

__all__ = ['calculate_precision_recall', 'calculate_ndcg', 'recall', 'ndcg', 'recalls_and_ndcgs_for_ks']

# Cell
import numpy as np
import torch
from sklearn.metrics import precision_score, recall_score, ndcg_score

# Cell
def calculate_precision_recall(X, y_true, y_pred, N, threshold):
    """Calculate the precision and recall scores.

    Args:
        X
        y_true
        y_pred
        N
        threshold

    Returns:
        precision_score (float)
        recall_score (float)
    """
    precision = 0
    recall = 0
    count = 0

    rec_true = np.array([1 if rating >= threshold else 0 for rating in y_true])
    rec_pred = np.zeros(y_pred.size)

    for user_id in np.unique(X[:,0]):
        indices = np.where(X[:,0] == user_id)[0]

        rec_true = np.array([1 if y_true[i] >= threshold else 0 for i in indices])

        if (np.count_nonzero(rec_true) > 0): # ignore test users without relevant ratings

            user_pred = np.array([y_pred[i] for i in indices])
            rec_pred = np.zeros(indices.size)

            for pos in np.argsort(user_pred)[-N:]:
                if user_pred[pos] >= threshold:
                    rec_pred[pos] = 1

            precision += precision_score(rec_true, rec_pred, zero_division=0)
            recall += recall_score(rec_true, rec_pred)
            count += 1

    return precision/count, recall/count

# Cell
def calculate_ndcg(X, y_true, y_pred, N):
    """Calculate the NDCG score.

    Args:
        X
        y_true
        y_pred
        N

    Returns:
        ndcg_score (float)
    """
    ndcg = 0
    count = 0

    for user_id in np.unique(X[:,0]):
        indices = np.where(X[:,0] == user_id)[0]

        user_true = np.array([y_true[i] for i in indices])
        user_pred = np.array([y_pred[i] for i in indices])

        user_true = np.expand_dims(user_true, axis=0)
        user_pred = np.expand_dims(user_pred, axis=0)

        if user_true.size > 1:
            ndcg += ndcg_score(user_true, user_pred, k=N, ignore_ties=False)
            count += 1

    return ndcg / count

# Cell
def recall(scores, labels, k):
    scores = scores.cpu()
    labels = labels.cpu()
    rank = (-scores).argsort(dim=1)
    cut = rank[:, :k]
    hit = labels.gather(1, cut)
    return (hit.sum(1).float() / torch.min(torch.Tensor([k]).to(hit.device), labels.sum(1).float())).mean().cpu().item()

# Cell
def ndcg(scores, labels, k):
    scores = scores.cpu()
    labels = labels.cpu()
    rank = (-scores).argsort(dim=1)
    cut = rank[:, :k]
    hits = labels.gather(1, cut)
    position = torch.arange(2, 2+k)
    weights = 1 / torch.log2(position.float())
    dcg = (hits.float() * weights).sum(1)
    idcg = torch.Tensor([weights[:min(int(n), k)].sum() for n in labels.sum(1)])
    ndcg = dcg / idcg
    return ndcg.mean()

# Cell
def recalls_and_ndcgs_for_ks(scores, labels, ks):
    metrics = {}

    scores = scores
    labels = labels
    answer_count = labels.sum(1)

    labels_float = labels.float()
    rank = (-scores).argsort(dim=1)
    cut = rank
    for k in sorted(ks, reverse=True):
       cut = cut[:, :k]
       hits = labels_float.gather(1, cut)
       metrics['Recall@%d' % k] = \
           (hits.sum(1) / torch.min(torch.Tensor([k]).to(labels.device), labels.sum(1).float())).mean().cpu().item()

       position = torch.arange(2, 2+k)
       weights = 1 / torch.log2(position.float())
       dcg = (hits * weights.to(hits.device)).sum(1)
       idcg = torch.Tensor([weights[:min(int(n), k)].sum() for n in answer_count]).to(dcg.device)
       ndcg = (dcg / idcg).mean()
       metrics['NDCG@%d' % k] = ndcg.cpu().item()

    return metrics