"""Step-size controllers.  Everything else in the integration is SciPy's
(the RK45 tableau, the stage combination, the embedded error estimate), so
these are the only thing under comparison.

The classical controllers are made of constants somebody chose - a safety
factor, growth clamps, PI gains.  The question this module asks is what
happens if the quantity a constant is GUESSING is measured instead.

    root finding                      step-size control
    ------------                      -----------------
    the scale of the value axis       the safety factor
    chosen  -> brittle, linear        chosen  -> works, three constants
    measured-> superlinear            measured-> works, none

What did NOT transfer is the credit ledger itself.  A first version banked an
accuracy margin exactly as the root-finding safeguard does; the margin turned
out to take no part in any decision and was removed.  The move that transfers
is the more general one - replace a chosen constant by the measured value of
what it approximates - not the particular bookkeeping that carried it there.
"""
from __future__ import annotations

import math


class Controller:
    name = "?"
    constants = 0

    def reset(self):
        pass

    def update(self, h, e, ok, order):
        raise NotImplementedError


class Elementary(Controller):
    """h_new = h * (1/e)^(1/(p+1)) with no safety factor and no clamps.

    Included because it is the controller the error model alone dictates.  It
    is what every practical code refuses to use, and the reason for that
    refusal is exactly what the constants in the others are buying."""

    name = "elementary (no constants)"
    constants = 0

    def update(self, h, e, ok, order):
        if e == 0.0:
            return h * 10.0
        return h * (1.0 / e) ** (1.0 / (order + 1))


class Classical(Controller):
    """The textbook controller: a safety factor and two growth clamps."""

    name = "classical (scipy: safety 0.9, clamps 0.2/10)"
    constants = 3
    # exactly the constants scipy.integrate._ivp.rk uses, so the driver can
    # be validated against solve_ivp with the same control law
    SAFETY, MINF, MAXF = 0.9, 0.2, 10.0

    def update(self, h, e, ok, order):
        if e == 0.0:
            fac = self.MAXF
        else:
            fac = self.SAFETY * (1.0 / e) ** (1.0 / (order + 1))
            fac = min(self.MAXF, max(self.MINF, fac))
        return h * fac


class GustafssonPI(Controller):
    """The PI controller used in modern codes: two gains, a safety factor and
    two clamps."""

    name = "PI (Gustafsson, 5 constants)"
    constants = 5
    SAFETY, MINF, MAXF = 0.9, 0.2, 10.0
    K_I, K_P = 0.3, 0.4

    def reset(self):
        self.e_prev = None

    def update(self, h, e, ok, order):
        if e == 0.0:
            return h * self.MAXF
        p1 = order + 1
        if self.e_prev is None or not ok:
            fac = self.SAFETY * (1.0 / e) ** (1.0 / p1)
        else:
            fac = (self.SAFETY * (1.0 / e) ** (self.K_I / p1)
                   * (self.e_prev / e) ** (self.K_P / p1))
        fac = min(self.MAXF, max(self.MINF, fac))
        if ok:
            self.e_prev = e
        return h * fac


class MeasuredSafety(Controller):
    """The safety factor measured instead of chosen.

    NOTE ON THE NAME.  This began as "the credit rule transferred", and it
    banked a margin exactly as the root-finding safeguard does.  The margin
    turned out to play no part in the decision - the step factor was computed
    from the measured model discrepancy alone - so the banking was vestigial
    and the name was unearned.  What transfers is not the credit ledger but
    the deeper move behind it: replace a chosen constant with the quantity it
    is guessing, measured from the run.  The ledger is left out because it did
    nothing here, and calling this a credit controller would be a mislabel.

    The elementary controller aims at err = tol exactly, and is refused in
    practice because the model err ~ C h^(p+1) is only approximate: aiming at
    the boundary means missing it about half the time, and every miss costs a
    whole rejected step.  The classical answer is a safety factor chosen by
    hand.

    The answer here is to measure the thing the safety factor is guessing.
    After proposing a step, the model predicts what its error will be; the
    step is taken and the error is observed.  The discrepancy

        delta = log2(observed / predicted)

    is exactly how wrong the model was, in bits, and the next step aims at
    2^-delta instead of at 1.  When the model is right, delta -> 0 and the
    controller becomes the elementary one; when the model has just been
    optimistic, the controller backs off by precisely the observed amount.

    The first step has nothing to have learned from, so it aims at the
    tolerance itself - the same licence the root-finding guard grants.
    """

    name = "measured safety (no constants)"
    constants = 0

    def reset(self):
        self.predicted = None      # error the model expected for this step
        self.delta = 0.0           # bits by which the model was last wrong

    def update(self, h, e, ok, order):
        p1 = order + 1
        if self.predicted is not None and self.predicted > 0.0 and e > 0.0:
            # only optimism matters: the model predicting too SMALL an error
            # is what causes rejections.  Pessimism costs nothing and is not
            # corrected for.
            self.delta = max(0.0, math.log2(e / self.predicted))
        if e == 0.0:
            fac = 10.0
        else:
            # aim at 2^-delta of the tolerance rather than at the tolerance
            fac = (2.0 ** (-self.delta) / e) ** (1.0 / p1)
        h_new = h * fac
        # record what the model now predicts for the step just chosen, so the
        # next call can measure how wrong it was
        self.predicted = e * (h_new / h) ** p1
        return h_new


ALL = [Elementary(), Classical(), GustafssonPI(), MeasuredSafety()]