"""Chapter 9, code 0: fixed-step methods for one first-order ODE.

Solve x' = f(t,x), given x(t0) = x0.
Based on the revised Fortran companion code.

key = 1: Euler                     (first order)
key = 2: modified Euler / Heun      (second order)
key = 3: classical Runge-Kutta      (fourth order)

To solve another equation:
    1. Change x0, t0, t_end, and dt in the driver.
    2. Edit the derivative function f below.
    3. Disable check_exact or update the exact formulas in the driver.

Default: x' = x-t*t+1, x(0) = 0.5, 0 <= t <= 4.
Exact: x(t) = (t+1)**2-0.5*exp(t).
Try dt = 0.2, 0.1, and 0.05 with each method key.
Halving dt reduces the final error asymptotically by factors of
about 2 (Euler), 4 (Heun), and 16 (RK4).
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 = 0.0
    x0 = 0.5
    t_end = 4.0
    dt = 0.1
    key = 3
    check_exact = True
    # ------------------------------------------------------------------
    max_steps = 100000
    if dt <= 0.0 or t_end <= t0:
        raise ValueError("Use dt > 0 and t_end > t0.")
    if key == 1:
        method = "Euler"
        filename = "ch09_euler.dat"
    elif key == 2:
        method = "Modified Euler (Heun)"
        filename = "ch09_heun.dat"
    elif key == 3:
        method = "Runge-Kutta 4th order"
        filename = "ch09_rk4.dat"
    else:
        raise ValueError("Choose key = 1, 2, or 3.")

    # Indexed times avoid repeated-addition drift.
    # Include t_end, shortening the final step if needed.
    if math.ceil((t_end-t0)/dt) > max_steps:
        raise ValueError("Maximum number of steps exceeded.")
    t = t0+np.arange(math.floor((t_end-t0)/dt)+1)*dt
    roundoff = 64.0*sys.float_info.epsilon*max(1.0, abs(t_end))
    if t_end-t[-1] <= roundoff:
        t[-1] = t_end
    else:
        t = np.append(t, t_end)
    if np.any(np.diff(t) <= 0.0):
        raise ValueError("Step too small to advance time.")
    steps = len(t)-1
    x = np.zeros(len(t))
    x[0] = x0

    for i in range(steps):
        if key == 1:
            x[i+1] = euler1(f, t[i], x[i], t[i+1])
        elif key == 2:
            x[i+1] = heun1(f, t[i], x[i], t[i+1])
        else:
            x[i+1] = rk4_scalar(f, t[i], x[i], t[i+1])

    # One plain text file per run; a rerun of that method overwrites it.
    print(f"Method: {method}")
    print(f"Fixed step: {dt:.5f}")
    with open(filename, "w", encoding="utf-8") as data:
        if check_exact:
            # Valid for this equation even when t0 or x0 changes.
            exact = (t+1.0)**2+(x0-(t0+1.0)**2)*np.exp(t-t0)
            abs_error = np.abs(x-exact)
            print(f"{'t':>10} {'numerical x':>18} {'exact x':>18} {'abs error':>14}")
            data.write("# t x_numerical x_exact abs_error\n")
            for i in range(len(t)):
                print(f"{t[i]:10.5f} {x[i]:18.12f} {exact[i]:18.12f} "
                      f"{abs_error[i]:14.5e}")
                data.write(f"{t[i]:.16e} {x[i]:.16e} {exact[i]:.16e} "
                           f"{abs_error[i]:.16e}\n")
        else:
            print(f"{'t':>10} {'numerical x':>18}")
            data.write("# t x_numerical\n")
            for i in range(len(t)):
                print(f"{t[i]:10.5f} {x[i]:18.12f}")
                data.write(f"{t[i]:.16e} {x[i]:.16e}\n")
    print(f"Steps: {steps}")
    print(f"Data file: {filename}")

    # Dense exact curve and actual numerical points.
    plt.figure()
    if check_exact:
        t_plot = np.linspace(t0, t_end, 1000)
        x_plot = (t_plot+1.0)**2+(x0-(t0+1.0)**2)*np.exp(t_plot-t0)
        plt.plot(t_plot, x_plot, "k-", linewidth=1.2, label="Exact")
        plt.plot(t, x, "ko--", markersize=4, label=method)
    else:
        plt.plot(t, x, "ko-", markersize=4, label=method)
    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.title(f"{method}, step = {dt:g}")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


# ---------------- Equation and numerical methods -----------------------
def f(t, x):
    """Supply the derivative x' = f(t,x) here."""
    return x-t*t+1.0


def euler1(f, ti, xi, tf):
    """One Euler step: slope at the beginning of the interval."""
    h = tf-ti
    return xi+h*f(ti, xi)


def heun1(f, ti, xi, tf):
    """Euler predictor, then ONE trapezoidal corrector."""
    h = tf-ti
    slope1 = f(ti, xi)
    predicted = xi+h*slope1
    slope2 = f(tf, predicted)
    return xi+0.5*h*(slope1+slope2)


def rk4_scalar(f, ti, xi, tf):
    """Classical RK4. Each k includes h."""
    h = tf-ti
    k1 = h*f(ti, xi)
    k2 = h*f(ti+h/2.0, xi+k1/2.0)
    k3 = h*f(ti+h/2.0, xi+k2/2.0)
    k4 = h*f(tf, xi+k3)
    return xi+(k1+2.0*k2+2.0*k3+k4)/6.0


# Execute the driver after the equation and methods have been defined.
if __name__ == "__main__":
    main()