"""The full sweep: every similarity function x every scale x every safeguard,
over both benchmark sets, against every baseline.

Purpose.  Every claim in the paper so far rests on one configuration
(similarity = ratio) measured on eight hand-picked problems, with only two
configurations carried to the 199-problem benchmark.  The convergence law of
stage 0c predicts the behaviour of ALL nine similarity functions; a law
confirmed on eight problems and one configuration is an anecdote.  This sweep
is what turns it into evidence, and what would refute it if it is wrong.

Grid
    similarity  9  (ratio sqrt cbrt rational ratio_soft exp_abs
                    rat_sq exp_sq cosine)
    scale       4  (start bracket derived derived3)
    safeguard   4  (none step bracket both)
    = 144 configurations, plus 10 baselines, over 199 problems.

Cost model and reporting rules are the same as bench_full.py: evaluations of
f; failures charged the full budget; the 2x-bisection bound checked only on
solved runs; ITP reported both at its recommended constants and as a
per-problem oracle over a (kappa1, n0) grid.

Run:  python -u sweep_full.py            full sweep
      python -u sweep_full.py --smoke    a reduced grid, for the pre-flight
"""
from __future__ import annotations

import json
import math
import os
import sys
import time
from multiprocessing import Pool

os.environ.setdefault("OMP_NUM_THREADS", "1")

from scipy.optimize._tstutils import (_APS_TESTS_DICTS,
                                      _CHANDRUPATLA_TESTS_DICTS)

from simroot.baselines import (anderson_bjorck, best_of, bisection, brent,
                               brent_scipy, chandrupatla, chandrupatla_scipy,
                               dekker, illinois, itp, pegasus, regula_falsi,
                               ridders, toms748)
from simroot.pure import PURE_SIMS, pure_root

XTOL = 1e-12
MAXFEV = 200
ACC = 1e-6

SCALES = ("start", "bracket", "derived", "derived3")
GUARD_MODES = ("none", "step", "bracket", "both", "project", "safe",
               "credit")
ITP_GRID = [(fac, n0) for fac in (1.0, 0.1, 0.01, 1e-4)
            for n0 in (0, 1, 2, 5)]


def load_problems():
    """(set, id, f, a, b, root).  Rebuilt inside each worker so nothing that
    has to cross a process boundary is a closure."""
    out = []
    for d in _APS_TESTS_DICTS:
        a, b = d["bracket"]
        args = d.get("args", ())
        out.append(("APS", d["ID"],
                    (lambda ff, aa: lambda x: ff(x, *aa))(d["f"], args),
                    float(a), float(b), float(d["root"])))
    for d in _CHANDRUPATLA_TESTS_DICTS:
        a, b = d["bracket"]
        out.append(("CHA", d["ID"], (lambda ff: lambda x: ff(x))(d["f"]),
                    float(a), float(b), float(d["root"])))
    return out


def itp_oracle(f, a, b, xtol, maxfev):
    best = None
    for fac, n0 in ITP_GRID:
        try:
            x, n, conv = itp(f, a, b, xtol, maxfev,
                             kappa1=fac * 0.2 / (b - a), n0=n0)
        except ValueError:
            continue
        if conv and (best is None or n < best[1]):
            best = (x, n, True)
    return best if best is not None else (0.5 * (a + b), maxfev, False)


# Competitors are taken at the better of SciPy's implementation and ours,
# problem by problem, wherever SciPy ships the method.  Ours turned out to be
# the cheaper of the two - 6991 against brentq's 7004, 5209 against
# _chandrupatla's 5228 - so quoting the library alone would have flattered us.
BASELINES = {
    "bisection": bisection, "regula_falsi": regula_falsi,
    "illinois": illinois, "pegasus": pegasus,
    "anderson_bjorck": anderson_bjorck, "ridders": ridders,
    "dekker": dekker,
    "brent": best_of(brent, brent_scipy),
    "chandrupatla": best_of(chandrupatla, chandrupatla_scipy),
    "toms748": toms748,
    "itp": itp, "itp_oracle": itp_oracle,
}


