"""Aggregate blame over a batch of failing test cases."""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Callable, Mapping, Sequence

import numpy as np

from .attribution import CaseBlame, _radius, _separated, analyze_case
from .noise import NoiseSource
from .pipeline import Case, Pipeline

__all__ = ["BatchBlame", "StepSummary", "blame_batch"]


@dataclass
class StepSummary:
    name: str
    fix_alone: float                     # mean of v({i}) - v(empty): share of failures fixed by i alone
    fix_alone_ci: tuple[float, float]
    shapley: float                       # mean Shapley value over analyzed failures
    shapley_ci: tuple[float, float]
    share: float                         # sum phi_i / sum (v(all) - v(empty))
    leave_one_out: float
    last_touch: int                      # number of cases where naive "last touch" blames i


@dataclass
class BatchBlame:
    pipeline: str
    n_cases: int                         # cases run factually (may stop early)
    n_failures: int
    cases: list[CaseBlame]
    steps: list[StepSummary]
    unresolved: int                      # failures not fixed even with every fix applied
    cost: float
    stop_reason: str
    settings: dict = field(default_factory=dict)

    def ranking(self) -> list[StepSummary]:
        return sorted(self.steps, key=lambda s: -s.shapley)


def _summaries(cases: list[CaseBlame], players: list[str], bound: str,
               delta: float) -> list[StepSummary]:
    gain = np.array([c.v_full - c.v_empty for c in cases])
    out = []
    for p in players:
        phi = np.array([c.shapley.values[p] for c in cases])
        fa = np.array([c.fix_alone[p] for c in cases])
        r_phi = _radius(bound, phi, 2.0, delta) if len(phi) else float("inf")
        r_fa = _radius(bound, fa, 2.0, delta) if len(fa) else float("inf")
        out.append(StepSummary(
            p, float(fa.mean()), (float(fa.mean() - r_fa), float(fa.mean() + r_fa)),
            float(phi.mean()), (float(phi.mean() - r_phi), float(phi.mean() + r_phi)),
            float(phi.sum() / gain.sum()) if gain.sum() != 0 else float("nan"),
            float(np.mean([c.leave_one_out[p] for c in cases])),
            sum(c.last_touch == p for c in cases)))
    return out


def blame_batch(pipeline: Pipeline, cases: Sequence[Case], *, fixes: Mapping[str, str] | None = None,
                players: Sequence[str] | None = None, mode: str = "recorded",
                n_worlds: int = 1, crn: bool = True, method: str = "auto",
                mc: Mapping | None = None, is_failure: Callable[[float], bool] = lambda s: s < 1.0,
                ci: str = "normal", delta: float = 0.05,
                stop_when_separated: bool = False, min_failures: int = 10,
                stop_bound: str = "bernstein", value_range: float = 2.0,
                growth: float = 1.5, budget: float | None = None,
                seed: int = 0) -> BatchBlame:
    """Run every case factually, analyze the failures, and aggregate.

    ``ci="normal"`` gives the usual mean +- 1.96 SE intervals across cases;
    ``ci="bernstein"`` gives finite-sample empirical-Bernstein intervals
    (per-case Shapley values lie in [-1, 1] for a 0/1 metric).

    With ``stop_when_separated=True``, failures are analyzed in the given
    order and analysis stops as soon as the top step's Shapley interval
    separates from all others (empirical-Bernstein, union bound over steps
    and geometric checkpoints, as in :func:`~agentblame.attribution.mc_shapley`),
    or when ``budget`` (pipeline cost) is exhausted. Shuffle ``cases`` first
    so that the analyzed prefix is a random sample. ``value_range`` is the
    range of a per-case Shapley value (2 for a 0/1 metric: phi in [-1, 1]; 1 if
    you know no fix can hurt). ``stop_bound="normal"`` uses asymptotic
    intervals instead: it stops much earlier but carries no finite-sample,
    anytime-valid guarantee.
    """
    c0 = pipeline.meter.total
    analyzed: list[CaseBlame] = []
    n_fail = n_seen = 0
    reason = "all failures analyzed"
    k = 0
    next_check = min_failures
    for case in cases:
        n_seen += 1
        factual = pipeline.run(case.inputs, noise=NoiseSource(case.seed))
        if not is_failure(factual.score):
            continue
        n_fail += 1
        cb = analyze_case(pipeline, case, fixes=fixes, players=players, mode=mode,
                          n_worlds=n_worlds, crn=crn, method=method, mc=mc,
                          seed=seed, factual=factual)
        analyzed.append(cb)
        if budget is not None and pipeline.meter.total - c0 >= budget:
            reason = "budget"
            break
        if stop_when_separated and len(analyzed) >= next_check:
            k += 1
            ps = analyzed[0].players
            dk = delta / (len(ps) * k * (k + 1))
            X = np.array([[c.shapley.values[p] for p in ps] for c in analyzed])
            means = X.mean(axis=0)
            rad = np.array([_radius(stop_bound, X[:, j], value_range, dk) for j in range(len(ps))])
            if _separated(means, rad, "top1"):
                reason = f"top step separated ({stop_bound}) after {len(analyzed)} failures"
                break
            next_check = max(len(analyzed) + 1, int(math.ceil(min_failures * growth ** k)))
    if not analyzed:
        return BatchBlame(pipeline.name, n_seen, 0, [], [], 0, pipeline.meter.total - c0,
                          "no failures")
    steps = _summaries(analyzed, analyzed[0].players, ci, delta)
    unresolved = sum(c.v_full - c.v_empty <= 0 for c in analyzed)
    settings = dict(mode=mode, n_worlds=n_worlds, crn=crn, method=method, ci=ci, delta=delta)
    return BatchBlame(pipeline.name, n_seen, n_fail, analyzed, steps, unresolved,
                      pipeline.meter.total - c0, reason, settings)