"""Chapter 9, code 1: adaptive RKF45 for one first-order ODE x' = f(t,x).

Based on the revised Fortran companion code.
The driver supplies initial data. All stages and adaptive decisions are
inside rkf45. Output contains the initial and every accepted point.

Example: x' = 100/(1+10000*t*t), x(-1) = atan(-100).
Exact: x(t) = atan(100*t). Rapid change near zero makes adaptation visible.
To solve another equation, change the initial data and function f.
Disable check_exact or update BOTH exact formulas (comparison and plot).
Requires NumPy and Matplotlib for arrays and plotting.
"""
import math
import sys

import matplotlib.pyplot as plt
import numpy as np


def main():
    # ---------------- Initial data: edit here ---------------------------
    t0 = -1.0
    x0 = math.atan(100.0*t0)
    t_end = 1.0
    h_first = 0.1
    tolerance = 1.0e-6       # Absolute local error tolerance
    check_exact = True
    # ------------------------------------------------------------------
    t, x, h, accepted, rejected = rkf45(f, t0, x0, t_end, h_first, tolerance)

    filename = "ch09_rkf45_scalar.dat"
    print("Adaptive Runge-Kutta-Fehlberg RKF45")
    print(f"Tolerance: {tolerance:.4e}")
    print(f"{'t':>10} {'numerical x':>18} {'h':>14}")
    with open(filename, "w", encoding="utf-8") as data:
        data.write("# t x h\n")
        for i in range(len(t)):
            print(f"{t[i]:10.6f} {x[i]:18.12f} {h[i]:14.6e}")
            data.write(f"{t[i]:.16e} {x[i]:.16e} {h[i]:.16e}\n")
    print(f"Accepted steps: {accepted:8d}")
    print(f"Rejected trials:{rejected:8d}")
    print(f"Final x:       {x[-1]:18.12f}")
    if check_exact:
        exact = math.atan(100.0*t_end)  # DEFAULT example only
        print(f"Exact final x: {exact:18.12f}")
        print(f"Abs error:     {abs(x[-1]-exact):12.4e}")
    print(f"Data file: {filename}")

    # Dense exact curve and actual numerical points.
    plt.figure()
    if check_exact:
        t_plot = np.linspace(t0, t_end, 2000)
        x_plot = np.arctan(100.0*t_plot)
        plt.plot(t_plot, x_plot, "k-", linewidth=1.2, label="Exact")
        plt.plot(t, x, "ko--", markersize=4, label="RKF45")
        plt.legend()
    else:
        plt.plot(t, x, "ko-", markersize=4)
    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.title("Adaptive RKF45: solution")
    plt.grid(True)
    plt.tight_layout()

    # h[i] is the step used to reach t[i]; omit initial h = 0.
    plt.figure()
    plt.semilogy(t[1:], h[1:], "ko-", markersize=4)
    plt.xlabel("t")
    plt.ylabel("Accepted step size h")
    plt.title("Adaptive RKF45: step size")
    plt.grid(True)
    plt.tight_layout()
    plt.show()


def f(t, x):
    """Supply x' = f(t,x) here. x is available for state-dependent ODEs."""
    return 100.0/(1.0+10000.0*t*t)


def rkf45(f, t0, x0, t_end, h_first, tolerance):
    """Adaptive Fehlberg RKF45; advance with the fifth-order estimate.

    The fourth/fifth-order difference estimates local error; tolerance
    does not bound accumulated global error. Rejection leaves t and x
    unchanged. Return arrays of initial and accepted points, plus counts.
    """
    max_attempts = 100000
    h_min = 1.0e-12
    if t_end <= t0 or h_first <= 0.0 or tolerance <= 0.0:
        raise ValueError("Use t_end > t0, h_first > 0, and tolerance > 0.")
    t = t0
    x = x0
    h = h_first
    accepted = 0
    rejected = 0
    attempts = 0
    t_values = [t0]
    x_values = [x0]
    h_values = [0.0]

    while t < t_end:
        # Ignore only a machine-rounding gap at the endpoint.
        if t_end-t <= 64.0*sys.float_info.epsilon*max(1.0, abs(t_end)):
            break
        if attempts >= max_attempts:
            raise RuntimeError("Maximum number of trials exceeded.")
        h = min(h, t_end-t)
        tf = t+h
        if h < h_min or tf <= t:
            raise RuntimeError("Step too small to advance time.")
        attempts += 1

        # Fehlberg stages: each k includes h.
        k1 = h*f(t, x)
        k2 = h*f(t+h/4.0, x+k1/4.0)
        k3 = h*f(t+3.0*h/8.0, x+3.0*k1/32.0+9.0*k2/32.0)
        k4 = h*f(t+12.0*h/13.0, x+1932.0*k1/2197.0
                 -7200.0*k2/2197.0+7296.0*k3/2197.0)
        k5 = h*f(t+h, x+439.0*k1/216.0-8.0*k2
                 +3680.0*k3/513.0-845.0*k4/4104.0)
        k6 = h*f(t+h/2.0, x-8.0*k1/27.0+2.0*k2
                 -3544.0*k3/2565.0+1859.0*k4/4104.0-11.0*k5/40.0)
        trial = (x+16.0*k1/135.0+6656.0*k3/12825.0
                 +28561.0*k4/56430.0-9.0*k5/50.0+2.0*k6/55.0)
        err = abs(k1/360.0-128.0*k3/4275.0
                  -2197.0*k4/75240.0+k5/50.0+2.0*k6/55.0)
        if err <= tolerance:
            if h == t_end-t:
                tf = t_end
            t = tf
            x = trial
            accepted += 1
            t_values.append(t)
            x_values.append(x)
            h_values.append(h)
        else:
            rejected += 1

        # Embedded error scales approximately as h^5.
        # Safety factor 0.9; limit growth and shrinkage.
        if err == 0.0:
            factor = 2.0
        else:
            factor = 0.9*(tolerance/err)**0.2
        if err <= tolerance:
            h *= min(2.0, factor)
        else:
            h *= max(0.2, factor)

    return (np.array(t_values), np.array(x_values), np.array(h_values),
            accepted, rejected)


if __name__ == "__main__":
    main()