"""
Exercise E.22 from "A primer on..."
Implement the Runge Kutta 4 method as a class. We implement it
as a sub-class of the ForwardEuler_v2 from the ODE lecture notes.
"""

import numpy as np
import matplotlib.pyplot as plt

class ForwardEuler_v2:
    def __init__(self, f):
        self.f = f

    def set_initial_condition(self,U0):
        self.U0 = float(U0)

    def solve(self, time_points):
        """Compute solution for array of time points"""
        self.t = np.asarray(time_points)
        N = len(self.t)
        self.u = np.zeros(N)
        self.u[0] = self.U0

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

    def advance(self):
        """Advance the solution one time step."""
        # Create local variables to get rid of "self." in
        # the numerical formula
        u, f, n, t = self.u, self.f, self.n, self.t
        #dt is not necessarily constant:
        dt = t[n+1]-t[n]
        unew = u[n] + dt*f(u[n], t[n])
        return unew

class RK4(ForwardEuler_v2):
    def advance(self):
        u, f, n, t = self.u, self.f, self.n, self.t
        dt = t[n+1]-t[n]
        k1 = f(u[n], t[n])
        k2 = f(u[n]+0.5*dt*k1,t[n]+0.5*dt)
        k3 = f(u[n]+0.5*dt*k2,t[n]+0.5*dt)
        k4 = f(u[n]+dt*k3,t[n]+dt)
        unew = u[n] + (dt/6)*(k1+2*k2+2*k3+k4)
        return unew


if __name__ == '__main__':
    """
    Demonstrate how the class the class is used,
    by solving the logistic growth problem.
    """

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

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

    problem = Logistic(0.2, 1, 0.1)

    #choose very few time steps, to see the difference between FE and RK4:
    time = np.linspace(0,40,11)
    method0 = ForwardEuler_v2(problem)
    method0.set_initial_condition(problem.U0)
    u0, t0 = method0.solve(time)
    plt.plot(t0,u0)


    method1 = RK4(problem)
    method1.set_initial_condition(problem.U0)
    u1, t1 = method1.solve(time)
    plt.plot(t1,u1)

    plt.show()


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