"""Locked split retrieval protocol shared by learned and baseline embeddings."""

from __future__ import annotations

import random
from collections import defaultdict
from collections.abc import Callable, Mapping, Sequence

import pandas as pd

from .metrics import expected_tie_aware_metrics
from .splits import validate_split


ScoreFunction = Callable[[Sequence[str], Sequence[str]], Mapping[str, float]]


def evaluate_retrieval(
    assignments: pd.DataFrame,
    split_name: str,
    methods: Mapping[str, ScoreFunction],
    reference_size: int,
    draws: int,
    seed: int,
    recall_at: Sequence[int] = (10, 50, 100),
    ndcg_at: Sequence[int] = (10, 50, 100),
) -> pd.DataFrame:
    validate_split(assignments)
    split = assignments[assignments["split"] == split_name]
    universe = sorted(split["bgc_id"].astype(str).unique())
    group_members: dict[str, list[str]] = defaultdict(list)
    for row in split.itertuples(index=False):
        group_members[str(row.group_id)].append(str(row.bgc_id))
    eligible = {
        group: sorted(set(members))
        for group, members in group_members.items()
        if len(set(members)) > reference_size
    }
    if not eligible:
        raise ValueError(f"No {split_name} groups have more than {reference_size} members")

    rows: list[dict[str, object]] = []
    for group_id, members in sorted(eligible.items()):
        for draw in range(draws):
            draw_seed = f"{seed}:{split_name}:{group_id}:{draw}"
            references = random.Random(draw_seed).sample(members, reference_size)
            relevant = set(members).difference(references)
            candidates = [identifier for identifier in universe if identifier not in references]
            for method_name, score_function in methods.items():
                scores = dict(score_function(candidates, references))
                if set(scores) != set(candidates):
                    raise ValueError(f"{method_name} did not score the complete candidate universe")
                metrics = expected_tie_aware_metrics(scores, relevant, recall_at, ndcg_at)
                rows.append(
                    {
                        "split": split_name,
                        "group_id": group_id,
                        "draw": draw,
                        "method": method_name,
                        "reference_ids": ";".join(references),
                        **metrics,
                    }
                )
    return pd.DataFrame(rows)