"""
Exercise E.21 from "A primer on ..."
Implementation of the Runge Kutta 4 method as a function,
based on the ForwardEuler function from the lecture notes.
"""
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 = f(u[n], t[n])
        k2 = f(u[n]+0.5*dt*k1,t[n]+0.5*dt)
        k3 = f(u[n]+0.5*dt*k2,t[n]+0.5*dt)
        k4 = f(u[n]+dt*k3,t[n]+dt)
        u[n+1] = u[n] + (dt/6)*(k1+2*k2+2*k3+k4)

    return u, t

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

U0 = 1
T = 3
N = 30
u, t = RK4(f, U0, T, N)
plt.plot(t,u)
plt.show()

"""
Terminal> python RK4_func.py
(output is a plot)
"""