# Plays many seeded bot matches at once (the game's --balance runner, one process per seed) and adds up their tallies,
# flagging any weapon outside 45-55% win rate. Seeds start at SEED_BASE (default 100), so a run is repeatable and a
# fresh base checks results don't just fit one set of seeds.
#
# python tools/balance_runs.py [matches per process] [processes] [-- extra game arguments, like race]

import math
import os
import pathlib
import re
import subprocess
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor

ROOT = pathlib.Path(__file__).resolve().parent.parent
GAME = ROOT / "build" / "smack-a-brother"
TARGET = 0.05


def run(seed, matches, extra, prefs):
    env = dict(os.environ, RPG_SEED=str(seed), SDL_VIDEO_DRIVER="dummy", SDL_AUDIO_DRIVER="none", XDG_DATA_HOME=prefs)
    result = subprocess.run([str(GAME), "--balance", str(matches)] + extra, env=env, capture_output=True, text=True)
    return result.stdout + result.stderr


def main():
    args = sys.argv[1:]
    extra = args[args.index("--") + 1:] if "--" in args else []
    args = args[:args.index("--")] if "--" in args else args
    matches = int(args[0]) if args else 2500
    processes = int(args[1]) if len(args) > 1 else os.cpu_count()
    base = int(os.environ.get("SEED_BASE", "100"))
    with tempfile.TemporaryDirectory() as prefs, ThreadPoolExecutor(processes) as pool:
        outputs = list(pool.map(lambda i: run(base + i, matches, extra, prefs), range(processes)))

    tallies = {"Weapons": {}, "Trinkets": {}}
    total = stalls = 0
    for output in outputs:
        section = None
        for line in output.splitlines():
            if line.startswith("Weapons"):
                section = "Weapons"
            elif line.startswith("Trinkets"):
                section = "Trinkets"
            row = re.match(r"\s+(.+?)\s+[\d.]+% \+-\s*[\d.]+\s+\((\d+) wins / (\d+) picks, (\d+) stalls\)", line)
            if row and section:
                tally = tallies[section].setdefault(row.group(1), [0, 0, 0])
                for i in range(3):
                    tally[i] += int(row.group(2 + i))
            summary = re.search(r"(\d+) matches in .*?(\d+) hit the", line)
            if summary:
                total += int(summary.group(1))
                stalls += int(summary.group(2))

    print(f"{total} matches, {stalls} stalled ({100 * stalls / max(total, 1):.1f}%), seeds {base}-{base + processes - 1}")
    for section, tally in tallies.items():
        print(section)
        for name, (wins, picks, stalled) in sorted(tally.items(), key=lambda item: -item[1][0] / max(item[1][1], 1)):
            rate = wins / max(picks, 1)
            margin = 1.96 * math.sqrt(rate * (1 - rate) / max(picks, 1))
            flag = "  <--" if section == "Weapons" and abs(rate - 0.5) > TARGET else ""
            print(f"  {name:16} {100 * rate:5.1f}% +-{100 * margin:4.1f}  ({picks} picks, {stalled} stalls){flag}")


if __name__ == "__main__":
    main()