"""Retrieval metrics with exact expectation over tied-score orderings."""

from __future__ import annotations

import math
from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence

import numpy as np


def _score_groups(scores: Mapping[str, float]) -> list[tuple[float, list[str]]]:
    grouped: dict[float, list[str]] = defaultdict(list)
    for identifier, score in scores.items():
        if not math.isfinite(float(score)):
            raise ValueError(f"Non-finite retrieval score for {identifier}")
        grouped[float(score)].append(str(identifier))
    return sorted(grouped.items(), key=lambda item: item[0], reverse=True)


def expected_tie_aware_metrics(
    scores: Mapping[str, float],
    relevant_ids: Iterable[str],
    recall_at: Sequence[int] = (10, 50, 100),
    ndcg_at: Sequence[int] = (10, 50, 100),
) -> dict[str, float]:
    relevant = set(relevant_ids)
    missing = relevant.difference(scores)
    if missing:
        raise ValueError(f"Relevant IDs absent from candidates: {sorted(missing)[:10]}")
    if not relevant:
        raise ValueError("At least one relevant candidate is required")

    groups = _score_groups(scores)
    total_candidates = len(scores)
    total_relevant = len(relevant)
    metrics: dict[str, float] = {}

    for cutoff in recall_at:
        effective = min(int(cutoff), total_candidates)
        expected_hits = 0.0
        consumed = 0
        for _, identifiers in groups:
            size = len(identifiers)
            relevant_count = sum(identifier in relevant for identifier in identifiers)
            slots = max(0, min(size, effective - consumed))
            expected_hits += relevant_count * slots / size
            consumed += size
            if consumed >= effective:
                break
        metrics[f"recall@{cutoff}"] = expected_hits / total_relevant
        metrics[f"precision@{cutoff}"] = expected_hits / effective if effective else 0.0

    reciprocal_sum = 0.0
    average_precision_sum = 0.0
    prior_items = 0
    prior_relevant = 0
    for _, identifiers in groups:
        size = len(identifiers)
        relevant_count = sum(identifier in relevant for identifier in identifiers)
        if relevant_count:
            reciprocal_mean = np.mean(
                [1.0 / (prior_items + position) for position in range(1, size + 1)]
            )
            reciprocal_sum += relevant_count * reciprocal_mean
            for position in range(1, size + 1):
                relevant_before = (
                    (position - 1) * (relevant_count - 1) / (size - 1) if size > 1 else 0.0
                )
                probability_relevant = relevant_count / size
                expected_precision = (
                    prior_relevant + 1 + relevant_before
                ) / (prior_items + position)
                average_precision_sum += probability_relevant * expected_precision
        prior_items += size
        prior_relevant += relevant_count
    metrics["mrr"] = reciprocal_sum / total_relevant
    metrics["map"] = average_precision_sum / total_relevant

    for cutoff in ndcg_at:
        effective = min(int(cutoff), total_candidates)
        dcg = 0.0
        consumed = 0
        for _, identifiers in groups:
            size = len(identifiers)
            relevant_count = sum(identifier in relevant for identifier in identifiers)
            slots = max(0, min(size, effective - consumed))
            relevance_probability = relevant_count / size
            for offset in range(slots):
                rank = consumed + offset + 1
                dcg += relevance_probability / math.log2(rank + 1)
            consumed += size
            if consumed >= effective:
                break
        ideal_count = min(effective, total_relevant)
        ideal = sum(1.0 / math.log2(rank + 1) for rank in range(1, ideal_count + 1))
        metrics[f"ndcg@{cutoff}"] = dcg / ideal if ideal else 0.0

    tied_candidates = sum(len(ids) for _, ids in groups if len(ids) > 1)
    metrics["tie_fraction"] = tied_candidates / total_candidates if total_candidates else 0.0
    metrics["candidate_count"] = float(total_candidates)
    metrics["relevant_count"] = float(total_relevant)
    return metrics