"""Integration driver whose only original part is the step-size controller.

The tableau, the stage combination and the embedded error estimate are
SciPy's - `scipy.integrate._ivp.rk.rk_step` with RK45's A, B, C and E - so
nothing about the method itself is a transcription of ours.  An earlier
version of this file carried a hand-typed Dormand-Prince tableau; that was an
unnecessary risk and an unnecessary thing for a reader to have to check.

What remains ours is the loop that decides the next step size, which is the
object under study.  Everything else is shared between controllers, so a
difference between two runs can only come from the controller.

The driver is validated against SciPy's own `solve_ivp(method="RK45")`: with
the classical controller - SciPy's own safety factor and clamps - the two
must agree closely on the number of right-hand-side evaluations.
"""
from __future__ import annotations

import math

import numpy as np
from scipy.integrate._ivp.rk import RK45, rk_step

ORDER = RK45.error_estimator_order      # 4; the step exponent is ORDER + 1


class Counter:
    """Counts evaluations of the right-hand side - the cost model."""

    __slots__ = ("f", "n")

    def __init__(self, f):
        self.f = f
        self.n = 0

    def __call__(self, t, y):
        self.n += 1
        return np.asarray(self.f(t, y), dtype=float)


def _norm(v, y, y_new, atol, rtol):
    scale = atol + rtol * np.maximum(np.abs(y), np.abs(y_new))
    return float(np.sqrt(np.mean((v / scale) ** 2)))


def integrate(f, t0, t1, y0, controller, atol=1e-8, rtol=1e-8,
              max_steps=200000):
    """Integrate with the given controller.  Returns what the comparison
    needs.  The initial step is chosen the same way for every controller, so
    none can win by guessing it better."""
    g = Counter(f)
    t = float(t0)
    y = np.asarray(y0, dtype=float)
    fv = g(t, y)
    K = np.empty((RK45.n_stages + 1, y.size), dtype=float)

    d0 = _norm(y, y, y, atol, rtol)
    d1 = _norm(fv, y, y, atol, rtol)
    h = 1e-6 if (d0 < 1e-5 or d1 < 1e-5) else 0.01 * d0 / d1
    h = min(h, abs(t1 - t0))
    controller.reset()

    accepted = rejected = 0
    while t < t1 and accepted + rejected < max_steps:
        h = min(h, t1 - t)
        if h <= 0.0:
            break
        y_new, f_new = rk_step(g, t, y, fv, h, RK45.A, RK45.B, RK45.C, K)
        err = np.dot(K.T, RK45.E) * h            # RK45's own error estimate
        e = _norm(err, y, y_new, atol, rtol)
        ok = e <= 1.0
        h_next = controller.update(h, e, ok, ORDER)
        if ok:
            accepted += 1
            t, y, fv = t + h, y_new, f_new
        else:
            rejected += 1
        if not math.isfinite(h_next) or h_next <= 0.0:
            break
        h = h_next
    return {"t": t, "y": list(y), "nfev": g.n, "accepted": accepted,
            "rejected": rejected,
            "reached_end": abs(t - t1) < 1e-12 * max(1.0, abs(t1))}


def validate_against_scipy(atol=1e-8, rtol=1e-8):
    """The driver against SciPy's own solver, with SciPy's own controller.

    This replaces the hand-written order check that the previous version
    needed: if the driver reproduces solve_ivp's evaluation count with the
    same control law, the tableau, the stage combination, the error estimate
    and the loop are all the library's behaviour."""
    from scipy.integrate import solve_ivp

    from .controllers import Classical

    out = []
    probs = [
        ("decay", lambda t, y: [-y[0]], 0.0, 10.0, [1.0]),
        ("forced", lambda t, y: [-y[0] + math.sin(t)], 0.0, 10.0, [0.0]),
        ("lotka", lambda t, y: [1.5 * y[0] - y[0] * y[1],
                                -3.0 * y[1] + y[0] * y[1]], 0.0, 15.0,
         [10.0, 5.0]),
        ("vdp", lambda t, y: [y[1], 5.0 * (1 - y[0] ** 2) * y[1] - y[0]],
         0.0, 25.0, [2.0, 0.0]),
    ]
    for name, f, t0, t1, y0 in probs:
        mine = integrate(f, t0, t1, y0, Classical(), atol=atol, rtol=rtol)
        ref = solve_ivp(f, (t0, t1), y0, method="RK45", atol=atol,
                        rtol=rtol)
        out.append((name, mine["nfev"], int(ref.nfev),
                    max(abs(a - b) for a, b in zip(mine["y"],
                                                   ref.y[:, -1]))))
    return out