"""
Exercise E.30 from "A primer on..."
Implement the 2nd order Runge-Kutta method,
often called the explicit midpoint method ot
the modified Euler method. The method can be
written as a function, based on the Forward Euler function. 
"""

import matplotlib.pyplot as plt
import numpy as np


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

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

    for n in range(N):
        t[n + 1] = t[n] + dt
        u[n + 1] = u[n] + dt * f(t[n], u[n])

    return t, u


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

    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)
        u[n + 1] = u[n] + k2
    return t, u



# solve a system of 2 ODEs with known solution:
def f(t, u):
    return np.array([u[1], -u[0]])


U0 = [0,1]
T = 20
N = 100
t1, u1 = forward_euler(f, U0, T, N)
plt.plot(t1, u1[:,0], label= 'FE: u')
plt.plot(t1, u1[:,1], label= 'FE: v')

t2, u2 = explicit_midpoint(f, U0, T, N)
plt.plot(t2, u2[:,0], label= 'EM: u')
plt.plot(t2, u2[:,1], label= 'EM: v')

plt.plot(t1, np.sin(t1),label='sin(t)')
plt.plot(t1, np.cos(t1), label= 'cos(t)')

plt.xlabel('t')
plt.legend()
plt.show()

"""
Terminal> python RungeKutta2_func.py
"""