## Code from Section 1.1 of "Solving Ordinary Differential 
## Equations in Python", to be used in Problem 1

import numpy as np

def forward_euler(f, u0, T, N):
    """Solve u¡¯=f(t, u), 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
    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



## Code from Section 1.4 of "Solving Ordinary Differential 
## Equations in Python", to be used in Problem 2

import numpy as np

class ForwardEuler:
	def __init__(self, f):
		self.f = lambda t, u: np.asarray(f(t, u), float)

	def set_initial_condition(self, u0):
		if np.isscalar(u0): # scalar ODE
			self.neq = 1 # no of equations
			u0 = float(u0)
		else: # system of ODEs
			self.neq = u0.size # no of equations
			u0 = np.asarray(u0)
		self.u0 = u0

	def solve(self, t_span, N):
		"""Compute solution for
		t_span[0] <= t <= t_span[1],
		using N steps."""
		t0, T = t_span
		self.dt = (T - t0) / N
		self.t = np.zeros(N + 1)
		if self.neq == 1:
			self.u = np.zeros(N + 1)
		else:
			self.u = np.zeros((N + 1, self.neq))
		msg = "Please set initial condition before calling solve"
		assert hasattr(self, "u0"), msg
	
		self.t[0] = t0
		self.u[0] = self.u0

		for n in range(N):
			self.n = n
			self.t[n + 1] = self.t[n] + self.dt
			self.u[n + 1] = self.advance()
		return self.t, self.u

	def advance(self):
		"""Advance the solution one time step."""
		u, dt, f, n, t = self.u, self.dt, self.f, self.n, self.t
		return u[n] + dt * f(t[n], u[n])



## Code from Section 2.2 of "Solving Ordinary Differential 
## Equations in Python", to be used in Problems 3,4 and 5.

import numpy as np

class ODESolver:
    def __init__(self, f):
        # Wrap user¡¯s f in a new function that always
        # converts list/tuple to array (or let array be array)
        self.model = f
        self.f = lambda t,u: np.asarray(f(t,u), float)

    def set_initial_condition(self, u0):
        if np.isscalar(u0): # scalar ODE
            self.neq = 1 # no of equations
            self.u0 = float(u0)
        else: # system of ODEs
            u0 = np.asarray(u0)
            self.neq = u0.size # no of equations
            self.u0 = u0

    def solve(self, t_span, N):
        """Compute solution for t_span[0] <= t <= t_span[1],
        using N steps."""
        t0, T = t_span
        self.dt = (T - t0) / N
        self.t = np.zeros(N + 1) # N steps ~ N+1 time points
        if self.neq == 1:
            self.u = np.zeros(N + 1)
        else:
            self.u = np.zeros((N + 1, self.neq))
        msg = "Please set initial condition before calling solve"
        assert hasattr(self, "u0"), msg
        self.t[0] = t0
        self.u[0] = self.u0
        for n in range(N):
            self.n = n
            self.t[n + 1] = self.t[n] + self.dt
            self.u[n + 1] = self.advance()
        return self.t, self.u

    def advance(self):
        raise NotImplementedError("Advance method is not implemented in the base class")
        
class ForwardEuler(ODESolver):
    def advance(self):
        u, f, n, t = self.u, self.f, self.n, self.t
        dt = self.dt
        return u[n] + dt * f(t[n], u[n])
        
class ExplicitMidpoint(ODESolver):
	def advance(self):
		u, f, n, t = self.u, self.f, self.n, self.t
		dt = self.dt
		dt2 = dt / 2.0
		k1 = f(t[n], u[n])
		k2 = f(t[n] + dt2, u[n] + dt2 * k1)
		return u[n] + dt * k2

class RungeKutta4(ODESolver):
	def advance(self):
		u, f, n, t = self.u, self.f, self.n, self.t
		dt = self.dt
		dt2 = dt / 2.0
		k1 = f(t[n], u[n],)
		k2 = f(t[n] + dt2, u[n] + dt2 * k1, )
		k3 = f(t[n] + dt2, u[n] + dt2 * k2, )
		k4 = f(t[n] + dt, u[n] + dt * k3, )
		return u[n] + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)