"""
Exercise E.22 from "A primer on..."
Implement the Runge-Kutta 4 method as a class, using the ForwardEuler
class as the starting point. The implementation here is
based on the ForwardEuler class from chapter 1 of "Solving ODEs in Python".
"""

import numpy as np

class ForwardEuler:
    def __init__(self, f):
        """Initialize the right-hand side function f """
        self.f = lambda t, u: np.asarray(f(t, u), float)

    def set_initial_condition(self, u0):
        """Store the initial condition as 
        an instance attribute. 
        """
        self.u0 = np.asarray(u0)
        self.neq = self.u0.size

    def solve(self, t_span, N):
        """Compute solution for
        t_span[0] <= t <= t_span[1],
        using N steps.
        Returns the solution and the 
        time points as arrays. 
        """
        t0, T = t_span
        self.dt = (T - t0) / N
        self.t = np.zeros(N + 1)
        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])



"""
RK4 as a subclass of ForwardEuler. The only difference between the
two methods is the update formula used to compute u[n + 1]. This
is isolated in the advance method, so this is the only method that
needs to be reimplemented in the subclass. All other methods
are inherited from ForwardEuler. 
"""

class RK4(ForwardEuler):
    def advance(self):
        u, dt, f, n, t = self.u, self.dt, self.f, self.n, self.t
        K1 = dt * f(t[n], u[n])
        K2 = dt * f(t[n] + 0.5 * dt, u[n] + 0.5 * K1)
        K3 = dt * f(t[n] + 0.5 * dt, u[n] + 0.5 * K2)
        K4 = dt * f(t[n] + dt, u[n] + K3)
        return u[n] + (1 / 6) * (K1 + 2 * K2 + 2 * K3 + K4)


"""
We use a logistic growth problem to demonstrate the
use of the classes and the accuracy. 
"""

class Logistic:
    def __init__(self, alpha, R):
        self.alpha, self.R = alpha, R

    def __call__(self, t, u):
        return self.alpha * u * (1 - u / self.R)


if __name__ == '__main__':
    """
    Demonstrate how the class is used,
    by solving the logistic growth problem.
    To test the class for a system of ODEs,
    run pendulum.py
    """
    import matplotlib.pyplot as plt

    problem = Logistic(alpha=0.2, R=1.0)
    fe_solver = ForwardEuler(problem)
    u0 = 0.1
    fe_solver.set_initial_condition(u0)
    T = 40
    t, u = fe_solver.solve(t_span=(0, T), N=400)

    rk4_solver = RK4(problem)
    rk4_solver.set_initial_condition(u0)
    t1, u1 = rk4_solver.solve(t_span=(0, T), N=10)

    plt.plot(t, u,label = "FE")
    plt.plot(t1,u1,label='RK4')
    plt.title('Logistic growth, Forward Euler and RK4')
    plt.xlabel('t')
    plt.ylabel('u')
    plt.legend()
    plt.show()

    """
    Terminal> python RK4_class.py 
    (output is a plot)
    """