"""Revision point 3 (2026-09-22): ONE scoring criterion for the whole paper.

The protocol tables were scored against the published root (|x - root| <= 1e-6), which charges
every method the full budget on six problems where f vanishes identically near the root
(section 7.3).  Section 9 uses the bracketing criterion: a sign change of f within 1e-6 of the
answer, or f(x) == 0 (corpus_roots_jl.solved).  The runs did not keep x, so they cannot be
re-scored from the files; every distinct configuration of every protocol sweep is run again
here, on APS and CHA, under the bracketing criterion, and written in the SAME row format
(label, spec, APS/CHA: total, solved, n, over, per[pid] = [n, ok]) so that every table and
check can switch files without changing code.

Built-in check: A', the base rule and the library rivals must reproduce the APS and CHA
lines of Table 10 exactly (those runs already use this criterion).

    python -u rescore_uniform.py                 all configurations -> sweep_uniform.json
    python -u rescore_uniform.py --smoke         12 configurations, local code check
    python -u rescore_uniform.py --workers N
"""
from __future__ import annotations

import glob
import json
import math
import os
import random
import sys
import time
from multiprocessing import get_context

os.environ.setdefault("OMP_NUM_THREADS", "1")

from corpus_roots_jl import solved
from sweep_full import MAXFEV, XTOL, label, load_problems, make_method

OUT = "sweep_uniform.json"


def sources():
    """Every row-format protocol sweep and the distinct specs in it (JSON-encoded)."""
    files = {}
    for p in sorted(glob.glob("sweep_*.json")):
        if "smoke" in p or p == OUT:
            continue
        with open(p, encoding="utf8") as fh:
            d = json.load(fh)
        if isinstance(d, list) and d and "spec" in d[0]:
            files[p] = [json.dumps(r["spec"]) for r in d]
    return files


def evaluate(spec_json):
    import scipy.optimize._tstutils as T
    spec = tuple(json.loads(spec_json))
    method = make_method(spec)
    res = {"label": label(spec), "spec": list(spec)}
    for setname in ("APS", "CHA"):
        total = nsolved = over = n_probs = 0
        per = {}
        for s, pid, f, a, b, _root in load_problems():
            if s != setname:
                continue
            n_probs += 1
            T._f6_cache.clear()
            random.seed(20260921)
            try:
                x, n, conv = method(f, a, b, XTOL, MAXFEV)
                x = float(x)
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            # no reset here: the random step function is checked on the realisation the
            # method met, exactly as in sweep_fourth.evaluate (Table 10)
            ok = bool(conv and solved(f, x))
            n = int(n)
            total += n if ok else MAXFEV
            per[pid] = [n, ok]
            if ok:
                nsolved += 1
                if n > 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4:
                    over += 1
        res[setname] = {"total": total, "solved": nsolved, "n": n_probs, "over": over,
                        "per": per}
    return spec_json, res


_CACHE = {}


def load_uniform(path):
    """The rows of protocol sweep `path`, re-scored under the bracketing criterion.  The
    15 MB file is read once per process; callers get fresh lists of the shared rows."""
    if "U" not in _CACHE:
        with open(OUT, encoding="utf8") as fh:
            _CACHE["U"] = json.load(fh)
    U = _CACHE["U"]
    return [U["by_spec"][s] for s in U["files"][path]]


def main():
    workers = int(sys.argv[sys.argv.index("--workers") + 1]) if "--workers" in sys.argv else 6
    files = sources()
    specs = sorted({s for v in files.values() for s in v})
    if "--smoke" in sys.argv:
        must = [json.dumps(list(x)) for x in (
            ("base", "chandrupatla"), ("base", "brent"), ("base", "toms748"),
            ("pure", "ratio", "derived3", "project"),
            ("pure", "ratio", "derived3", "project", "x",
             "contra+flatanalogmempt+exp+half+close+sub", "recent", "distance", "none",
             "pattern"))]
        specs = [s for s in must if s in specs] + [s for s in specs if s not in must][:7]
    print(f"{len(files)} sweep files, {sum(map(len, files.values()))} rows, "
          f"{len(specs)} distinct configurations, {workers} workers", flush=True)
    t0 = time.time()
    by_spec = {}
    with get_context("spawn").Pool(workers, maxtasksperchild=40) as pool:
        for i, (s, r) in enumerate(pool.imap_unordered(evaluate, specs, chunksize=2), 1):
            by_spec[s] = r
            if i % 25 == 0 or i == len(specs):
                el = time.time() - t0
                print(f"  {i}/{len(specs)} done, {el:.0f} s, ETA {el / i * (len(specs) - i):.0f} s",
                      flush=True)
    files = {p: v for p, v in files.items() if all(s in by_spec for s in v)}
    with open(OUT if "--smoke" not in sys.argv else "sweep_uniform_smoke.json", "w",
              encoding="utf8") as fh:
        json.dump({"criterion": "sign change within 1e-6 or f(x) == 0",
                   "files": files, "by_spec": by_spec}, fh)
    # the built-in check against Table 10 (APS, CHA lines, same criterion)
    want = {'["pure", "ratio", "derived3", "project", "x", "contra+flatanalogmempt+exp+half+close'
            '+sub", "recent", "distance", "none", "pattern"]': (2259, 1074),
            '["base", "chandrupatla"]': (2595, 1505), '["base", "brent"]': (2707, 3238),
            '["base", "toms748"]': (2636, 3088)}
    for s, (aps, cha) in want.items():
        if s in by_spec:
            got = (by_spec[s]["APS"]["total"], by_spec[s]["CHA"]["total"])
            print(f"  check vs Table 10: {by_spec[s]['label']:<60} {got} want {(aps, cha)} "
                  + ("OK" if got == (aps, cha) else "MISMATCH"), flush=True)
    print("rescore done", flush=True)


if __name__ == "__main__":
    main()