#!/usr/bin/env python3
"""gridfit22p.py — sweep new-arch env params over corpus, report TOTALs."""
import subprocess, re, sys, itertools, json

BASE = {
    "RT_LUT_OFF": "1", "RT_IIR12": "0", "RT_NOWARP": "1",
    "RT_NOBLEND": "1", "RT_NOIIR3": "1", "RT_ENV": "live",
}

def run(pool, scale, floor):
    env = {**BASE, "RT_POOL": str(pool), "RT_SCALE_M": str(scale), "RT_FLOOR": str(floor)}
    import os
    e = {**os.environ, **env}
    r = subprocess.run(["python3", "scripts/corpus_structural.py"],
                       capture_output=True, text=True, env=e, cwd="/home/m/re-tools")
    m = re.search(r"TOTAL\s+62\s+([\d.]+)\s+([\d.]+)", r.stdout)
    groups = dict(re.findall(r"(t1kq|t1k|al|res|dual|comb)\s+\d+\s+([\d.]+)", r.stdout))
    if not m:
        return None, None, groups
    return float(m.group(1)), float(m.group(2)), groups

if __name__ == "__main__":
    grid = []
    scales = [1.05, 1.10, 1.144]
    pools = [0] if len(sys.argv) < 2 else None
    results = []
    for pool, scale, floor in itertools.product([0], scales, [1]):
        tot, mx, g = run(pool, scale, floor)
        if tot is None:
            print(f"pool={pool} scale={scale} floor={floor}: FAILED"); continue
        print(f"pool={pool} scale={scale:.3f} floor={floor}: TOTAL={tot:.3f} max={mx:.2f}  {g}")
        results.append((tot, pool, scale, floor, g))
    results.sort()
    best = results[0]
    print("\nBEST: TOTAL=%.3f pool=%d scale=%.3f floor=%d" % (best[0], best[1], best[2], best[3]))
    # second stage around best scale with pool variants
    _, bp, bs, bf, _ = best
    for pool in [3, 5]:
        for ds in [-0.04, 0.0, 0.04]:
            sc = round(bs + ds, 3)
            tot, mx, g = run(pool, sc, bf)
            if tot is None: continue
            print(f"[stage2] pool={pool} scale={sc:.3f}: TOTAL={tot:.3f}  {g}")
            results.append((tot, pool, sc, bf, g))
    results.sort()
    b = results[0]
    json.dump({"total": b[0], "pool": b[1], "scale": b[2], "floor": b[3], "groups": b[4]},
              open("/tmp/opencode/gridfit_best.json", "w"), indent=1)
    print("\nFINAL BEST: TOTAL=%.3f pool=%d scale=%.3f floor=%d" % (b[0], b[1], b[2], b[3]))