#!/usr/bin/env python3
"""Aggregate per-index far-hit probe counters by stratum, LOD and direction."""
from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l104/probes"


def aggregate(run: dict) -> dict[tuple[int, str], dict[str, int]]:
    result: dict[tuple[int, str], dict[str, int]] = defaultdict(
        lambda: {"queries": 0, "candidate_tests": 0, "hit_queries": 0})
    for index in run["stats"]:
        for direction in ("down", "up"):
            row = result[(index["lod"], direction)]
            for target, source in (("queries", "queries"),
                                   ("candidate_tests", "candidate_tests"),
                                   ("hit_queries", "hit_queries")):
                row[target] += index[f"{direction}_{source}"]
    return result


def main() -> int:
    runs = json.loads((OUT / "runs.json").read_text(encoding="utf-8"))
    lookup = {(run["stratum"], run["cache"]): aggregate(run) for run in runs}
    summary = []
    for stratum in ("SW", "FI", "TN", "VS"):
        before, after = lookup[(stratum, "off")], lookup[(stratum, "on")]
        for lod, direction in sorted(set(before) | set(after)):
            a = before.get((lod, direction), {"queries": 0, "candidate_tests": 0,
                                              "hit_queries": 0})
            b = after.get((lod, direction), {"queries": 0, "candidate_tests": 0,
                                             "hit_queries": 0})
            if a["queries"] != b["queries"] or a["hit_queries"] != b["hit_queries"]:
                raise RuntimeError(f"query/result drift for {stratum} LOD{lod} {direction}: {a} != {b}")
            summary.append({
                "stratum": stratum, "lod": lod, "direction": direction,
                "queries": a["queries"], "hit_queries": a["hit_queries"],
                "baseline_candidate_tests": a["candidate_tests"],
                "candidate_tests_after": b["candidate_tests"],
                "candidate_test_reduction_pct": 100 * (
                    1 - b["candidate_tests"] / a["candidate_tests"])
                    if a["candidate_tests"] else 0.0,
                "baseline_tests_per_query": a["candidate_tests"] / a["queries"]
                    if a["queries"] else 0.0,
                "tests_per_query_after": b["candidate_tests"] / b["queries"]
                    if b["queries"] else 0.0,
            })
    output = {"query_and_hit_results_identical": True, "summary": summary}
    (OUT / "summary.json").write_text(json.dumps(output, indent=2) + "\n", encoding="utf-8")
    print("[L104ProbeStatsSummary] " + json.dumps(output, sort_keys=True), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())