"""POSTHOC_lag_branch_2026-09-22: Chandrupatla on sign(f)|f|^(1/m_hat) with Lagouanelle's estimate
restricted to the branch m >= 2 (K3-L2), post hoc, with the harness of sweep_fourth.py.  The
registered K3 run (sweep_estimated.py / .json) is not touched.

    python -u sweep_estimated_posthoc.py            run -> sweep_estimated_posthoc.json, report
    python -u sweep_estimated_posthoc.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

sweep_full.EXTRA_BASELINES["k3_lag2"] = functools.partial(chandrupatla_transformed, estimator="lag2")
sweep_fourth.CONFIGS["k3_lag2"] = ("base", "k3_lag2")
SETS = ("APS", "CHA", "NEW", "RJL", "APP", "D")


def job(which):
    return sweep_fourth.evaluate(("k3_lag2", 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 held_on_simple_roots():
    """How many APS runs end holding m_hat != 1 (fresh)."""
    n = 0
    for s, pid, f, a, b, _r in sweep_full.load_problems():
        if s != "APS":
            continue
        st = {}
        chandrupatla_transformed(f, a, b, 1e-12, 200, estimator="lag2", stats=st)
        n += st.get("m_hat", 1.0) != 1.0
    return n


def report(R):
    seen = json.load(open("sweep_fourth_seen.json", encoding="utf8"))
    four = json.load(open("sweep_fourth_fourth.json", encoding="utf8"))
    ap = dict(seen["Aprime"])
    ap.update(four["Aprime"])
    K1 = json.load(open("sweep_controls.json", encoding="utf8"))["K1"]
    grp = sum(tot(R[s], lambda p, _m=K1[s]["mstar"]: nonsimple(_m[p])) for s in ("CHA", "NEW", "D"))
    ps = list(ap["D"])
    w = sum(cost(ap["D"][p]) < cost(R["D"][p]) for p in ps)
    l = sum(cost(ap["D"][p]) > cost(R["D"][p]) for p in ps)
    print(f"K3-L2 (post hoc): corpus 4 {tot(R['D'])}  six corpora {sum(tot(R[s]) for s in SETS)}  "
          f"multiple/fractional {grp}  APS {tot(R['APS'])}")
    print(f"  A' on corpus 4 against it: {w}/{len(ps) - w - l}/{l}  p={sign_p(w, l):.2g}")
    print(f"  per corpus: " + " ".join(f"{s} {tot(R[s])}" for s in SETS))
    print(f"  violations {sum(v[2] for s in SETS for v in R[s].values())}")
    print(f"  APS runs ending with m_hat != 1: {held_on_simple_roots()} of 154")


def main():
    if "--report" not in sys.argv:
        with Pool(2) as pool:
            parts = pool.map(job, ("seen", "fourth"))
        R = {}
        for p in parts:
            R.update(p)
        with open("sweep_estimated_posthoc.json", "w", encoding="utf8") as fh:
            json.dump(R, fh)
    with open("sweep_estimated_posthoc.json", encoding="utf8") as fh:
        R = json.load(fh)
    report(R)
    print("posthoc done")


if __name__ == "__main__":
    main()