Blame for autojev/evaluate.py

"""Accuracy and calibration on fixed Choice/Noul examples, with paired comparisons."""

import argparse
from collections import defaultdict
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime, timezone
import hashlib
import json
import math
import os
from pathlib import Path
import time
from typing import NotRequired, TypedDict, cast

import httpx
import numpy as np
from numpy.typing import NDArray

from autojev.types import Example, JSONValue, Label, Question


ECE_BINS = 15
CALIBRATION_MARGIN = 0.010
JEV_ENDPOINT = "https://openrouter.ai/api/alpha/decisions"
type FloatArray = NDArray[np.float64]
type Infer = Callable[[Sequence[Example]], Sequence[Sequence[float]]]
type UsageCallback = Callable[[Mapping[str, JSONValue]], None]


class Prediction(TypedDict):
    id: str
    suite: str
    dataset: str
    family: str
    label: Label
    options: list[Label]
    probabilities: list[float]
    prediction: Label
    correct: bool
    confidence: float
    temperature: float
    logits: NotRequired[list[float]]
    soft_target: NotRequired[list[float]]
    answer: NotRequired[dict[str, JSONValue]]
    model: NotRequired[str]
    provider: NotRequired[str]
    request_id: NotRequired[str]
    latency_ms: NotRequired[float]
    normalization_total: NotRequired[float]
    created: NotRequired[str]
    request_sha256: NotRequired[str]


class ReliabilityBin(TypedDict):
    lower: float
    upper: float
    count: int
    confidence: float | None
    accuracy: float | None


class Metrics(TypedDict):
    count: int
    accuracy: float
    ece: float
    brier: float
    nll: float
    reliability: list[ReliabilityBin]
    zero_probability_count: int
    soft_count: NotRequired[int]
    soft_nll: NotRequired[float]
    soft_brier: NotRequired[float]


class Effect(TypedDict):
    local: float
    reference: float
    difference: float
    difference_ci95: list[float]


class Comparison(TypedDict):
    count: int
    families: int
    effects: dict[str, Effect]
    by_suite: dict[str, dict[str, Effect]]
    empirical_target_met: bool
    statistically_supported: bool
    calibration_margin: float
    bootstrap_replicates: int
    bootstrap_seed: int


def options(question: Question) -> list[Label]:
    if question["type"] == "noul":
        return [False, True]
    if question["type"] == "choice":
        return list(question["criteria"])
    if question["type"] == "score":  # team-1: ordered rubric levels "0".."n-1", same as model.options
        return [str(i) for i in range(len(question["criteria"]))]
    raise ValueError("The accuracy/calibration experiment includes Choice, Noul and Score only")


def hard_label(row: Example) -> Label:
    """Soft SFT targets never silently replace the evaluation reference label."""
    label = row.get("label", row["target"])
    if not isinstance(label, (str, int, bool)):
        raise ValueError(f"An explicit hard label is required: {row['id']}")
    return label


def label_index(labels: Sequence[Label], label: Label) -> int:
    for index, value in enumerate(labels):
        if type(value) is type(label) and value == label:
            return index
    raise ValueError(f"Reference or prediction {label!r} is not an option")


