"""Stage 11 - are our implementations of the competitors faithful?

The proposed rule's margin over Brent is 3% and over TOMS 748 is 0.5%, and
Chandrupatla's own algorithm - added on revision - now beats
both of them and us.  A conclusion that rests on margins this small is only as
good as the competitors' implementations, so each is checked against SciPy's,
by evaluation count, problem by problem.

SciPy references used:
    brentq                     compiled C, same algorithm as our brent
    toms748                    pure Python, the one we already call
    _chandrupatla              pure Python, SciPy's implementation of the
                               method whose problem set is half this benchmark

A disagreement in the count is not automatically a bug in ours - the stopping
rules differ in their relative-tolerance terms - so the comparison is reported
per problem and in aggregate, and the tolerances are set to make the two as
comparable as they can be.
"""
from __future__ import annotations

import math
import sys

import numpy as np
from scipy.optimize import brentq
from scipy.optimize._chandrupatla import _chandrupatla
from scipy.optimize._tstutils import (_APS_TESTS_DICTS,
                                      _CHANDRUPATLA_TESTS_DICTS)

from simroot.baselines import Counter, _Budget, brent, chandrupatla, toms748

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


def load():
    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 counted(f):
    st = {"n": 0}

    def g(x):
        st["n"] += 1
        return f(x)

    return g, st


def scipy_brentq(f, a, b):
    g, st = counted(f)
    try:
        x = brentq(g, a, b, xtol=XTOL, rtol=8.9e-16, maxiter=MAXFEV)
        return x, st["n"], True
    except Exception:
        return float("nan"), st["n"], False


def scipy_chandrupatla(f, a, b):
    g, st = counted(f)

    def vec(x):
        return np.asarray([g(float(xi)) for xi in np.atleast_1d(x)]).reshape(
            np.shape(x))

    try:
        res = _chandrupatla(vec, a, b, xatol=XTOL, xrtol=8.9e-16,
                            fatol=0.0, frtol=0.0, maxiter=MAXFEV)
        return float(np.asarray(res.x).ravel()[0]), st["n"], bool(
            np.all(np.asarray(res.success)))
    except Exception:
        return float("nan"), st["n"], False


def main():
    probs = load()
    pairs = [
        ("brent", lambda f, a, b: brent(f, a, b, XTOL, MAXFEV), scipy_brentq,
         "scipy.brentq (C)"),
        ("chandrupatla",
         lambda f, a, b: chandrupatla(f, a, b, XTOL, MAXFEV),
         scipy_chandrupatla, "scipy._chandrupatla"),
    ]
    for name, mine, theirs, label in pairs:
        print("=" * 78, flush=True)
        print(f"{name} : ours against {label}", flush=True)
        print("=" * 78, flush=True)
        t_mine = t_theirs = 0
        ok_mine = ok_theirs = 0
        diffs = []
        for _s, pid, f, a, b, root in probs:
            xm, nm, cm = mine(f, a, b)
            xt, nt, ct = theirs(f, a, b)
            gm = bool(cm and math.isfinite(xm) and abs(xm - root) <= ACC)
            gt = bool(ct and math.isfinite(xt) and abs(xt - root) <= ACC)
            t_mine += nm if gm else MAXFEV
            t_theirs += nt if gt else MAXFEV
            ok_mine += gm
            ok_theirs += gt
            if gm and gt and nm != nt:
                diffs.append((pid, nm, nt))
        print(f"   ours   : {ok_mine}/{len(probs)} solved, {t_mine} evaluations",
              flush=True)
        print(f"   theirs : {ok_theirs}/{len(probs)} solved, {t_theirs} "
              f"evaluations", flush=True)
        print(f"   problems where the counts differ: {len(diffs)} of "
              f"{len(probs)}", flush=True)
        if diffs:
            worse = sum(1 for _p, nm, nt in diffs if nm > nt)
            print(f"      ours dearer on {worse}, cheaper on "
                  f"{len(diffs) - worse}", flush=True)
            diffs.sort(key=lambda t: abs(t[1] - t[2]), reverse=True)
            for pid, nm, nt in diffs[:8]:
                print(f"      {pid:12s} ours {nm:>4d}   theirs {nt:>4d}",
                      flush=True)
        rel = abs(t_mine - t_theirs) / max(t_theirs, 1)
        verdict = ("faithful" if rel < 0.05 and abs(ok_mine - ok_theirs) <= 1
                   else "SUSPECT - the paper must not lean on this competitor")
        print(f"   aggregate difference {rel * 100:.1f}%  =>  {verdict}\n",
              flush=True)


if __name__ == "__main__":
    main()