"""
Ex. E.2 from "A primer on..."
Solve the same model as in ex E.1, but now
using the class implementation of the forward Euler
method and a class representation of the right hand side 
function. 
"""


import numpy as np
import matplotlib.pyplot as plt


#class version of Forward Euler, copied here for convenience
class ForwardEuler_v0:
    def __init__(self, f):
        self.f = f  # , self.U0, self.T, self.N = f, U0, T, N

    def set_initial_condition(self, 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 / N
        self.t = np.zeros(N + 1)  # N steps ~ N+1 time points
        self.u = np.zeros(N + 1)

        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."""
        # Create local variables to get rid of "self." in
        # the numerical formula
        u, dt, f, n, t = self.u, self.dt, self.f, self.n, self.t

        return u[n] + dt * f(t[n], u[n])


#class representation of the right hand side function,
#including the growth rate as a flexible parameter
class RHS:
    def __init__(self, r):
        self.r = r

    def __call__(self, t, u):
        return self.r * u

# create an instance of the class, which represents
# the right hand side function f = 0.1*u:
problem = RHS(0.1) 

#create a solver instance:
solver = ForwardEuler_v0(problem)


solver.set_initial_condition(u0 = 0.2)
T = 20
N = 4

t, u = solver.solve((0,T), N)
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()
plt.show()








plt.plot(t, u, label = f'dt={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()
plt.show()