def make_prediction(
    row: Example, probabilities: Sequence[float], *, prediction: Label | None = None,
    temperature: float = 1.0,
) -> Prediction:
    labels = options(row["question"])
    values = [float(value) for value in probabilities]
    if len(values) != len(labels) or any(not math.isfinite(p) or not 0 <= p <= 1 for p in values):
        raise ValueError(f"Invalid probability vector: {row['id']}")
    if not math.isclose(sum(values), 1, abs_tol=1e-8):
        raise ValueError(f"Probabilities must be normalized: {row['id']}")
    label = hard_label(row)
    label_index(labels, label)
    chosen = labels[max(range(len(labels)), key=values.__getitem__)] if prediction is None else prediction
    chosen_index = label_index(labels, chosen)
    family = row.get("family")
    dataset = row["source"].get("dataset", row["suite"])
    if not isinstance(family, str) or not family or not isinstance(dataset, str):
        raise ValueError(f"Missing dataset/family provenance: {row['id']}")
    result: Prediction = {
        "id": row["id"], "suite": row["suite"], "dataset": dataset, "family": family,
        "label": label, "options": labels, "probabilities": values, "prediction": chosen,
        "correct": type(chosen) is type(label) and chosen == label,
        "confidence": values[chosen_index], "temperature": temperature,
    }
    human = row["source"].get("human_distribution")
    soft: list[float] | None = None
    if isinstance(human, dict):
        keys = [str(value).lower() if isinstance(value, bool) else str(value) for value in labels]
        if set(human) != set(keys):
            raise ValueError(f"Human distribution does not match options: {row['id']}")
        soft = [float(cast(float, human[key])) for key in keys]
    elif isinstance(row["target"], list):
        soft = [float(value) for value in row["target"]]
    if soft is not None:
        if len(soft) != len(labels) or any(not math.isfinite(p) or not 0 <= p <= 1 for p in soft) or not math.isclose(sum(soft), 1, abs_tol=1e-6):
            raise ValueError(f"Invalid soft distribution: {row['id']}")
        result["soft_target"] = soft
    return result


def evaluate_logits(
    rows: Sequence[Example], logits: Sequence[Sequence[float]], temperature: float = 1.0,
) -> list[Prediction]:
    if len(rows) != len(logits) or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Logit rows must match examples and temperature must be positive")
    result: list[Prediction] = []
    for row, values in zip(rows, logits, strict=True):
        count = len(options(row["question"]))
        raw = np.asarray(values[:count], dtype=np.float64)
        if len(raw) != count or not np.isfinite(raw).all():
            raise ValueError(f"Missing or nonfinite valid logits: {row['id']}")
        shifted = (raw - raw.max()) / temperature
        probabilities = np.exp(shifted)
        probabilities /= probabilities.sum()
        prediction = make_prediction(row, probabilities.tolist(), temperature=temperature)
        prediction["logits"] = raw.tolist()
        result.append(prediction)
    validate_coverage(rows, result)
    return result


def predict_local(rows: Sequence[Example], infer: Infer, temperature: float = 1.0) -> list[Prediction]:
    """The model adapter returns raw logits in the supplied row/option order."""
    return evaluate_logits(rows, infer(rows), temperature)


def fit_temperature(logits: Sequence[Sequence[float]], target_indices: Sequence[int]) -> float:
    """Fit one scalar on the separate temperature fold using hard-label NLL."""
    if not logits or len(logits) != len(target_indices):
        raise ValueError("Temperature fitting needs nonempty matching logits and hard labels")
    width = max(map(len, logits))
    values = np.full((len(logits), width), -np.inf, dtype=np.float64)
    for index, (row, target) in enumerate(zip(logits, target_indices, strict=True)):
        if not row or not 0 <= target < len(row) or not np.isfinite(row).all():
            raise ValueError("Temperature fitting requires finite, unpadded valid logits")
        values[index, :len(row)] = row
    values -= values.max(axis=1, keepdims=True)
    target_logits = values[np.arange(len(values)), np.asarray(target_indices)]

    def loss(log_temperature: float) -> float:
        inverse = math.exp(-log_temperature)
        return float(np.mean(np.log(np.exp(values * inverse).sum(axis=1)) - target_logits * inverse))

    # A broad, fixed interval avoids selecting temperature bounds after seeing dev results.
    low, high = math.log(0.05), math.log(20.0)
    ratio = (math.sqrt(5) - 1) / 2
    left, right = high - ratio * (high - low), low + ratio * (high - low)
    left_loss, right_loss = loss(left), loss(right)
    for _ in range(80):
        if left_loss <= right_loss:
            high, right, right_loss = right, left, left_loss
            left = high - ratio * (high - low)
            left_loss = loss(left)
        else:
            low, left, left_loss = left, right, right_loss
            right = low + ratio * (high - low)
            right_loss = loss(right)
    candidates = [math.log(0.05), (low + high) / 2, math.log(20.0), 0.0]
    return math.exp(min(candidates, key=loss))


