"""Scoring and aggregation.

Two deliberate choices, both of which must be reported alongside any number:

1. Parse failures count as incorrect, not as dropped samples. Dropping them
   inflates scores for models with unreliable output formats. The dropped count
   is surfaced separately so the reader can see the rate.
2. Intervals are bootstrap percentile CIs. Buckets can be small once results are
   split by target size, and a point estimate alone hides that.
"""

from __future__ import annotations

import random
from collections import defaultdict
from dataclasses import dataclass, field, asdict
from typing import Iterable, Sequence


@dataclass
class SampleResult:
    """Outcome for one sample. `pred_point` is in original-image pixels."""

    sample_id: str
    correct: bool
    pred_point: tuple[float, float] | None
    gt_box: tuple[float, float, float, float]
    area_fraction: float
    size_bucket: str
    data_type: str  # "text" or "icon" in ScreenSpot-Pro
    application: str = ""
    parse_failed: bool = False
    off_image: bool = False  # Prediction physically outside the image
    # The model's own predicted extent, when it emitted one, in the same
    # original-image pixels as `pred_point`. None for point-only models. `correct`
    # never reads it — ScreenSpot-Pro scores point-in-ground-truth-box — so this
    # cannot shift any accuracy number. It is recorded so predicted-vs-true IoU
    # stays computable from a results file without re-running the benchmark.
    pred_box: tuple[float, float, float, float] | None = None
    raw_response: str = ""
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    latency_s: float = 0.0
    error: str = ""


@dataclass
class Interval:
    point: float
    low: float
    high: float

    def __str__(self) -> str:
        return f"{self.point:.3f} [{self.low:.3f}, {self.high:.3f}]"


@dataclass
class GroupSummary:
    label: str
    n: int
    n_correct: int
    accuracy: Interval
    parse_failures: int = 0
    cost_usd: float = 0.0
    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def cost_per_correct(self) -> float | None:
        """Dollars per correct click — the cost axis the papers omit."""
        return None if self.n_correct == 0 else self.cost_usd / self.n_correct


def bootstrap_accuracy(
    flags: Sequence[bool],
    n_resamples: int = 10_000,
    confidence: float = 0.95,
    seed: int = 0,
) -> Interval:
    """Percentile bootstrap CI for a proportion."""
    n = len(flags)
    if n == 0:
        return Interval(float("nan"), float("nan"), float("nan"))

    point = sum(flags) / n
    if n == 1:
        return Interval(point, point, point)

    rng = random.Random(seed)
    as_int = [1 if f else 0 for f in flags]
    means = []
    for _ in range(n_resamples):
        total = 0
        for _ in range(n):
            total += as_int[rng.randrange(n)]
        means.append(total / n)

    means.sort()
    alpha = (1.0 - confidence) / 2.0
    lo_idx = max(0, min(n_resamples - 1, int(alpha * n_resamples)))
    hi_idx = max(0, min(n_resamples - 1, int((1.0 - alpha) * n_resamples)))
    return Interval(point, means[lo_idx], means[hi_idx])


def summarize(
    results: Iterable[SampleResult],
    label: str = "overall",
    n_resamples: int = 10_000,
    seed: int = 0,
) -> GroupSummary:
    items = list(results)
    flags = [r.correct for r in items]
    return GroupSummary(
        label=label,
        n=len(items),
        n_correct=sum(flags),
        accuracy=bootstrap_accuracy(flags, n_resamples=n_resamples, seed=seed),
        parse_failures=sum(1 for r in items if r.parse_failed),
        cost_usd=sum(r.cost_usd for r in items),
        input_tokens=sum(r.input_tokens for r in items),
        output_tokens=sum(r.output_tokens for r in items),
    )


def summarize_by(
    results: Iterable[SampleResult],
    key: str,
    n_resamples: int = 10_000,
    seed: int = 0,
) -> list[GroupSummary]:
    """Group by a SampleResult attribute — `size_bucket`, `data_type`, ..."""
    buckets: dict[str, list[SampleResult]] = defaultdict(list)
    for r in results:
        buckets[str(getattr(r, key))].append(r)
    return [
        summarize(v, label=k, n_resamples=n_resamples, seed=seed)
        for k, v in sorted(buckets.items())
    ]


def compare(
    baseline: Sequence[SampleResult],
    variant: Sequence[SampleResult],
    n_resamples: int = 10_000,
    confidence: float = 0.95,
    seed: int = 0,
) -> Interval:
    """Paired bootstrap CI for the accuracy delta (variant - baseline).

    Pairing on sample_id is what makes a scaffold's gain measurable at this scale:
    the two runs see identical samples, so per-sample variance cancels. If the
    interval contains zero the gain is not distinguishable from noise.
    """
    base_by_id = {r.sample_id: r.correct for r in baseline}
    deltas = [
        (1 if r.correct else 0) - (1 if base_by_id[r.sample_id] else 0)
        for r in variant
        if r.sample_id in base_by_id
    ]

    n = len(deltas)
    if n == 0:
        return Interval(float("nan"), float("nan"), float("nan"))

    point = sum(deltas) / n
    if n == 1:
        return Interval(point, point, point)

    rng = random.Random(seed)
    means = []
    for _ in range(n_resamples):
        total = 0
        for _ in range(n):
            total += deltas[rng.randrange(n)]
        means.append(total / n)

    means.sort()
    alpha = (1.0 - confidence) / 2.0
    lo_idx = max(0, min(n_resamples - 1, int(alpha * n_resamples)))
    hi_idx = max(0, min(n_resamples - 1, int((1.0 - alpha) * n_resamples)))
    return Interval(point, means[lo_idx], means[hi_idx])


def result_to_dict(r: SampleResult) -> dict:
    return asdict(r)