"""Two checks on diag_model_order.py before its result may be stated.

  C1  WHAT Chandrupatla does on the problems where we lose, by family: the share of his
      steps that are bisections, from the traced mirror whose evaluation count is asserted
      equal to the original's.  If the flat families are won by bisecting, the gap there
      is not a better model on his side.
  C2  the 5 states where "exactly one negative Lagrange weight" failed: print them.  The
      sign argument is exact in real arithmetic, so each exception must be a
      floating-point degeneracy (a weight that underflowed to 0, or f values so close
      that a difference cancelled) - or the argument is wrong and must be withdrawn.

    python -u diag_model_order_checks.py
"""
from __future__ import annotations

import math
from collections import defaultdict

from scipy.optimize._tstutils import _CHANDRUPATLA_TESTS_DICTS

from diag_model_order import ACC, MAXFEV, OURS, XTOL, ch_cost, iqi
from explain_chandrupatla import chandrupatla_traced
from simroot.baselines import chandrupatla
from simroot.pure import pure_root


def main():
    fam = defaultdict(lambda: [0, 0, 0, 0, 0])     # bisect, iqi, ours n, ch n, problems
    exceptions = []
    for d in _CHANDRUPATLA_TESTS_DICTS:
        pid = str(d["ID"])
        f = d["f"]
        a, b = (float(v) for v in d["bracket"])
        r = float(d["root"])
        st = {"trace": []}
        x, n, conv = pure_root(f, a, b, XTOL, MAXFEV, stats=st, **OURS)
        if not (conv and abs(float(x) - r) <= ACC):
            continue
        ch = ch_cost(f, a, b, r)
        if ch is None or n <= ch:
            continue
        _xc, nc, okc, tr = chandrupatla_traced(f, a, b, XTOL, MAXFEV)
        _xo, no, _oko = chandrupatla(f, a, b, XTOL, MAXFEV)
        assert nc == no, f"mirror drifted on {pid}: {nc} vs {no}"
        k = pid.split(".")[0]
        cell = fam[k]
        cell[0] += tr["bisect"]
        cell[1] += tr["iqi"]
        cell[2] += n
        cell[3] += ch
        cell[4] += 1
        for s in st["trace"]:
            if s["xp2"] is None:
                continue
            pts = ((s["x"], s["fx"]), (s["xp"], s["fp"]), (s["xp2"], s["fp2"]))
            fs = [p[1] for p in pts]
            if not (min(fs) < 0.0 < max(fs)):
                continue
            _xq, w = iqi(*pts)
            if w is not None and sum(v < 0.0 for v in w) != 1:
                exceptions.append((pid, fs, w))

    print("C1  Chandrupatla on the problems where we lose, by family")
    print(f"  {'family':<6} {'probs':>5} {'ours':>6} {'CH':>6} {'gap':>5} "
          f"{'CH bisect share':>16}")
    tot_gap = 0
    for k in sorted(fam):
        bi, iq, on, cn, p = fam[k]
        tot_gap += on - cn
        share = 100.0 * bi / max(bi + iq, 1)
        print(f"  {k:<6} {p:5d} {on:6d} {cn:6d} {on - cn:5d} {share:15.0f}%")
    print(f"  total gap on the losing problems: {tot_gap}")

    print(f"\nC2  the {len(exceptions)} states breaking 'exactly one negative weight'")
    for pid, fs, w in exceptions:
        tiny = min(abs(v) for v in fs)
        print(f"  {pid:<8} f = [{', '.join(f'{v:.3e}' for v in fs)}]   "
              f"w = [{', '.join(f'{v:.3e}' for v in w)}]   "
              f"zero weights: {sum(v == 0.0 for v in w)}   min|f| {tiny:.1e}")


if __name__ == "__main__":
    main()