def metrics(rows: Sequence[Prediction]) -> Metrics:
    if not rows:
        raise ValueError("Cannot score an empty evaluation")
    bins: list[list[Prediction]] = [[] for _ in range(ECE_BINS)]
    brier, nll, soft_brier, soft_nll = [], [], [], []
    for row in rows:
        index = label_index(row["options"], row["label"])
        probabilities = row["probabilities"]
        nll.append(-math.log(max(probabilities[index], 1e-12)))
        brier.append(sum((p - int(i == index)) ** 2 for i, p in enumerate(probabilities)))
        bins[min(ECE_BINS - 1, int(row["confidence"] * ECE_BINS))].append(row)
        if "soft_target" in row:
            soft = row["soft_target"]
            soft_brier.append(sum((p - q) ** 2 for p, q in zip(probabilities, soft, strict=True)))
            soft_nll.append(-sum(q * math.log(max(p, 1e-12)) for p, q in zip(probabilities, soft, strict=True)))
    reliability: list[ReliabilityBin] = []
    ece = 0.0
    for index, group in enumerate(bins):
        confidence = sum(row["confidence"] for row in group) / len(group) if group else None
        accuracy = sum(row["correct"] for row in group) / len(group) if group else None
        if confidence is not None and accuracy is not None:
            ece += len(group) / len(rows) * abs(confidence - accuracy)
        reliability.append({"lower": index / ECE_BINS, "upper": (index + 1) / ECE_BINS,
                            "count": len(group), "confidence": confidence, "accuracy": accuracy})
    result: Metrics = {
        "count": len(rows), "accuracy": sum(row["correct"] for row in rows) / len(rows),
        "ece": ece, "brier": sum(brier) / len(rows), "nll": sum(nll) / len(rows),
        "reliability": reliability,
        "zero_probability_count": sum(p == 0 for row in rows for p in row["probabilities"]),
    }
    if soft_nll:
        result.update({"soft_count": len(soft_nll), "soft_nll": sum(soft_nll) / len(soft_nll), "soft_brier": sum(soft_brier) / len(soft_brier)})
    return result


def calibration_ok(local: Metrics, reference: Metrics, margin: float = CALIBRATION_MARGIN) -> bool:
    return local["ece"] <= reference["ece"] + margin and local["brier"] <= reference["brier"] + margin


def selection_key(summary: Metrics, step: int) -> tuple[float, float, int]:
    """Minimize this key among checkpoints passing calibration_ok."""
    return -summary["accuracy"], summary["brier"], step


def validate_coverage(rows: Sequence[Example], predictions: Sequence[Prediction]) -> None:
    expected = {row["id"]: row for row in rows}
    actual = {row["id"]: row for row in predictions}
    if len(expected) != len(rows) or len(actual) != len(predictions) or expected.keys() != actual.keys():
        raise ValueError("Require exactly one successful prediction for every fixed evaluation ID")
    for identifier, row in expected.items():
        saved = actual[identifier]
        rebuilt = make_prediction(row, saved["probabilities"], prediction=saved["prediction"], temperature=saved["temperature"])
        for field in ("suite", "dataset", "family", "label", "options", "correct", "confidence"):
            if saved[field] != rebuilt[field]:
                raise ValueError(f"Saved prediction differs from its evaluation example: {identifier}/{field}")


def _statistics(row: Prediction) -> FloatArray:
    """Additive sufficient statistics allow exact ECE recomputation per bootstrap draw."""
    output = np.zeros(4 + 3 * ECE_BINS, dtype=np.float64)
    index = label_index(row["options"], row["label"])
    probabilities = row["probabilities"]
    output[:4] = [1, int(row["correct"]), sum((p - int(i == index)) ** 2 for i, p in enumerate(probabilities)),
                  -math.log(max(probabilities[index], 1e-12))]
    bucket = min(ECE_BINS - 1, int(row["confidence"] * ECE_BINS))
    output[4 + 3 * bucket:7 + 3 * bucket] = [1, row["confidence"], int(row["correct"])]
    return output


def _stat_metrics(values: FloatArray) -> dict[str, FloatArray]:
    return {"accuracy": values[..., 1] / values[..., 0], "brier": values[..., 2] / values[..., 0],
            "nll": values[..., 3] / values[..., 0],
            "ece": np.abs(values[..., 5::3] - values[..., 6::3]).sum(axis=-1) / values[..., 0]}


