"""
Exercise E.24 from "A primer on..."
Write a test function for the forward Euler 
solver. Solving a system of ODEs with linear 
solution, which should be reproduced exactly
by most numerical solvers. 
"""

import numpy as np
import matplotlib.pyplot as plt


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 test_forward_euler():
    def f(t, u):
        v = u[0]
        w = u[1]
        dv = 3 + (3 + 4 * t - w)**3
        dw = 4 + (2 + 3 * t - v)**4
        return np.array([dv, dw])

    tol = 1e-10
    u0 = [2, 3]
    T = 5
    expected = [17, 23]
    N = 5
    t, u = forward_euler(f, u0, T, N)
    computed = u[-1]
    for e, c in zip(expected, computed):
        assert(abs(e-c) < tol)

test_forward_euler()