#!/usr/bin/env python3

"""Compute bootstrap confidence intervals for Oak benchmark JSONL files."""

from __future__ import annotations

import argparse

import json

import math

import random

import statistics

from pathlib import Path

from typing import Any, Iterable

from oakbench.rows import row_returncode

from oakbench.populations import population_conflicts

from oakbench.stats import percentile_nearest

SUBOP_SUFFIXES = (".add", ".commit")

# Mirrors oakbench.reporting.TAIL_PERCENTILE_MIN_SAMPLES; kept import-free so

# this script stays runnable against result files on machines without the repo.

TAIL_MIN_SAMPLES = {"p95": 20, "p99": 100}

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("jsonl", nargs="+", type=Path)

    parser.add_argument("--metric", default="elapsed_ms")

    parser.add_argument("--bootstrap-samples", type=int, default=2000)

    parser.add_argument("--confidence", type=float, default=0.95)

    parser.add_argument("--seed", type=int, default=20260609)

    parser.add_argument("--include-subops", action="store_true")

    return parser.parse_args()

def iter_rows(paths: Iterable[Path]) -> Iterable[dict[str, Any]]:

    for path in paths:

        with path.open() as fh:

            for raw_line in fh:

                line = raw_line.strip()

                if line:

                    yield json.loads(line)

def bootstrap_ci(

    values: list[float],

    statistic: str,

    samples: int,

    confidence: float,

    rng: random.Random,

) -> tuple[float, float]:

    if len(values) <= 1:

        value = values[0] if values else float("nan")

        return value, value

    stats: list[float] = []

    n = len(values)

    for _ in range(samples):

        sample = [values[rng.randrange(n)] for _ in range(n)]

        if statistic == "median":

            stats.append(float(statistics.median(sample)))

        elif statistic == "p90":

            stats.append(percentile_nearest(sample, 90.0))

        elif statistic == "p95":

            stats.append(percentile_nearest(sample, 95.0))

        elif statistic == "p99":

            stats.append(percentile_nearest(sample, 99.0))

        elif statistic == "mean":

            stats.append(float(statistics.mean(sample)))

        else:

            raise ValueError(statistic)

    stats.sort()

    alpha = 1.0 - confidence

    low_index = max(0, min(len(stats) - 1, math.floor((alpha / 2.0) * len(stats))))

    high_index = max(0, min(len(stats) - 1, math.ceil((1.0 - alpha / 2.0) * len(stats)) - 1))

    return stats[low_index], stats[high_index]

def main() -> int:

    args = parse_args()

    grouped: dict[tuple[str, str, str], list[float]] = {}

    population_rows: dict[tuple[str, str, str], list[dict[str, Any]]] = {}

    outcomes: dict[tuple[str, str, str], list[dict[str, Any]]] = {}

    for row in iter_rows(args.jsonl):

        operation = str(row.get("operation") or "")

        if not args.include_subops and operation.endswith(SUBOP_SUFFIXES):

            continue

        key = (str(row.get("subject")), str(row.get("scenario")), operation)

        outcomes.setdefault(key, []).append(row)

        grouped.setdefault(key, [])

        if row_returncode(row) != 0:

            continue

        population_rows.setdefault(key, []).append(row)

        try:

            value = float(row[args.metric])

        except (KeyError, TypeError, ValueError):

            continue

        grouped.setdefault(key, []).append(value)

    conflicts = [(key, population_conflicts(rows)) for key, rows in population_rows.items()]

    conflicts = [(key, fields) for key, fields in conflicts if fields]

    if conflicts:

        for key, fields in conflicts:

            print(f"INVALID incompatible population {'/'.join(key)}: {', '.join(fields)}")

        print("Split inputs by measurement identity before computing statistics.")

        return 1

    rng = random.Random(args.seed)

    print(f"# Bootstrap Statistics: `{args.metric}`")

    print()

    print(f"- Inputs: {', '.join(f'`{path}`' for path in args.jsonl)}")

    print(f"- Bootstrap samples: {args.bootstrap_samples}")

    print(f"- Confidence: {args.confidence:.3f}")

    print("- Latency statistics condition on successful measured trials; all outcomes remain below.")

    print()

    print("| Subject | Scenario / operation | Attempts (excluding skips) | Successful | Failed | Unknown outcome | Skipped | Metric unmeasured |")

    print("| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |")

    for key, rows in sorted(outcomes.items()):

        codes = [row_returncode(row) for row in rows]

        skipped = codes.count(77)

        success = codes.count(0)

        unknown = sum(row.get("returncode") is None for row in rows)

        failed = len(rows) - skipped - success - unknown

        print(f"| `{key[0]}` | `{key[1]}/{key[2]}` | {len(rows)-skipped} | {success} | {failed} | {unknown} | {skipped} | {success-len(grouped[key])} |")

    print()

    print(

        "| Subject | Scenario / operation | n | median | median CI | p90 | p90 CI | "

        "p95 | p99 | max | mean | mean CI |"

    )

    print("| --- | --- | ---: | ---: | --- | ---: | --- | ---: | ---: | ---: | ---: | --- |")

    for (subject, scenario, operation), values in sorted(grouped.items()):

        if not values:

            print(f"| `{subject}` | `{scenario}/{operation}` | 0 | unmeasured | | | | | | | | |")

            continue

        median_value = float(statistics.median(values))

        p90_value = percentile_nearest(values, 90.0)

        mean_value = float(statistics.mean(values))

        median_ci = bootstrap_ci(values, "median", args.bootstrap_samples, args.confidence, rng)

        p90_ci = bootstrap_ci(values, "p90", args.bootstrap_samples, args.confidence, rng)

        mean_ci = bootstrap_ci(values, "mean", args.bootstrap_samples, args.confidence, rng)

        # Tail honesty: a nearest-rank percentile below 1/(1-p) samples is just

        # the max relabelled. Blank cells mean "unmeasurable at this n", never 0.

        n = len(values)

        p95_cell = f"{percentile_nearest(values, 95.0):.3f}" if n >= TAIL_MIN_SAMPLES["p95"] else ""

        p99_cell = f"{percentile_nearest(values, 99.0):.3f}" if n >= TAIL_MIN_SAMPLES["p99"] else ""

        print(

            f"| `{subject}` | `{scenario}/{operation}` | {n} | "

            f"{median_value:.3f} | [{median_ci[0]:.3f}, {median_ci[1]:.3f}] | "

            f"{p90_value:.3f} | [{p90_ci[0]:.3f}, {p90_ci[1]:.3f}] | "

            f"{p95_cell} | {p99_cell} | {max(values):.3f} | "

            f"{mean_value:.3f} | [{mean_ci[0]:.3f}, {mean_ci[1]:.3f}] |"

        )

    print()

    print(

        f"p95/p99 cells are blank below n={TAIL_MIN_SAMPLES['p95']}/n={TAIL_MIN_SAMPLES['p99']}: "

        "at smaller n the nearest-rank tail percentile equals the max, so claiming it would overstate precision."

    )

    return 0

if __name__ == "__main__":

    raise SystemExit(main())