def compare(
    local: Sequence[Prediction], reference: Sequence[Prediction], *, replicates: int = 10000,
    seed: int = 20260920, margin: float = CALIBRATION_MARGIN,
) -> Comparison:
    """Paired family bootstrap preserves each suite's fixed contribution to the headline."""
    if replicates < 1000:
        raise ValueError("Use at least 1,000 paired bootstrap draws")
    left, right = {row["id"]: row for row in local}, {row["id"]: row for row in reference}
    if not left or len(left) != len(local) or len(right) != len(reference) or left.keys() != right.keys():
        raise ValueError("Comparison requires complete unique matching prediction IDs")
    suites = sorted({row["suite"] for row in local})
    suite_index = {suite: index for index, suite in enumerate(suites)}
    counts = np.asarray([sum(row["suite"] == suite for row in local) for suite in suites], dtype=np.float64)
    families: dict[tuple[str, str], dict[int, FloatArray]] = {}
    for identifier, row in left.items():
        other = right[identifier]
        if any(row[field] != other[field] for field in ("suite", "dataset", "family", "label", "options")):
            raise ValueError(f"Paired prediction provenance differs: {identifier}")
        group = families.setdefault((row["dataset"], row["family"]), {})
        index = suite_index[row["suite"]]
        group.setdefault(index, np.zeros((2, 4 + 3 * ECE_BINS), dtype=np.float64))
        group[index] += np.stack([_statistics(row), _statistics(other)])
    strata: defaultdict[tuple[int, ...], list[dict[int, FloatArray]]] = defaultdict(list)
    for key in sorted(families):
        family = families[key]
        strata[tuple(sorted(family))].append(family)
    rng = np.random.default_rng(seed)
    samples = np.zeros((replicates, len(suites), 2, 4 + 3 * ECE_BINS), dtype=np.float64)
    for membership in sorted(strata):
        clusters = strata[membership]
        values = np.stack([np.stack([cluster[index] for index in membership]) for cluster in clusters])
        flattened = values.reshape(len(clusters), -1)
        for start in range(0, replicates, 128):
            size = min(128, replicates - start)
            draws = rng.multinomial(len(clusters), np.full(len(clusters), 1 / len(clusters)), size=size)
            totals = (draws @ flattened).reshape(size, len(membership), 2, -1)
            for position, index in enumerate(membership):
                samples[start:start + size, index] += totals[:, position]
    # Whole-family draws vary record counts. Retain the predeclared task mixture.
    weights = counts[None, :, None, None] / samples[..., :1]
    pooled = (samples * weights).sum(axis=1)
    pooled_metrics = _stat_metrics(pooled)
    suite_metrics = _stat_metrics(samples)

    def effects(a: Sequence[Prediction], b: Sequence[Prediction], draws: Mapping[str, FloatArray]) -> dict[str, Effect]:
        a_metrics, b_metrics = metrics(a), metrics(b)
        result: dict[str, Effect] = {}
        for name in ("accuracy", "ece", "brier", "nll"):
            a_value = a_metrics[name]
            b_value = b_metrics[name]
            difference = draws[name][..., 0] - draws[name][..., 1]
            result[name] = {"local": a_value, "reference": b_value, "difference": a_value - b_value,
                            "difference_ci95": np.quantile(difference, [0.025, 0.975]).tolist()}
        return result

    overall = effects(local, reference, pooled_metrics)
    by_suite = {suite: effects([r for r in local if r["suite"] == suite], [r for r in reference if r["suite"] == suite],
                              {name: values[:, index] for name, values in suite_metrics.items()})
                for index, suite in enumerate(suites)}
    return {
        "count": len(local), "families": len(families), "effects": overall, "by_suite": by_suite,
        "empirical_target_met": overall["accuracy"]["difference"] > 0 and all(overall[name]["difference"] <= margin for name in ("ece", "brier")),
        "statistically_supported": overall["accuracy"]["difference_ci95"][0] > 0 and all(overall[name]["difference_ci95"][1] <= margin for name in ("ece", "brier")),
        "calibration_margin": margin, "bootstrap_replicates": replicates, "bootstrap_seed": seed,
    }


