import evaluate as ev
from sklearn.metrics import ndcg_score
import datasets

@ev.utils.file_utils.add_start_docstrings("_DESCRIPTION", "_KWARGS_DESCRIPTION")
class nDCG(ev.Metric):
    def _info(self):
        return ev.MetricInfo(
            module_type="metric",
            description="nDCG",
            citation="None",
            inputs_description="None",
            features=datasets.Features({
                'predictions': datasets.Sequence(datasets.Value('float')),
                'references': datasets.Sequence(datasets.Value('float'))
            }),
            homepage="none",
        )

    def _compute(self, predictions, references, sample_weight=None, k=5):
        """Returns the scores"""
        score = ndcg_score(references, predictions, k=k, sample_weight=sample_weight)
        return {
            "nDCG@"+str(k): score
        }