"""Офлайн-метрики на отложенной по времени выборке: NDCG@10, Recall@50 (docs/recommendations.md#1)."""

from __future__ import annotations

import math
from collections.abc import Callable, Iterable


def ndcg_at_k(ranked: list[str], relevant: set[str], k: int = 10) -> float:
    dcg = sum(1.0 / math.log2(i + 2) for i, t in enumerate(ranked[:k]) if t in relevant)
    ideal = sum(1.0 / math.log2(i + 2) for i in range(min(k, len(relevant))))
    return dcg / ideal if ideal > 0 else 0.0


def recall_at_k(ranked: list[str], relevant: set[str], k: int = 50) -> float:
    if not relevant:
        return 0.0
    return len(set(ranked[:k]) & relevant) / len(relevant)


def evaluate(
    users: Iterable[str],
    relevant: dict[str, set[str]],
    rank: Callable[[str], list[str]],
) -> dict[str, float]:
    """Средние метрики по пользователям, у которых в тесте есть хотя бы один релевантный трек."""
    n = 0
    ndcg = rec = 0.0
    for u in users:
        rel = relevant.get(u)
        if not rel:
            continue
        ranked = rank(u)
        ndcg += ndcg_at_k(ranked, rel, 10)
        rec += recall_at_k(ranked, rel, 50)
        n += 1
    if n == 0:
        return {"users": 0, "ndcg@10": 0.0, "recall@50": 0.0}
    return {"users": n, "ndcg@10": round(ndcg / n, 4), "recall@50": round(rec / n, 4)}