def remote_prediction(row: Example, client: httpx.Client, *, model: str, url: str, token: str,
                      usage_callback: UsageCallback | None = None) -> Prediction:
    if row.get("images"):
        raise ValueError("The paired Jev comparison is text only")
    labels = options(row["question"])
    key = row["source"].get("question_key", "q")
    if not isinstance(key, str):
        raise ValueError("Question key must be a string")
    payload = {"model": model, "state": row["state"], "questions": {key: row["question"]}}
    started = time.perf_counter()
    response: httpx.Response | None = None
    for attempt in range(3):
        response = client.post(url, json=payload, headers={"Authorization": f"Bearer {token}"})
        if response.status_code not in (429, 500, 502, 503, 504, 529):
            break
        if attempt < 2:
            time.sleep(2 ** attempt)
    assert response is not None
    response.raise_for_status()
    body = cast(dict[str, JSONValue], response.json())
    if usage_callback is not None:
        usage_callback(body)
    answers = cast(dict[str, JSONValue], body["answers"])
    answer = cast(dict[str, JSONValue], answers[key])
    if answer["type"] != row["question"]["type"]:
        raise ValueError(f"Response question type differs: {row['id']}")
    chosen: Label
    if row["question"]["type"] == "noul":
        probability = float(cast(float, answer["noul"]))
        values = [1 - probability, probability]
        chosen = probability > 0.5
    else:
        distribution = cast(dict[str, float], answer["probabilities"])
        if set(distribution) != set(labels):
            raise ValueError(f"Response options differ: {row['id']}")
        values = [distribution[str(label)] for label in labels]
        chosen = cast(str, answer["choice"])
    if any(not math.isfinite(value) or not 0 <= value <= 1 for value in values) or abs(sum(values) - 1) > 0.025:
        raise ValueError(f"Invalid provider probabilities: {row['id']}")
    total = sum(values)
    result = make_prediction(row, [value / total for value in values], prediction=chosen)
    result.update({"answer": answer, "latency_ms": (time.perf_counter() - started) * 1000, "normalization_total": total})
    result["created"] = datetime.now(timezone.utc).isoformat()
    result["request_sha256"] = hashlib.sha256(json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode()).hexdigest()
    for field in ("model", "provider"):
        if isinstance(body.get(field), str):
            result[field] = cast(str, body[field])
    if isinstance(body.get("id"), str):
        result["request_id"] = cast(str, body["id"])
    return result


def read_rows(path: Path) -> list[Example]:
    with path.open(encoding="utf-8") as stream:
        return [cast(Example, json.loads(line)) for line in stream if line.strip()]


def read_predictions(path: Path) -> list[Prediction]:
    with path.open(encoding="utf-8") as stream:
        return [cast(Prediction, json.loads(line)) for line in stream if line.strip()]


