"""
Implementation of the ForwardEuler method as a function.
"""
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 f(t, u):
    return np.array([u[1], -u[0]])


U0 = [1,0]
T = 20
N = 2000
t, u = forward_euler(f, U0, T, N)
plt.plot(t, u[:,0], label= 'u')
plt.plot(t, u[:,1], label= 'v')
plt.xlabel('t')
plt.legend()
plt.show()