"""Exer E.21. Function implementation of the RungeKutta4 method.
The code is based on the file ForwardEuler_func.py from the book,
but with some of the test code removed for simplicity."""

import numpy as np

def ForwardEuler(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[k] is the solution at time t[k]
    u[0] = U0
    t[0] = 0
    dt = T/float(n)
    for k in range(n):
        t[k+1] = t[k] + dt
        u[k+1] = u[k] + dt*f(u[k], t[k])
    return u, t

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[k] is the solution at time t[k]
    u[0] = U0
    t[0] = 0
    dt = T/float(n)
    for k in range(n):
        K1 = dt*f(u[k],t[k])
        K2 = dt*f(u[k]+0.5*K1,t[k]+0.5*dt)
        K3 = dt*f(u[k]+0.5*K2,t[k]+0.5*dt)
        K4 = dt*f(u[k]+K3,t[k]+dt)

        t[k+1] = t[k] + dt
        u[k+1] = u[k]+1.0/6*(K1+2*K2+2*K3+K4)
        #u[k+1] = u[k] + dt*f(u[k], t[k])
    return u, t



# Problem: u'=u
def f(u, t):
    return u

u, t = ForwardEuler(f, U0=1, T=4, n=20)
u2, t2 = RK4(f, U0=1, T=4, n=5)

# Compare numerical solution and exact solution in a plot
import matplotlib.pyplot as plt
u_exact = np.exp(t)
plt.plot(t, u, t2, u2, t, u_exact)
plt.xlabel('t')
plt.ylabel('u')
plt.legend(['FE', 'RK4','exact']),
plt.title("Solution of the ODE u'=u, u(0)=1")
plt.show()