"""PREREG_modab_2026-09-21: A' against ModAB (the authors' package pymodab) - on the paper's
protocol (published roots), on the six corpora of section 9 (uniform criterion), and on
ModAB's own benchmark (Roots.jl f01-f91: all 91 and the 78 out of the development sets).
A' and ModAB are frozen; nothing is selected by the result.

    python -u sweep_modab.py            run -> sweep_modab.json, then the verdicts
    python -u sweep_modab.py --report
"""
from __future__ import annotations

import json
import math
import sys
from multiprocessing import Pool

import sweep_fourth
from corpus_roots_jl import build as build_rjl
from corpus_roots_jl import solved
from sweep_full import APRIME, MAXFEV, XTOL
from sweep_full import evaluate as eval_protocol
from sweep_full import make_method
from sweep_oos import sign_p

sweep_fourth.CONFIGS["modab"] = ("base", "modab")
SETS = ("APS", "CHA", "NEW", "RJL", "APP", "D")
APRIME_SPEC = ("pure", "ratio", "derived3", "project", "x", APRIME, "recent", "distance",
               "none", "pattern")


def eval_rjl_all(spec):
    method = make_method(spec)
    out = {}
    for pid, f, a, b, ins in build_rjl()[0]:
        try:
            x, n, conv = method(f, a, b, XTOL, MAXFEV)
            x = float(x)
        except Exception:
            x, n, conv = float("nan"), MAXFEV, False
        ok = bool(conv and solved(f, x))
        over = ok and int(n) > 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4
        out[pid] = [int(n), ok, bool(over), bool(ins)]
    return out


def job(j):
    kind, name = j
    if kind == "protocol":
        return kind, name, eval_protocol(("base", "modab"))
    if kind in ("seen", "fourth"):
        return kind, name, sweep_fourth.evaluate((name, kind))[1]
    spec = APRIME_SPEC if name == "Aprime" else ("base", name)
    return kind, name, eval_rjl_all(spec)


def cost(v):
    return v[0] if v[1] else MAXFEV


def h2h(a, b, keep=None):
    ps = [p for p in a if keep is None or keep(p)]
    w = sum(cost(a[p]) < cost(b[p]) for p in ps)
    l = sum(cost(a[p]) > cost(b[p]) for p in ps)
    return w, len(ps) - w - l, l


def report(R):
    ap_proto = {r["label"]: r for r in json.load(open("sweep_aprime.json", encoding="utf8"))}[
        f"pure+extP+{APRIME}.ratio.derived3.project"]
    full = {r["label"]: r for r in json.load(open("sweep_full.json", encoding="utf8"))}
    mp = R["protocol"]["modab"]
    print("PAPER PROTOCOL (published roots)")
    for s in ("APS", "CHA"):
        a = {p: v for p, v in ap_proto[s]["per"].items()}
        m = {p: v for p, v in mp[s]["per"].items()}
        w, t, l = h2h(a, m)
        print(f"  {s}: A' {ap_proto[s]['total']}  ModAB {mp[s]['total']} (solved {mp[s]['solved']}, "
              f"over {mp[s]['over']})  Chandrupatla {full['chandrupatla'][s]['total']}   "
              f"A' vs ModAB {w}/{t}/{l} p={sign_p(w, l):.2g}")
    seen = json.load(open("sweep_fourth_seen.json", encoding="utf8"))
    four = json.load(open("sweep_fourth_fourth.json", encoding="utf8"))
    for k in seen:
        seen[k].update(four[k])
    md = dict(R["seen"]["modab"])
    md.update(R["fourth"]["modab"])
    print("\nSIX CORPORA (uniform criterion): A' / ModAB / Chandrupatla totals; ModAB violations; "
          "A' vs ModAB")
    tot_a = tot_m = 0
    for s in SETS:
        a, m = seen["Aprime"][s], md[s]
        ta, tm = sum(cost(v) for v in a.values()), sum(cost(v) for v in m.values())
        tc = sum(cost(v) for v in seen["chandrupatla"][s].values())
        tot_a += ta
        tot_m += tm
        w, t, l = h2h(a, m)
        print(f"  {s:<4} {ta:6d} / {tm:6d} / {tc:6d}   viol {sum(v[2] for v in m.values()):3d}   "
              f"{w}/{t}/{l} p={sign_p(w, l):.2g}")
    print(f"  all  {tot_a:6d} / {tot_m:6d}")
    a, m = seen["Aprime"]["D"], md["D"]
    w, t, l = h2h(a, m)
    m1 = (sum(cost(v) for v in a.values()) < sum(cost(v) for v in m.values())
          and w > l and sign_p(w, l) < 0.05)
    print(f"\nM1 (corpus 4): {'PASSED' if m1 else 'NOT passed'} - A' "
          f"{sum(cost(v) for v in a.values())} vs ModAB {sum(cost(v) for v in m.values())}, "
          f"{w}/{t}/{l} p={sign_p(w, l):.2g}")
    for cls in ("D1", "D2", "D3", "D4", "D5", "D6"):
        k = lambda p, c=cls: p.startswith(c)
        w, t, l = h2h(a, m, k)
        print(f"    {cls}: A' {sum(cost(a[p]) for p in a if k(p))}  ModAB "
              f"{sum(cost(m[p]) for p in m if k(p))}   {w}/{t}/{l}")
    rj = R["rjl"]
    print("\nMODAB'S OWN BENCHMARK (Roots.jl f01-f91)")
    for title, keep in (("all 91", None), ("78 out of the development sets", lambda p: not rj["Aprime"][p][3])):
        sel = [p for p in rj["Aprime"] if keep is None or keep(p)]
        tots = {k: sum(cost(rj[k][p]) for p in sel) for k in rj}
        w, t, l = h2h({p: rj["Aprime"][p] for p in sel}, rj["modab"])
        print(f"  {title}: " + "  ".join(f"{k} {tots[k]}" for k in ("Aprime", "modab", "chandrupatla"))
              + f"   A' vs ModAB {w}/{t}/{l} p={sign_p(w, l):.2g}   ModAB viol "
              f"{sum(rj['modab'][p][2] for p in sel)}")


def main():
    if "--report" not in sys.argv:
        jobs = [("protocol", "modab"), ("seen", "modab"), ("fourth", "modab"),
                ("rjl", "Aprime"), ("rjl", "modab"), ("rjl", "chandrupatla")]
        with Pool(len(jobs)) as pool:
            out = pool.map(job, jobs)
        R = {}
        for kind, name, res in out:
            R.setdefault(kind, {})[name] = res
        with open("sweep_modab.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_modab.json", encoding="utf8") as fh:
        R = json.load(fh)
    report(R)
    print("modab done")


if __name__ == "__main__":
    main()