"""PREREG_estimated_transform_2026-09-22 (control K3): Chandrupatla on sign(f)|f|^(1/m_hat), m_hat
estimated during the run by Lagouanelle (K3L) or by A's analogy order (K3A), on the six corpora with
the harness of sweep_fourth.py (same problems, seeding, charging, uniform criterion).  Report E1-E4.

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

import functools
import json
import sys
from multiprocessing import Pool

import sweep_fourth
import sweep_full
from simroot.transformed import chandrupatla_transformed
from sweep_controls import nonsimple
from sweep_oos import sign_p

for _name, _est in (("k3_lag", "lag"), ("k3_analogy", "analogy")):
    sweep_full.EXTRA_BASELINES[_name] = functools.partial(chandrupatla_transformed, estimator=_est)
    sweep_fourth.CONFIGS[_name] = ("base", _name)
NAMES = ("k3_lag", "k3_analogy")
SETS = ("APS", "CHA", "NEW", "RJL", "APP", "D")


def job(j):
    name, which = j
    return name, sweep_fourth.evaluate((name, which))[1]


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


def tot(r, keep=None):
    return sum(cost(v) for p, v in r.items() if keep is None or keep(p))


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):
    seen = json.load(open("sweep_fourth_seen.json", encoding="utf8"))
    four = json.load(open("sweep_fourth_fourth.json", encoding="utf8"))
    ref = {}
    for k in ("Aprime", "chandrupatla"):
        ref[k] = dict(seen[k])
        ref[k].update(four[k])
    ap, ch = ref["Aprime"], ref["chandrupatla"]
    K1 = json.load(open("sweep_controls.json", encoding="utf8"))["K1"]
    print("E1 - corpus 4 (decisive): K3 total vs A' total; A' w/t/l against K3, p")
    verdict = []
    for k in NAMES:
        w, t, l = h2h(ap["D"], R[k]["D"])
        k3_not_worse = tot(R[k]["D"]) <= tot(ap["D"])
        k3_wins = l > w and sign_p(w, l) < 0.05
        a_wins = tot(ap["D"]) < tot(R[k]["D"]) and w > l and sign_p(w, l) < 0.05
        verdict.append("withdraw" if (k3_not_worse or k3_wins) else ("stands" if a_wins else "intermediate"))
        print(f"  {k}: {tot(R[k]['D'])} vs A' {tot(ap['D'])}   A' {w}/{t}/{l} p={sign_p(w, l):.2g}  -> {verdict[-1]}")
    overall = "withdraw" if "withdraw" in verdict else ("stands" if all(v == "stands" for v in verdict)
                                                        else "intermediate")
    print(f"  outcome for the claim of 6.6: {overall}")
    print("\nE2 - six corpora totals (violations of the 2x bound in brackets)")
    for k, r in (("K3L", R["k3_lag"]), ("K3A", R["k3_analogy"]), ("A'", ap), ("Chandrupatla(f)", ch)):
        cells = " ".join(f"{s} {tot(r[s])}" + (f"({sum(v[2] for v in r[s].values())})" if k.startswith("K3") else "")
                         for s in SETS)
        print(f"  {k:<16} {cells}  | all {sum(tot(r[s]) for s in SETS)}")
    print("\nE3 - multiple/fractional roots (group of Table 11): K3L / K3A / A' / oracle (Chandrupatla on g)")
    for s in ("CHA", "NEW", "D"):
        m = K1[s]["mstar"]
        keep = lambda p, _m=m: nonsimple(_m[p])  # noqa: E731
        orc = tot(K1[s]["chandrupatla_g"], keep)
        a_ = tot(ap[s], keep)
        row = []
        for k in NAMES:
            v = tot(R[k][s], keep)
            share = (a_ - v) / (a_ - orc) if a_ != orc else float("nan")
            row.append(f"{k} {v} (closes {100 * share:+.0f}% of A'->oracle)")
        print(f"  {s}: n={sum(1 for p in m if keep(p))}  " + "  ".join(row) + f"  A' {a_}  oracle {orc}")
    print("\nE4 - the other problems: K3 against Chandrupatla on f (cost of false readings)")
    for s in ("CHA", "NEW", "D"):
        m = K1[s]["mstar"]
        keep = lambda p, _m=m: not nonsimple(_m[p])  # noqa: E731
        print(f"  {s}: " + "  ".join(f"{k} {tot(R[k][s], keep)}" for k in NAMES)
              + f"  Chandrupatla(f) {tot(ch[s], keep)}  A' {tot(ap[s], keep)}")
    for k in NAMES:
        keep = lambda p: True  # noqa: E731
        print(f"  {k} vs Chandrupatla(f), APS: {tot(R[k]['APS'])} vs {tot(ch['APS'])}")


def main():
    if "--report" not in sys.argv:
        jobs = [(n, w) for n in NAMES for w in ("seen", "fourth")]
        with Pool(len(jobs)) as pool:
            out = pool.map(job, jobs)
        R = {}
        for name, res in out:
            R.setdefault(name, {}).update(res)
        with open("sweep_estimated.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_estimated.json", encoding="utf8") as fh:
        R = json.load(fh)
    report(R)
    print("estimated done")


if __name__ == "__main__":
    main()