"""Recompute the source-scale summary from committed raw trajectory arrays."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np


ROOT = Path(__file__).resolve().parents[2]


def read_target(target: str) -> dict[str, object]:
    path = ROOT / "outputs" / f"authored_2d_{target}.executed.ipynb"
    raw_path = ROOT / "outputs" / f"authored_2d_{target}.raw.npz"
    if not raw_path.is_file():
        raise RuntimeError(f"{target}: missing raw source trajectory artifact: {raw_path}")
    with np.load(raw_path) as raw:
        expected_shapes = {
            "KL_bw_runs": (100, 201),
            "KL_cbo_runs": (100, 201),
            "KL_svgd_runs": (100, 801),
            "KL_fr_runs": (100, 401),
        }
        raw_shapes = {name: list(raw[name].shape) for name in expected_shapes if name in raw.files}
        if raw_shapes != {name: list(shape) for name, shape in expected_shapes.items()} or not all(
            np.isfinite(raw[name]).all() for name in expected_shapes
        ):
            raise RuntimeError(f"{target}: malformed/non-finite raw trajectory arrays: {raw_shapes}")
        values = {
            "ntests": expected_shapes["KL_cbo_runs"][0],
            "cbo_start": float(np.median(raw["KL_cbo_runs"][:, 0])),
            "cbo_final": float(np.median(raw["KL_cbo_runs"][:, -1])),
            "bw_start": float(np.median(raw["KL_bw_runs"][:, 0])),
            "bw_final": float(np.median(raw["KL_bw_runs"][:, -1])),
            "svgd_start": float(np.median(raw["KL_svgd_runs"][:, 0])),
            "svgd_final": float(np.median(raw["KL_svgd_runs"][:, -1])),
            "fr_start": float(np.median(raw["KL_fr_runs"][:, 0])),
            "fr_final": float(np.median(raw["KL_fr_runs"][:, -1])),
        }
    result: dict[str, object] = {
        "notebook": str(path.relative_to(ROOT)) if path.is_file() else None,
        "raw_trajectories": str(raw_path.relative_to(ROOT)),
        "raw_shapes": raw_shapes,
        "raw_values_finite": True,
    }
    result.update(values)
    result["cbo_beats_bw_final"] = bool(result["cbo_final"] < result["bw_final"])
    result["all_four_baselines_present"] = True
    return result


def main() -> None:
    targets = {target: read_target(target) for target in "ABCD"}
    passed = all(
        row["ntests"] == 100 and row["all_four_baselines_present"] and row["cbo_beats_bw_final"]
        for row in targets.values()
    )
    output = {
        "protocol": "committed arrays generated by author experiment_2D.ipynb cells 0--5; summary recomputed from raw values",
        "targets": targets,
        "c5_cbo_beats_bw_on_all_targets": passed,
        "scope": "Median final-KL comparison only; this does not claim universal superiority over every baseline or target.",
    }
    (ROOT / "outputs" / "authored_2d_summary.json").write_text(json.dumps(output, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(output, indent=2))
    if not passed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()