"""
Ex. 3.21 from "A primer on..."
Implement the Runge Kutta 4th order method as a function, 
using the forward_euler function as starting point. The only 
necessary change is the update formula for the solution, where
the single line of the FE method is replaced by five lines
to compute the stage derivatives (k1-k4) and update the solution.
"""

import numpy as np
import matplotlib.pyplot as plt

def RK4(f, u0, T, N):
    """Solve u'=f(u,t), u(0)=U0, with n steps until t=T."""
    t = np.zeros(N + 1)
    u = np.zeros(N + 1)  # u[n] is the solution at time t[n]

    u[0] = u0
    t[0] = 0
    dt = T / N

    for n in range(N):
        t[n + 1] = t[n] + dt
        k1 = dt * f(t[n], u[n])
        k2 = dt * f(t[n] + 0.5 * dt, u[n] + 0.5 * k1)
        k3 = dt * f(t[n] + 0.5 * dt, u[n] + 0.5 * k2)
        k4 = dt * f(t[n] + dt, u[n] + k3)
        u[n + 1] = u[n] + (1/6) * (k1 + 2 * k2 + 2 * k3 + k4)

    return t, u


# demo: solve u' = u, u(0) = 1
def f(t, u):
    return u


U0 = 1
T = 4
N = 20
t, u = RK4(f, U0, T, N)
plt.plot(t, u, label=f'$\Delta t$ = {T/N}')
plt.plot(t, np.exp(t), label='Exact')
plt.title('Exp. growth, Runge-Kutta 4')
plt.xlabel('t')
plt.ylabel('u')
plt.legend()
plt.show()