def make_method(spec):
    """spec is ('base', name) or ('pure', similarity, scale, guard[, space]).

    The optional 5th element is the SPACE the step is read in (PREREG
    amendment 1, 2026-09-20): omitted or "x" reproduces the published grid
    bit-for-bit, "log" is the logarithmic-march mode."""
    if spec[0] == "base":
        return BASELINES[spec[1]]
    sim, scale, guard = spec[1], spec[2], spec[3]
    space = spec[4] if len(spec) > 4 else "x"
    prestep = spec[5] if len(spec) > 5 else "none"
    points = spec[6] if len(spec) > 6 else "recent"
    direction = spec[7] if len(spec) > 7 else "distance"

    def m(f, a, b, xtol, maxfev):
        return pure_root(f, a, b, xtol, maxfev, similarity=sim,
                         scale=scale, guard=guard, space=space, prestep=prestep,
                         points=points, direction=direction)
    return m


def label(spec):
    if spec[0] == "base":
        return spec[1]
    tag = "pure"
    if len(spec) > 4 and spec[4] != "x":
        tag = spec[4]
    if len(spec) > 6 and spec[6] != "recent":
        tag = f"{tag}+shape"
    if len(spec) > 7 and spec[7] != "distance":
        tag = f"{tag}+dissim"
    if len(spec) > 5 and spec[5] != "none":
        # name the gate by WHICH gate it is: "progress" is a measured negative and must
        # never be confused with "slope" in a results file
        tag = f"{tag}+{spec[5]}" if tag != "pure" else spec[5]
    return f"{tag}.{spec[1]}.{spec[2]}.{spec[3]}"


def evaluate(spec):
    """Run one configuration over all 199 problems.  Returns a plain dict so
    it can cross the process boundary."""
    method = make_method(spec)
    probs = load_problems()
    res = {"label": label(spec), "spec": list(spec)}
    for setname in ("APS", "CHA"):
        total = solved = over = 0
        n_probs = 0
        per = {}
        for s, pid, f, a, b, root in probs:
            if s != setname:
                continue
            n_probs += 1
            try:
                x, n, conv = method(f, a, b, XTOL, MAXFEV)
            except Exception:
                x, n, conv = float("nan"), MAXFEV, False
            # scipy hands back numpy scalars, whose booleans are numpy.bool_
            # and do not serialise; coerce at the boundary rather than at the
            # dump, so nothing numpy-shaped survives into the results
            ok = bool(conv and math.isfinite(x) and abs(x - root) <= ACC)
            n = int(n)
            total += n if ok else MAXFEV
            per[pid] = [n, ok]
            if ok:
                solved += 1
                if n > 2 * math.ceil(math.log2(abs(b - a) / XTOL)) + 4:
                    over += 1
        res[setname] = {"total": total, "solved": solved, "n": n_probs,
                        "over": over, "per": per}
    return res


def _preflight_log_mode():
    """Fail LOUDLY if the log mode is not actually available.

    `evaluate` catches every exception and charges the full budget, so a stale
    deployment (an old simroot/pure.py without `space`) looks exactly like a method
    that solves nothing: on 2026-09-20 that produced a whole sweep of 39800 = 199 x
    MAXFEV before the cause was found. One call up front turns that into an error."""
    x, _n, ok = pure_root(lambda t: t - 1.0, 0.0, 3.0, 1e-10, 100,
                          similarity="ratio", scale="derived3", guard="credit",
                          space="log")
    if not (ok and abs(x - 1.0) < 1e-8):
        raise SystemExit("log mode preflight FAILED: the mode is present but does not solve "
                         "a linear problem — refusing to run a sweep that would report zeros")