def write_json(path: Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n")


class Arguments(argparse.Namespace):
    data: Path
    output: Path
    predictions: Path | None
    reference: Path | None
    jev: bool
    local: bool
    checkpoint: Path | None
    batch_size: int
    temperature: float | None
    model: str
    resume: bool


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True, help="JSON summary; predictions use the same stem")
    parser.add_argument("--predictions", type=Path, help="Recompute saved predictions without inference")
    parser.add_argument("--reference", type=Path, help="Paired Jev predictions")
    parser.add_argument("--jev", action="store_true", help="Query Jev's native OpenRouter decisions endpoint")
    parser.add_argument("--local", action="store_true", help="Evaluate the pinned Qwen base or a local checkpoint")
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--temperature", type=float, help="Override the checkpoint temperature; use 1 for raw results")
    parser.add_argument("--model", default="typesafe/jev-1.13")
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args(namespace=Arguments())
    if sum((args.jev, args.local, args.predictions is not None)) != 1:
        parser.error("Choose exactly one of --jev, --local or --predictions")
    if args.batch_size < 1:
        parser.error("--batch-size must be positive")
    if not args.local and (args.checkpoint is not None or args.temperature is not None):
        parser.error("--checkpoint and --temperature require --local")
    if args.resume and not args.jev:
        parser.error("--resume applies to interrupted remote evaluations")
    rows = read_rows(args.data)
    identity = {"dataset_sha256": hashlib.sha256(args.data.read_bytes()).hexdigest(), "model": args.model,
                "endpoint": JEV_ENDPOINT, "ece_bins": ECE_BINS,
                "evaluator_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
    if args.predictions is not None:
        predictions = read_predictions(args.predictions)
    elif args.local:
        import torch
        from autojev.model import DecisionModel

        destination = args.output.with_suffix(".predictions.jsonl")
        if destination.exists():
            raise FileExistsError(f"Refusing to overwrite an evaluation: {destination}")
        model = DecisionModel(checkpoint=args.checkpoint)
        scale = model.temperature if args.temperature is None else args.temperature

        @torch.inference_mode()
        def infer(batch_rows: Sequence[Example]) -> list[list[float]]:
            output: list[list[float]] = []
            for start in range(0, len(batch_rows), args.batch_size):
                batch = model.prepare(batch_rows[start:start + args.batch_size])
                logits: list[list[float]] = model(batch).cpu().tolist()
                output.extend(values[:count] for values, count in zip(logits, batch.counts, strict=True))
            return output

        predictions = predict_local(rows, infer, scale)
        identity.update({"model": model.base_model, "revision": model.revision, "temperature": scale,
                         "checkpoint": str(args.checkpoint.resolve()) if args.checkpoint else None, "endpoint": "local"})
        if args.checkpoint is not None:
            identity["checkpoint_config_sha256"] = hashlib.sha256((args.checkpoint / "decision_config.json").read_bytes()).hexdigest()
        write_json(args.output.with_suffix(".manifest.json"), identity)
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_text("".join(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n" for row in predictions))
    else:
        destination = args.output.with_suffix(".predictions.jsonl")
        manifest_path = args.output.with_suffix(".manifest.json")
        if args.resume and destination.exists():
            if not manifest_path.exists() or json.loads(manifest_path.read_text()) != identity:
                raise ValueError("Cannot resume predictions from a different dataset/model/protocol")
            predictions = read_predictions(destination)
            known = {row["id"] for row in predictions}
            validate_coverage([row for row in rows if row["id"] in known], predictions)
        else:
            if destination.exists():
                raise FileExistsError(f"Refusing to overwrite an evaluation: {destination}")
            predictions = []
            write_json(manifest_path, identity)
        done = {row["id"] for row in predictions}
        destination.parent.mkdir(parents=True, exist_ok=True)
        with httpx.Client(timeout=120) as client, destination.open("a") as stream:
            for row in rows:
                if row["id"] in done:
                    continue
                prediction = remote_prediction(row, client, model=args.model, url=JEV_ENDPOINT, token=os.environ["OPENROUTER_API_KEY"])
                stream.write(json.dumps(prediction, ensure_ascii=False, allow_nan=False) + "\n")
                stream.flush()
                predictions.append(prediction)
                if len(predictions) % 50 == 0:
                    print(f"Evaluated {len(predictions)}/{len(rows)}", flush=True)
    validate_coverage(rows, predictions)
    result: dict[str, object] = {
        "created": datetime.now(timezone.utc).isoformat(), "dataset_sha256": identity["dataset_sha256"],
        "inference": identity if args.predictions is None else {"predictions_sha256": hashlib.sha256(args.predictions.read_bytes()).hexdigest()},
        "overall": metrics(predictions), "by_suite": {suite: metrics([row for row in predictions if row["suite"] == suite])
                                                     for suite in sorted({row["suite"] for row in predictions})},
        "ece_bins": ECE_BINS, "labels": "hard reference labels; soft diagnostics are separate",
        "observed_models": sorted({row["model"] for row in predictions if "model" in row}),
        "observed_providers": sorted({row["provider"] for row in predictions if "provider" in row}),
    }
    if args.reference is not None:
        reference = read_predictions(args.reference)
        validate_coverage(rows, reference)
        result["comparison"] = compare(predictions, reference)
    write_json(args.output, result)
    print(json.dumps(result["overall"], indent=2))


if __name__ == "__main__":
    main()