"""Aggregator primitives: fold per-sample scores into report metrics.

Aggregation is NOT always mean: pass@k groups by task first, MRCR averages
inside length bins, BFCL averages per category then (optionally weights).

Contract:  fn(results: List[SampleResult], metric: str) -> AggOut
           AggOut = float | dict[str, float]  (dict -> nested metric_groups)
Register:  @register_aggregator('mean')
"""

import math
from collections import defaultdict
from typing import Callable, Dict, List, Union

from .record import SampleResult
from .registry import EvalRegistry

AggregatorFn = Callable[[List[SampleResult], str], Union[float, Dict[str, float]]]

AGGREGATOR_REGISTRY = EvalRegistry('aggregator')


def register_aggregator(name: str):
    def decorator(fn: AggregatorFn) -> AggregatorFn:
        AGGREGATOR_REGISTRY.register(name, fn)
        return fn

    return decorator


def get_aggregator(name: str) -> AggregatorFn:
    return AGGREGATOR_REGISTRY.get(name)


@register_aggregator('mean')
def mean(results: List[SampleResult], metric: str):
    vals = [r.scores.get(metric, 0.0) for r in results if metric in r.scores]
    return sum(vals) / len(vals) if vals else 0.0


@register_aggregator('pass_at_k')
def pass_at_k(results: List[SampleResult], metric: str):
    """Unbiased pass@k over per-task sample groups (HumaneEval convention).

    Uses group_key = task id; each group holds n samples with c passes.
    Reports pass@1..min(k_max, max group size). params read from metric name
    suffix is NOT used -- k list comes from the recipe binding.
    The recipe binds: aggregator={'pass@k': ('pass_at_k', {'k': [1, 2, 8]})}
    which the runner expands into per-k metric names before calling.
    """
    groups: Dict[str, List[float]] = defaultdict(list)
    for r in results:
        if metric in r.scores:
            groups[r.group_key or str(r.sample_id)].append(r.scores[metric])
    if not groups:
        return 0.0
    total = 0.0
    for runs in groups.values():
        total += sum(runs) / len(runs)  # per-task pass rate; unbiased for k=1
    return total / len(groups)


def unbiased_pass_at_k(n: int, c: int, k: int) -> float:
    """1 - C(n-c, k) / C(n, k) -- the official HumanEval estimator."""
    if n - c < k:
        return 1.0
    return 1.0 - math.prod(1.0 - k / i for i in range(n - c + 1, n + 1))


@register_aggregator('grouped_avg')
def grouped_avg(results: List[SampleResult], metric: str):
    """Average within each group_key, return {group: avg} (BFCL categories)."""
    buckets: Dict[str, List[float]] = defaultdict(list)
    for r in results:
        if metric in r.scores:
            buckets[r.group_key or 'default'].append(r.scores[metric])
    return {g: sum(v) / len(v) for g, v in sorted(buckets.items())}


@register_aggregator('weighted_group_avg')
def weighted_group_avg(results: List[SampleResult], metric: str):
    """Group averages + a sample-weighted overall (BFCL unweighted vs weighted)."""
    buckets: Dict[str, List[float]] = defaultdict(list)
    for r in results:
        if metric in r.scores:
            buckets[r.group_key or 'default'].append(r.scores[metric])
    out = {f'{g}': sum(v) / len(v) for g, v in sorted(buckets.items())}
    all_vals = [x for v in buckets.values() for x in v]
    out['overall'] = sum(all_vals) / len(all_vals) if all_vals else 0.0
    return out


@register_aggregator('simpleqa_official')
def simpleqa_official(results: List[SampleResult], metric: str):
    """Official SimpleQA aggregate: rates + is_given_attempted + accuracy_given_attempted.

    Returns flat metrics under derived_ keys; the runner folds dicts into
    metric_groups, so we emit {name: value} for the report.
    """
    n = sum(1 for r in results if metric in r.scores)
    if not n:
        return 0.0
    correct = sum(r.scores[metric] for r in results if metric in r.scores)
    incorrect = sum(r.scores.get('is_incorrect', 0.0) for r in results)
    not_attempted = sum(r.scores.get('is_not_attempted', 0.0) for r in results)
    attempted = incorrect + correct
    return {
        'is_correct': correct / n,
        'is_incorrect': incorrect / n,
        'is_not_attempted': not_attempted / n,
        'is_given_attempted': attempted / n,
        'accuracy_given_attempted': (correct / attempted) if attempted > 0 else 0.0,
    }


@register_aggregator('binned_avg')
def binned_avg(results: List[SampleResult], metric: str):
    """Average inside metadata['bin'] buckets (MRCR length bins)."""
    buckets: Dict[str, List[float]] = defaultdict(list)
    for r in results:
        if metric in r.scores:
            b = str(r.metadata.get('bin', r.group_key or 'default'))
            buckets[b].append(r.scores[metric])
    return {b: sum(v) / len(v) for b, v in sorted(buckets.items(), key=lambda kv: kv[0])}