#!/usr/bin/env python3
"""Sums the raw swing histograms eval.c's PersistMarginSafetyHistogramsToFile
writes (one file per shard of a sharded CALIBRATE_MARGIN_SAFETY run)
bucket-by-bucket and bin-by-bin, then computes EXACT combined mean/
percentiles from the summed bins -- percentiles don't combine validly
shard-by-shard (each shard's own in-process percentile is over a
different, non-representative sub-sample), only the raw counts do.

File format (one line per tier x material-bucket):
    <tier> <bucket> <count> <exceeded> <sum> <sumSq> <max> <bin0> ... <bin299>
SWING_HIST_BIN_WIDTH=10, SWING_HIST_BINS=300 (covers [0, 3000) cp) --
read from each file's own header comment rather than hardcoded, in
case a future build changes them.

Usage:
    python3 aggregate_marginsafety.py /path/to/calib_out/hist_*.txt
"""
import argparse
import glob
import sys
from collections import defaultdict


def parse_header(path):
    with open(path) as f:
        for line in f:
            if line.startswith("# SWING_HIST_BINS="):
                parts = dict(kv.split("=") for kv in line[2:].split())
                return int(parts["SWING_HIST_BINS"]), int(parts["SWING_HIST_BIN_WIDTH"])
    sys.exit(f"{path}: no SWING_HIST_BINS header found")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                  formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("files", nargs="+", help="hist_*.txt files (glob already expanded by shell, "
                                               "or pass a glob pattern quoted)")
    args = ap.parse_args()

    files = []
    for pat in args.files:
        matched = glob.glob(pat)
        files.extend(matched if matched else [pat])
    if not files:
        sys.exit("no input files")

    n_bins, bin_width = parse_header(files[0])

    # (tier, bucket) -> dict(count, exceeded, sum, sumSq, max, bins=[n_bins])
    agg = defaultdict(lambda: {"count": 0, "exceeded": 0, "sum": 0, "sumSq": 0,
                                "max": 0, "bins": [0] * n_bins})

    for path in files:
        with open(path) as f:
            for line in f:
                if line.startswith("#") or not line.strip():
                    continue
                parts = line.split()
                tier, bucket = parts[0], int(parts[1])
                count, exceeded, s, sumsq, mx = (int(x) for x in parts[2:7])
                bins = [int(x) for x in parts[7:7 + n_bins]]
                a = agg[(tier, bucket)]
                a["count"] += count
                a["exceeded"] += exceeded
                a["sum"] += s
                a["sumSq"] += sumsq
                a["max"] = max(a["max"], mx)
                for i, v in enumerate(bins):
                    a["bins"][i] += v

    def percentile(bins, count, p):
        """Exact percentile from binned counts -- rank = p/100 * count,
        walk cumulative counts to find which bin holds that rank, report
        the bin's lower edge (a slight underestimate, bin_width=10cp
        granularity -- good enough for choosing a margin, not claimed
        exact to the cp)."""
        if count == 0:
            return None
        rank = (p / 100.0) * count
        cum = 0
        for i, c in enumerate(bins):
            cum += c
            if cum >= rank:
                return i * bin_width
        return (len(bins) - 1) * bin_width  # fell off the end -- all mass in last bin

    tiers = sorted(set(t for t, b in agg.keys()))
    print(f"Aggregated {len(files)} shard files, {n_bins} bins x {bin_width}cp "
          f"(covers [0, {n_bins * bin_width}) cp)\n")

    results = {}
    for tier in tiers:
        print(f"=== {tier} ===")
        print(f"{'bucket':>6} {'n':>10} {'mean':>8} {'p50':>6} {'p98':>6} "
              f"{'p99.9':>7} {'max':>7} {'exceeded':>9}")
        for bucket in range(8):
            a = agg.get((tier, bucket))
            if a is None or a["count"] == 0:
                print(f"{bucket:>6} (no data)")
                continue
            mean = a["sum"] / a["count"]
            p50 = percentile(a["bins"], a["count"], 50)
            p98 = percentile(a["bins"], a["count"], 98)
            p999 = percentile(a["bins"], a["count"], 99.9)
            results[(tier, bucket)] = {"p98": p98, "p999": p999, "mean": mean,
                                        "max": a["max"], "n": a["count"]}
            print(f"{bucket:>6} {a['count']:>10} {mean:>8.1f} {p50:>6} {p98:>6} "
                  f"{p999:>7} {a['max']:>7} {a['exceeded']:>9}")
        print()

    if "nnue_lazy" in tiers:
        p98s = [results.get(("nnue_lazy", b), {}).get("p98", "?") for b in range(8)]
        print(f"Suggested iNnueLazyMarginByMatBucket (p98 per bucket, matching "
              f"the existing table's own documented choice of statistic):")
        print(f"    {{{', '.join(str(x) for x in p98s)}}}")
    if "nnue_superlazy" in tiers:
        p98s = [results.get(("nnue_superlazy", b), {}).get("p98", "?") for b in range(8)]
        print(f"\nFor comparison, NNUE super-lazy (material-only) p98 per bucket:")
        print(f"    {{{', '.join(str(x) for x in p98s)}}}")


if __name__ == "__main__":
    main()