"""
Ex E.1 from "A primer on... " Solve a simple ODE using the 
function implementation of the Forward Euler method described
earlier. 
"""

import numpy as np
import matplotlib.pyplot as plt

#copy the function code here for convenience:
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)
    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
        u[n + 1] = u[n] + dt * f(t[n], u[n])

    return t, u

#a) Write the ODE on the form u' = f(t, u)
# u' = u/10, u0 = 0.2

#b) implement the right hand side as a function:
def rhs(t, u):
    return 0.1 * u

#c)solve the ODE with dt = 5 and T = 20, i.e. N = 4
N = 4
T = 20
t, u = forward_euler(rhs, u0 = 0.2, T = T, N = N)

#d) plot the numerical solution and the analytical
plt.plot(t, u, label=f'$\Delta t = $ {T/N}')
t_fine = np.linspace(0, T, 200)
plt.plot(t_fine, 0.2 * np.exp(0.1 * t_fine), label = 'Exact')
plt.legend()
#e) save to file: 
plt.savefig('forward_euler.pdf')
plt.show()

#f) Explore the effect of reducing the time step:
for N in [4, 8, 16, 32, 64, 128]:
    t, u = forward_euler(rhs, u0 = 0.2, T = T, N = N)
    plt.plot(t, u, label=f'$\Delta t = $ {T/N}')

plt.plot(t_fine, 0.2 * np.exp(0.1 * t_fine), label = 'Exact')
plt.legend()
plt.show()