"""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('perf_stats')
def perf_stats(results: List[SampleResult], metric: str):
    """Performance profile over per-sample usage: latency/ttft percentiles,
    throughput, token stats. Attach to any metric (reads SampleResult.usage).

    Report shape: metric_groups['perf'] = {p50_latency_s, p95_latency_s, ...}
    """
    import statistics

    def _pct(vals, q):
        if not vals:
            return None
        vals = sorted(vals)
        k = max(0, min(len(vals) - 1, int(round(q / 100 * (len(vals) - 1)))))
        return round(vals[k], 3)

    lat = [float((r.usage or {}).get('latency_s', 0) or 0) for r in results
           if (r.usage or {}).get('latency_s')]
    ttft = [float(r.usage['ttft_s']) for r in results
            if (r.usage or {}).get('ttft_s') is not None]
    itl = [float(r.usage['itl_mean_s']) for r in results
           if (r.usage or {}).get('itl_mean_s') is not None]
    n = len(lat)
    in_toks = [int((r.usage or {}).get('input_tokens', 0) or 0) for r in results]
    out_toks = [int((r.usage or {}).get('output_tokens', 0) or 0) for r in results]
    in_tok, out_tok = sum(in_toks), sum(out_toks)
    retried = sum(1 for r in results if (r.usage or {}).get('retries'))
    ok = sum(1 for r in results
             if (r.usage or {}).get('http_status') in (None, 200))  # None=unmeasured
    wall = sum(lat)
    # TPOT per request: (latency - ttft) / max(output_tokens - 1, 1)
    tpots = []
    for r in results:
        u = r.usage or {}
        lt, tf, ot = u.get('latency_s'), u.get('ttft_s'), u.get('output_tokens')
        if lt and tf is not None and ot and ot > 1:
            tpots.append((lt - tf) / (ot - 1))
    out = {
        'n_requests': n,
        'success_rate': round(ok / n, 4) if n else None,
        'latency_mean_s': round(statistics.mean(lat), 3) if lat else None,
        'latency_p50_s': _pct(lat, 50), 'latency_p90_s': _pct(lat, 90),
        'latency_p95_s': _pct(lat, 95), 'latency_p99_s': _pct(lat, 99),
        'ttft_mean_s': round(statistics.mean(ttft), 3) if ttft else None,
        'ttft_p90_s': _pct(ttft, 90), 'ttft_p99_s': _pct(ttft, 99),
        'tpot_mean_s': round(statistics.mean(tpots), 4) if tpots else None,
        'tpot_p90_s': _pct(tpots, 90), 'tpot_p99_s': _pct(tpots, 99),
        'input_tokens_mean': round(statistics.mean(in_toks), 1) if in_toks else 0,
        'output_tokens_mean': round(statistics.mean(out_toks), 1) if out_toks else 0,
        'total_tokens': in_tok + out_tok,
        'output_tps': round(out_tok / wall, 2) if wall else None,    # tokens/s
        'request_qps': round(n / wall, 4) if wall else None,        # req/s
        'wall_latency_s': round(wall, 1),
        'retry_rate': round(retried / n, 3) if n else None,
    }
    return out


@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])}