"""PREREG_close_sub_2026-09-21: A' = A + close + sub against A, its two ablations and the
library rivals.

    python -u sweep_fourth.py --seen     development: the five SEEN corpora (descriptive)
    python -u sweep_fourth.py --fourth   the ONE check on corpus_fourth.py -> sweep_fourth.json
    python -u sweep_fourth.py --report   verdicts from sweep_fourth.json
"""
from __future__ import annotations

import json
import math
import sys
from multiprocessing import Pool

from corpus_roots_jl import solved
from sweep_full import MAXFEV, XTOL, load_problems, make_method
from sweep_oos import CONFIGS as OOS
from sweep_oos import sign_p


def spec(pre):
    return ("pure", "ratio", "derived3", "project", "x", pre, "recent", "distance", "none",
            "pattern")


CONFIGS = {
    "base": OOS["base"],
    "A": OOS["A"],
    "Aprime": spec("contra+flatanalogmempt+exp+half+close+sub"),
    "Aprime-close": spec("contra+flatanalogmempt+exp+half+sub"),
    "Aprime-sub": spec("contra+flatanalogmempt+exp+half+close"),
    "chandrupatla": OOS["chandrupatla"],
    "brent": OOS["brent"],
    "toms748": OOS["toms748"],
}


def corpus(which):
    """[(set, id, f, a, b)] - success is always a sign change within 1e-6."""
    out = []
    if which == "fourth":
        from corpus_fourth import build
        out += [("D", pid, f, a, b) for _c, pid, f, a, b, _r in build()[0]]
        return out
    out += [(s, pid, f, a, b) for s, pid, f, a, b, _r in load_problems()]
    from corpus_new import build as b1
    out += [("NEW", pid, f, a, b) for _c, pid, f, a, b, _r in b1()[0]]
    from corpus_roots_jl import build as b2
    out += [("RJL", pid, f, a, b) for pid, f, a, b, ins in b2()[0] if not ins]
    from corpus_applied import build as b3
    out += [("APP", pid, f, a, b) for pid, f, a, b in b3()[0]]
    return out


def evaluate(job):
    name, which = job
    import random

    import scipy.optimize._tstutils as T
    method = make_method(CONFIGS[name])
    res = {}
    for s, pid, f, a, b in corpus(which):
        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
        ok = bool(conv and solved(f, x))
        over = ok and int(n) > 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4
        res.setdefault(s, {})[pid] = [int(n), ok, bool(over)]
    return name, res


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


def h2h(r, o, s, keep=lambda p: True):
    ps = [p for p in r[s] if keep(p)]
    w = sum(cost(r[s][p]) < cost(o[s][p]) for p in ps)
    l = sum(cost(r[s][p]) > cost(o[s][p]) for p in ps)
    return w, len(ps) - w - l, l


def tot(r, s, keep=lambda p: True):
    return sum(cost(v) for p, v in r[s].items() if keep(p))


def report_seen(R):
    print("SEEN corpora (development, descriptive): totals; A' vs Chandrupatla w/t/l")
    for s in ("APS", "CHA", "NEW", "RJL", "APP"):
        w, t, l = h2h(R["Aprime"], R["chandrupatla"], s)
        print(f"  {s:<4} " + "  ".join(f"{k} {tot(R[k], s)}" for k in CONFIGS)
              + f"   A' vs Ch {w}/{t}/{l} p={sign_p(w, l):.2g}")


def report_fourth(R):
    s, ch, ap = "D", R["chandrupatla"], R["Aprime"]
    print("FOURTH corpus (the one check)")
    for k in CONFIGS:
        r = R[k]
        print(f"  {k:<14} total {tot(r, s):6d}  solved {sum(v[1] for v in r[s].values())}/"
              f"{len(r[s])}  violations {sum(v[2] for v in r[s].values())}")
    w, t, l = h2h(ap, ch, s)
    f1 = tot(ap, s) < tot(ch, s)
    f2 = w > l and sign_p(w, l) < 0.05
    f3 = (sum(v[2] for v in ap[s].values()) == 0
          and sum(v[1] for v in ap[s].values()) >= sum(v[1] for v in ch[s].values()))
    f4 = tot(ap, s) < tot(R["A"], s)
    print(f"  F1 {'✓' if f1 else '✗'} ({tot(ap, s)} vs {tot(ch, s)}, "
          f"{100 * (tot(ap, s) - tot(ch, s)) / tot(ch, s):+.1f}%)  F2 {'✓' if f2 else '✗'} "
          f"({w}/{t}/{l} p={sign_p(w, l):.2g})  F3 {'✓' if f3 else '✗'}  F4 {'✓' if f4 else '✗'} "
          f"(A' {tot(ap, s)} vs A {tot(R['A'], s)})  -> "
          + ("A' TRANSFERS" if f1 and f2 and f3 else "A' does NOT transfer"))
    for comp in ("close", "sub"):
        d = tot(R[f"Aprime-{comp}"], s) - tot(ap, s)
        print(f"  ablation - {comp}: {d:+d} ({'earned' if d > 0 else 'NOT earned'})")
    nt = lambda p: p[:2] in ("D1", "D2", "D4", "D5")
    w, t, l = h2h(ap, ch, s, nt)
    print(f"  without the targeted classes (D1 D2 D4 D5): A' {tot(ap, s, nt)} vs Ch {tot(ch, s, nt)}"
          f"  {w}/{t}/{l} p={sign_p(w, l):.2g}")
    for rn in ("brent", "toms748"):
        w, t, l = h2h(ap, R[rn], s)
        print(f"  vs {rn:<8} {tot(ap, s)} vs {tot(R[rn], s)}  {w}/{t}/{l} p={sign_p(w, l):.2g}")
    print("  per class: A' / A / Chandrupatla / Brent / TOMS748   (A' vs Ch w/t/l)")
    for cls in ("D1", "D2", "D3", "D4", "D5", "D6"):
        k = lambda p, c=cls: p.startswith(c)
        w, t, l = h2h(ap, ch, s, k)
        print(f"    {cls}: " + " / ".join(f"{tot(R[n], s, k):5d}" for n in
                                         ("Aprime", "A", "chandrupatla", "brent", "toms748"))
              + f"   {w}/{t}/{l}")


def main():
    which = "fourth" if "--fourth" in sys.argv else ("seen" if "--seen" in sys.argv else None)
    if which is not None:
        with Pool(len(CONFIGS)) as pool:
            R = dict(pool.map(evaluate, [(n, which) for n in CONFIGS]))
        with open(f"sweep_fourth_{which}.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    else:
        which = "fourth"
    with open(f"sweep_fourth_{which}.json", encoding="utf8") as fh:
        R = json.load(fh)
    (report_fourth if which == "fourth" else report_seen)(R)
    print("fourth done")


if __name__ == "__main__":
    main()