"""Metric aggregation and slicing utilities.

This module provides functions for aggregating metrics across different
dimensions (difficulty, domain, tier) and computing summary statistics.
"""

from collections import defaultdict
from typing import Callable, Sequence

from embergraphrouter.evals.metrics.core import (
    accuracy,
    frac_of_oracle,
    mean_with_ci,
    model_distribution,
    tier_distribution,
)
from embergraphrouter.evals.runners.base import (
    AggregateMetrics,
    ExampleResult,
    SliceMetrics,
)


def aggregate_by_slice(
    results: Sequence[ExampleResult],
    slice_key: str | Callable[[ExampleResult], str],
) -> dict[str, SliceMetrics]:
    """Aggregate results by a slicing dimension.

    Args:
        results: Sequence of evaluation results.
        slice_key: Either a string key to extract from metadata, or a
            callable that returns the slice value for each result.

    Returns:
        Dictionary mapping slice values to SliceMetrics.
    """
    # Group results by slice value
    slices: dict[str, list[ExampleResult]] = defaultdict(list)

    for r in results:
        if callable(slice_key):
            value = slice_key(r)
        elif slice_key == "difficulty":
            value = r.difficulty
        elif slice_key == "domain":
            value = r.domain
        elif slice_key == "tier":
            value = r.tier or "unknown"
        else:
            value = r.metadata.get(slice_key, "unknown")

        slices[str(value)].append(r)

    # Compute metrics for each slice
    result_dict: dict[str, SliceMetrics] = {}
    for slice_value, slice_results in slices.items():
        result_dict[slice_value] = SliceMetrics(
            slice_name=slice_key if isinstance(slice_key, str) else "custom",
            slice_value=slice_value,
            n_examples=len(slice_results),
            quality=compute_quality_metrics(slice_results),
            efficiency=compute_efficiency_metrics(slice_results),
            routing=compute_routing_metrics(slice_results),
        )

    return result_dict


def compute_quality_metrics(
    results: Sequence[ExampleResult],
) -> dict[str, AggregateMetrics]:
    """Compute quality metrics for a set of results.

    Args:
        results: Sequence of evaluation results.

    Returns:
        Dictionary of quality metrics with confidence intervals.
    """
    if not results:
        return {}

    metrics: dict[str, AggregateMetrics] = {}

    # Accuracy
    metrics["accuracy"] = accuracy(results)

    # Score (if available and different from binary accuracy)
    scores = [r.score for r in results]
    if any(s not in (0.0, 1.0) for s in scores):
        metrics["score"] = mean_with_ci(scores)

    return metrics


def compute_efficiency_metrics(
    results: Sequence[ExampleResult],
) -> dict[str, AggregateMetrics]:
    """Compute efficiency metrics for a set of results.

    Args:
        results: Sequence of evaluation results.

    Returns:
        Dictionary of efficiency metrics with confidence intervals.
    """
    if not results:
        return {}

    metrics: dict[str, AggregateMetrics] = {}

    # Cost
    costs = [r.cost_usd for r in results]
    if any(c > 0 for c in costs):
        metrics["cost_usd"] = mean_with_ci(costs)

    # Latency
    latencies = [r.latency_ms for r in results]
    if any(lat > 0 for lat in latencies):
        metrics["latency_ms"] = mean_with_ci(latencies)

    # Tokens
    tokens_in = [r.tokens_in for r in results if r.tokens_in > 0]
    if tokens_in:
        metrics["tokens_in"] = mean_with_ci([float(t) for t in tokens_in])

    tokens_out = [r.tokens_out for r in results if r.tokens_out > 0]
    if tokens_out:
        metrics["tokens_out"] = mean_with_ci([float(t) for t in tokens_out])

    return metrics


def compute_routing_metrics(
    results: Sequence[ExampleResult],
) -> dict[str, object]:
    """Compute routing-specific metrics for a set of results.

    Args:
        results: Sequence of evaluation results.

    Returns:
        Dictionary of routing metrics (distributions, fractions, etc.)
    """
    if not results:
        return {}

    metrics: dict[str, object] = {}

    # Tier distribution
    tier_dist = tier_distribution(results)
    if tier_dist:
        metrics["tier_distribution"] = tier_dist

    # Model distribution
    model_dist = model_distribution(results)
    if model_dist:
        metrics["model_distribution"] = model_dist

    # Fraction of oracle (if oracle data available)
    frac_oracle = frac_of_oracle(results)
    if frac_oracle > 0:
        metrics["frac_of_oracle"] = frac_oracle

    # Routing signal statistics (if available)
    mean_logits = []
    logit_margins = []
    for r in results:
        signals = r.routing_signals
        if "mean_logit" in signals:
            mean_logits.append(signals["mean_logit"])
        if "logit_margin" in signals:
            logit_margins.append(signals["logit_margin"])

    if mean_logits:
        metrics["mean_logit_stats"] = mean_with_ci(mean_logits).to_dict()
    if logit_margins:
        metrics["logit_margin_stats"] = mean_with_ci(logit_margins).to_dict()

    return metrics


def compute_tier_analysis(
    results: Sequence[ExampleResult],
) -> dict[str, dict[str, object]]:
    """Compute detailed per-tier analysis.

    Provides accuracy and cost breakdown by tier to understand routing value.

    Args:
        results: Sequence of evaluation results with tier information.

    Returns:
        Dictionary mapping tier names to analysis results.
    """
    tier_slices = aggregate_by_slice(results, "tier")

    analysis: dict[str, dict[str, object]] = {}
    default_metrics = AggregateMetrics(mean=0.0)
    for tier_name, slice_metrics in tier_slices.items():
        accuracy_metrics = slice_metrics.quality.get("accuracy", default_metrics)
        cost_metrics = slice_metrics.efficiency.get("cost_usd", default_metrics)
        latency_metrics = slice_metrics.efficiency.get("latency_ms", default_metrics)

        analysis[tier_name] = {
            "n_examples": slice_metrics.n_examples,
            "pct_of_total": slice_metrics.n_examples / len(results) if results else 0.0,
            "accuracy": accuracy_metrics.to_dict(),
            "avg_cost_usd": cost_metrics.to_dict(),
            "avg_latency_ms": latency_metrics.to_dict(),
        }

    return analysis


def compute_difficulty_analysis(
    results: Sequence[ExampleResult],
) -> dict[str, dict[str, object]]:
    """Compute detailed per-difficulty analysis.

    Shows how routing performs across different difficulty levels.

    Args:
        results: Sequence of evaluation results with difficulty metadata.

    Returns:
        Dictionary mapping difficulty levels to analysis results.
    """
    diff_slices = aggregate_by_slice(results, "difficulty")

    analysis: dict[str, dict[str, object]] = {}
    for diff_level, slice_metrics in diff_slices.items():
        analysis[diff_level] = {
            "n_examples": slice_metrics.n_examples,
            "pct_of_total": slice_metrics.n_examples / len(results) if results else 0.0,
            "accuracy": slice_metrics.quality.get("accuracy", AggregateMetrics(mean=0.0)).to_dict(),
            "tier_distribution": slice_metrics.routing.get("tier_distribution", {}),
        }

    return analysis


__all__ = [
    "aggregate_by_slice",
    "compute_quality_metrics",
    "compute_efficiency_metrics",
    "compute_routing_metrics",
    "compute_tier_analysis",
    "compute_difficulty_analysis",
]
