"""Compare alternating fixed-work local runs without claiming a native gain."""

import argparse
import json
import math
from pathlib import Path
import re
import statistics


def load(directory):
    manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8-sig"))
    log = (directory / "probe.stdout.log").read_text(encoding="utf-8-sig")
    match = re.search(r"arm-ms=([0-9.]+) instructions=(\d+) journal-pages=128 nat-arp-replies=32", log)
    if not match or manifest["ProbeExitCode"] != 0 or int(match.group(2)) != 655360000:
        raise ValueError(f"Probe did not finish the exact work: {directory}")
    collection = manifest.get("Collection", "WPRSystemAndNative" if "Profile" in manifest else None)
    requested = collection != "NoTrace"
    if f"trace-requested={int(requested)} trace-enabled={int(requested)}" not in log:
        raise ValueError(f"Probe trace state does not match collection: {directory}")
    if collection == "WPRSystemAndNative":
        decoder = json.loads((directory / "decoder-summary.json").read_text(encoding="utf-8-sig"))
        if any(decoder.get(key) != 0 for key in
               ("status", "native_decode_errors", "events_lost", "buffers_lost")):
            raise ValueError(f"WPR trace has decode errors or lost events: {directory}")
    return {"directory": str(directory), "collection": collection,
            "hash": manifest["ProbeSHA256"], "pdb_hash": manifest["ProbePDBSHA256"],
            "affinity_mask": manifest.get("ProbeAffinityMask", 0),
            "observed_affinity_mask": manifest.get("ProbeObservedAffinityMask"),
            "arm_ms": float(match.group(1)), "wall_ms": 1000 * manifest["ProbeWallSeconds"],
            "cpu_ms": 1000 * manifest["ProbeCPUSeconds"]}


def upper_95(values):
    # Two-sided Student 95% upper bound; report but never equate with a
    # guarantee for a different (native Wii) workload.
    critical = {4: 3.182, 5: 2.776, 6: 2.571, 7: 2.447, 8: 2.365}
    return statistics.mean(values) + critical[len(values)] * statistics.stdev(values) / math.sqrt(len(values))


def compare(directories):
    runs = [load(path) for path in directories]
    on_mode = runs[1]["collection"] if len(runs) >= 2 else None
    if (on_mode not in ("PrivateUserMode", "WPRSystemAndNative") or
            len(runs) < 4 or len(runs) > 16 or len(runs) % 2 or
            any(r["collection"] != kind for r, kind in
                zip(runs, ("NoTrace", on_mode) * (len(runs) // 2)))):
        raise ValueError("Provide at least two OFF/ON pairs in alternating order")
    if len({(r["hash"], r["pdb_hash"]) for r in runs}) != 1:
        raise ValueError("The compared executables or PDBs differ")
    if len({r["affinity_mask"] for r in runs}) != 1 or any(
            r["affinity_mask"] and r["observed_affinity_mask"] != r["affinity_mask"]
            for r in runs):
        raise ValueError("Probe CPU affinity differs across compared runs")
    pairs = []
    for off, on in zip(runs[::2], runs[1::2]):
        pairs.append({"off": off["directory"], "on": on["directory"],
                      "arm_phase_percent": 100 * (on["arm_ms"] / off["arm_ms"] - 1),
                      "whole_probe_percent": 100 * (on["wall_ms"] / off["wall_ms"] - 1),
                      "cpu_percent": 100 * (on["cpu_ms"] / off["cpu_ms"] - 1),
                      "non_arm_wall_delta_ms": (on["wall_ms"] - on["arm_ms"]) -
                                               (off["wall_ms"] - off["arm_ms"])})
    arm_differences = [pair["arm_phase_percent"] for pair in pairs]
    wall_differences = [pair["whole_probe_percent"] for pair in pairs]
    cpu_differences = [pair["cpu_percent"] for pair in pairs]
    baseline_range = 100 * (max(r["arm_ms"] for r in runs[::2]) /
                            min(r["arm_ms"] for r in runs[::2]) - 1)
    wall_baseline_range = 100 * (max(r["wall_ms"] for r in runs[::2]) /
                                 min(r["wall_ms"] for r in runs[::2]) - 1)
    bounds = (dict(arm_phase_percent=upper_95(arm_differences),
                   whole_probe_percent=upper_95(wall_differences),
                   cpu_percent=upper_95(cpu_differences)) if len(pairs) >= 4 else None)
    qualified = (on_mode == "WPRSystemAndNative" and bounds is not None and
                 baseline_range <= 2 and wall_baseline_range <= 2 and
                 all(value <= 2 for value in bounds.values()))
    return {"schema": 1, "on_collection": on_mode, "runs": runs, "pairs": pairs,
            "mean_arm_phase_percent": statistics.mean(arm_differences),
            "off_arm_phase_range_percent": baseline_range,
            "off_wall_range_percent": wall_baseline_range,
            "upper_95_percent": bounds,
            "mean_non_arm_wall_delta_ms": statistics.mean(p["non_arm_wall_delta_ms"] for p in pairs),
            "two_percent_budget_qualified": qualified,
            "note": "Fixed-work local collector cost only, not a native Wii speed estimate. Four or more OFF/ON pairs, low baseline spread, and upper bounds below 2% are required for this conservative local gate."}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("captures", nargs="+", type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    result = compare(args.captures)
    with args.output.open("x", encoding="utf-8") as target:
        json.dump(result, target, indent=2)
    print(json.dumps({key: result[key] for key in
                      ("on_collection", "pairs", "mean_arm_phase_percent",
                       "off_arm_phase_range_percent", "off_wall_range_percent",
                       "upper_95_percent", "mean_non_arm_wall_delta_ms",
                       "two_percent_budget_qualified")}, indent=2))