"""PREREG_out_of_sample_2026-09-21: the two strongest variants, their leave-one-out ablations
and the three library rivals, on APS, CHA and the out-of-sample corpus (corpus_new.py).

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

import json
import math
import sys
from multiprocessing import Pool

from corpus_new import build as build_new
from sweep_full import ACC, MAXFEV, XTOL, load_problems, make_method

P = "project"


def spec(pre, ext="none"):
    return ("pure", "ratio", "derived3", P, "x", pre, "recent", "distance", "none", ext)


CONFIGS = {
    "base": spec("none"),
    "A": spec("contra+flatanalogmempt+exp+half", "pattern"),
    "A-G2": spec("contra+flatanalogmempt+exp+half"),
    "A-contra": spec("flatanalogmempt+exp+half", "pattern"),
    "A-half": spec("contra+flatanalogmempt+exp", "pattern"),
    "A-exp": spec("contra+flatanalogmempt+half", "pattern"),
    "A-point": spec("contra+flatanalogmem+exp+half", "pattern"),
    "A-regime": spec("contra+half", "pattern"),
    "B": spec("defect+flatanalogmempt+exp"),
    "B-defect": spec("flatanalogmempt+exp"),
    "B-exp": spec("defect+flatanalogmempt"),
    "B-point": spec("defect+flatanalogmem+exp"),
    "B-regime": spec("defect"),
    "chandrupatla": ("base", "chandrupatla"),
    "brent": ("base", "brent"),
    "toms748": ("base", "toms748"),
}


def problems():
    out = [(s, pid, f, a, b, r) for s, pid, f, a, b, r in load_problems()]
    new, _ = build_new()
    out += [("NEW", pid, f, a, b, r) for _cls, pid, f, a, b, r in new]
    return out


def _reset_random_problem():
    """SciPy's original.05.00 (our C0.05) is a random step function: unseeded random() values
    kept in a module-level cache shared by every run in the process.  Addendum to the PREREG
    (after two unseeded runs gave the same verdicts): clear the cache and seed before EVERY
    run, so each method meets the same function and the run is reproducible."""
    import random

    import scipy.optimize._tstutils as T
    T._f6_cache.clear()
    random.seed(20260921)


def evaluate(name):
    method = make_method(CONFIGS[name])
    res = {}
    for s, pid, f, a, b, root in problems():
        _reset_random_problem()
        try:
            x, n, conv = method(f, a, b, XTOL, MAXFEV)
        except Exception:
            x, n, conv = float("nan"), MAXFEV, False
        ok = bool(conv and math.isfinite(x) and abs(x - root) <= ACC)
        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 total(r, s, keep=None):
    return sum((n if ok else MAXFEV) for pid, (n, ok, _o) in r[s].items()
               if keep is None or keep(pid))


def h2h(r, o, s, keep=None):
    w = t = l = 0
    for pid, (n, ok, _o) in r[s].items():
        if keep is not None and not keep(pid):
            continue
        c1 = n if ok else MAXFEV
        n2, ok2, _ = o[s][pid]
        c2 = n2 if ok2 else MAXFEV
        w += c1 < c2
        t += c1 == c2
        l += c1 > c2
    return w, t, l


def sign_p(w, l):
    n = w + l
    if n == 0:
        return 1.0
    k = min(w, l)
    return min(1.0, 2.0 * sum(math.comb(n, i) for i in range(k + 1)) / 2.0 ** n)


def report(R):
    ch = R["chandrupatla"]
    print("TOTALS                 APS    CHA    NEW  (NEW C0-C2  NEW C3-C5)")
    neutral = lambda pid: pid[:2] in ("C0", "C1", "C2")
    ours_home = lambda pid: pid[:2] in ("C3", "C4", "C5")
    for k in ("base", "A", "B", "chandrupatla", "brent", "toms748"):
        r = R[k]
        print(f"  {k:<18} {total(r, 'APS'):6d} {total(r, 'CHA'):6d} {total(r, 'NEW'):6d}"
              f"  ({total(r, 'NEW', neutral):6d}  {total(r, 'NEW', ours_home):6d})")
    print("\nOUT OF SAMPLE (NEW) - verdicts")
    for v in ("A", "B"):
        r = R[v]
        o1 = total(r, "NEW") < total(ch, "NEW")
        w, t, l = h2h(r, ch, "NEW")
        o2 = w > l and sign_p(w, l) < 0.05
        viol = sum(o for _n, _ok, o in r["NEW"].values())
        solved = sum(ok for _n, ok, _o in r["NEW"].values())
        solved_ch = sum(ok for _n, ok, _o in ch["NEW"].values())
        o3 = viol == 0 and solved >= solved_ch
        o4 = total(r, "NEW") < total(R["base"], "NEW")
        print(f"  {v}: O1 {'✓' if o1 else '✗'} ({total(r, 'NEW')} vs {total(ch, 'NEW')})  "
              f"O2 {'✓' if o2 else '✗'} ({w}/{t}/{l} p={sign_p(w, l):.2g})  "
              f"O3 {'✓' if o3 else '✗'} (viol {viol}, solved {solved} vs {solved_ch})  "
              f"O4 {'✓' if o4 else '✗'}  -> "
              + ("TRANSFERS" if o1 and o2 and o3 else "does NOT transfer"))
        for rn in ("brent", "toms748"):
            w2, t2, l2 = h2h(r, R[rn], "NEW")
            print(f"       vs {rn:<8} total {total(r, 'NEW')} vs {total(R[rn], 'NEW')}  "
                  f"per problem {w2}/{t2}/{l2} p={sign_p(w2, l2):.2g}")
        print("       per class (ours / Chandrupatla, per problem w/t/l):")
        for cls in ("C0", "C1", "C2", "C3", "C4", "C5"):
            keep = lambda pid, c=cls: pid.startswith(c)
            w3, t3, l3 = h2h(r, ch, "NEW", keep)
            print(f"         {cls}: {total(r, 'NEW', keep):5d} / {total(ch, 'NEW', keep):5d}"
                  f"   {w3}/{t3}/{l3} p={sign_p(w3, l3):.2g}")
    print("\nABLATION  delta = (without) - (with); + = the component earns its place")
    for v, parts in (("A", ("G2", "contra", "half", "exp", "point", "regime")),
                     ("B", ("defect", "exp", "point", "regime"))):
        r = R[v]
        for p in parts:
            o = R[f"{v}-{p}"]
            d = [total(o, s) - total(r, s) for s in ("APS", "CHA", "NEW")]
            earned = any(x > 0 for x in d)
            print(f"  {v} - {p:<7} APS {d[0]:+5d}  CHA {d[1]:+5d}  NEW {d[2]:+5d}   "
                  + ("earned" if earned else "NOT earned"))


def main():
    if "--report" not in sys.argv:
        with Pool(min(8, len(CONFIGS))) as pool:
            R = dict(pool.map(evaluate, list(CONFIGS)))
        with open("sweep_oos.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_oos.json", encoding="utf8") as fh:
        R = json.load(fh)
    report(R)
    print("oos done")


if __name__ == "__main__":
    main()