def build_grid_log(smoke):
    """The logarithmic-space grid (PREREG 2026-09-20). Baselines are NOT recomputed:
    their numbers are already in sweep_full.json and re-running known expensive
    controls is waste. Only the derived scales are admissible in u-space."""
    _preflight_log_mode()
    sims = ["ratio", "cbrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "project") if smoke else GUARD_MODES
    return [("pure", sim, sc, gu, "log")
            for sim in sims for sc in ("derived", "derived3") for gu in guards]


def build_grid_cap(smoke):
    """Scale as a SELECTION between the two computed ones (PREREG 2026-09-21): identical to
    the slope+dissimilar configuration except scale = derived3_cap, so one thing changes."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "safe") if smoke else GUARD_MODES
    return [("pure", sim, "derived3_cap", gu, "x", "slope", "recent", "dissimilar")
            for sim in sims for gu in guards]


def build_grid_combo(smoke):
    """A+B together (PREREG 2026-09-21 §2: licensed only because each mode alone improved
    APS — A 2527 and B 2726 against the published 2806)."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "project") if smoke else GUARD_MODES
    return [("pure", sim, "derived3", gu, "x", "slope", "shape", "dissimilar")
            for sim in sims for gu in guards]


def build_grid_points(smoke):
    """Mode A (PREREG 2026-09-21): the triple chosen by shape, on top of the slope gate that
    already reached the project target. Three-point scale only — a triple must exist."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "project") if smoke else GUARD_MODES
    return [("pure", sim, "derived3", gu, "x", "slope", "shape", "distance")
            for sim in sims for gu in guards]


def build_grid_direction(smoke):
    """Mode B (PREREG 2026-09-21): the far end read by dissimilarity, on top of the slope gate."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "project") if smoke else GUARD_MODES
    return [("pure", sim, sc, gu, "x", "slope", "recent", "dissimilar")
            for sim in sims for sc in ("derived", "derived3") for gu in guards]


def build_grid_slope(smoke):
    """The slope-stability gate grid (PREREG 2026-09-21). Only derived scales: there must be
    a slope to judge. Baselines are not recomputed — they are in sweep_full.json."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "safe") if smoke else GUARD_MODES
    return [("pure", sim, sc, gu, "x", "slope")
            for sim in sims for sc in ("derived", "derived3") for gu in guards]


def build_grid_gate(smoke):
    """The pre-step gate grid (PREREG 2026-09-21): the published x-space step with the
    gate switched on. Baselines are not recomputed - they are already in sweep_full.json."""
    sims = ["ratio", "sqrt"] if smoke else sorted(PURE_SIMS)
    guards = ("credit", "safe") if smoke else GUARD_MODES
    return [("pure", sim, sc, gu, "x", "progress")
            for sim in sims for sc in ("derived", "derived3") for gu in guards]


def build_grid(smoke):
    specs = [("base", nm) for nm in BASELINES]
    sims = sorted(PURE_SIMS)
    scales = SCALES
    guards = GUARD_MODES
    if smoke:
        sims = ["ratio", "exp_abs"]
        scales = ("derived", "derived3")
        guards = ("step", "both")
        specs = [("base", nm) for nm in ("bisection", "brent", "toms748")]
    for sim in sims:
        for sc in scales:
            for gu in guards:
                specs.append(("pure", sim, sc, gu))
    return specs


def main():
    smoke = "--smoke" in sys.argv
    logmode = "--log" in sys.argv          # the u-space grid, written to its OWN file
    gatemode = "--gate" in sys.argv        # the pre-step gate grid, likewise
    slopemode = "--slope" in sys.argv      # the slope-stability gate grid
    capmode = "--cap" in sys.argv          # scale selected between the two computed ones
    combomode = "--combo" in sys.argv      # modes A and B together
    pointsmode = "--points" in sys.argv    # mode A: triple by shape
    dirmode = "--direction" in sys.argv    # mode B: far end by dissimilarity
    specs = (build_grid_cap(smoke) if capmode else
             build_grid_combo(smoke) if combomode else
             build_grid_points(smoke) if pointsmode else
             build_grid_direction(smoke) if dirmode else
             build_grid_slope(smoke) if slopemode else
             build_grid_gate(smoke) if gatemode else
             build_grid_log(smoke) if logmode else build_grid(smoke))
    t0 = time.time()
    print(f"sweep: {len(specs)} configurations x 199 problems "
          f"({'SMOKE' if smoke else 'FULL'}{', LOG-SPACE' if logmode else ''})", flush=True)

    workers = max(1, min(8, (os.cpu_count() or 2) - 1))
    with Pool(workers) as pool:
        results = []
        for i, r in enumerate(pool.imap_unordered(evaluate, specs), 1):
            results.append(r)
            if i % 10 == 0 or i == len(specs):
                print(f"  {i}/{len(specs)} done, "
                      f"{time.time() - t0:.0f}s elapsed", flush=True)

    # the published sweep_full.json is what verify_article.py checks the paper
    # against, so the log grid never writes into it
    if capmode or combomode or pointsmode or dirmode:
        base = ("sweep_cap" if capmode else
                "sweep_combo" if combomode else
                "sweep_points" if pointsmode else "sweep_direction")
        out = f"{base}_smoke.json" if smoke else f"{base}.json"
        with open(out, "w", encoding="utf8") as fh:
            json.dump(results, fh)
        print(f"wrote {out}  ({time.time() - t0:.0f}s total)", flush=True)
        report(results)
        return
    if slopemode:
        out = "sweep_slope_smoke.json" if smoke else "sweep_slope.json"
        with open(out, "w", encoding="utf8") as fh:
            json.dump(results, fh)
        print(f"\nwrote {out}  ({time.time() - t0:.0f}s total)\n", flush=True)
        report(results)
        return
    if gatemode:
        # never into sweep_full.json: that file is what verify_article.py checks the
        # paper against, and the smoke run of 2026-09-21 did land in sweep_smoke.json
        # before this branch existed
        out = "sweep_gate_smoke.json" if smoke else "sweep_gate.json"
        with open(out, "w", encoding="utf8") as fh:
            json.dump(results, fh)
        print(f"\nwrote {out}  ({time.time() - t0:.0f}s total)\n", flush=True)
        report(results)
        return
    out = ("sweep_log_smoke.json" if smoke else "sweep_log.json") if logmode else \
          ("sweep_smoke.json" if smoke else "sweep_full.json")
    with open(out, "w", encoding="utf8") as fh:
        json.dump(results, fh)
    print(f"\nwrote {out}  ({time.time() - t0:.0f}s total)\n", flush=True)
    report(results)


def report(results, top=20):
    for setname in ("APS", "CHA"):
        rows = sorted(results,
                      key=lambda r: (r[setname]["n"] - r[setname]["solved"],
                                     r[setname]["total"]))
        n = rows[0][setname]["n"]
        print("=" * 84)
        print(f"{setname}: {n} problems   (top {top} of {len(rows)})")
        print("=" * 84)
        print(f"{'configuration':40s}{'solved':>9s}{'total fev':>12s}"
              f"{'over 2x':>9s}")
        print("-" * 84)
        for r in rows[:top]:
            d = r[setname]
            print(f"{r['label']:40s}{d['solved']:>6d}/{d['n']:<3d}"
                  f"{d['total']:>12d}{d['over']:>9d}")
        print()

    # the law of stage 0c, tested across the whole grid: similarities whose
    # (1 - S) has zero slope at 0 must be the slow ones whenever the scale is
    # chosen rather than derived.
    print("=" * 84)
    print("THE STAGE-0C LAW ACROSS THE GRID")
    print("=" * 84)
    quad = ("rat_sq", "exp_sq", "cosine")     # (1-S) ~ c2*z^2, slope c = 0
    lin = ("ratio", "rational", "exp_abs")    # (1-S) ~ z,      slope c = 1
    for scale in SCALES:
        for group, name in ((lin, "slope c=1 (linear)"),
                            (quad, "slope c=0 (quadratic)")):
            tot = solved = cnt = 0
            for r in results:
                if r["spec"][0] != "pure":
                    continue
                sim, sc, gu = r["spec"][1], r["spec"][2], r["spec"][3]
                if sc != scale or gu != "step" or sim not in group:
                    continue
                cnt += 1
                tot += r["APS"]["total"] + r["CHA"]["total"]
                solved += r["APS"]["solved"] + r["CHA"]["solved"]
            if cnt:
                print(f"  scale={scale:9s} {name:26s} "
                      f"mean total {tot / cnt:9.0f}   "
                      f"mean solved {solved / cnt:6.1f}/199")


if __name__ == "__main__":
    main()