"""
Ex E.21 - implement the RK4 method as a class. We
let it be a subclass of the ForwardEuler class, and
inherit as much as possible.
"""
import numpy as np
import matplotlib.pyplot as plt


class ForwardEuler:
    """
    Class for solving an ODE,

      du/dt = f(u, t)

    by the ForwardEuler solver.

    Class attributes:
    t: array of time values
    u: array of solution values (at time points t)
    k: step number of the most recently computed solution
    f: callable object implementing f(u, t)
    dt: time step (assumed constant)
    """
    def __init__(self, f):
        if not callable(f):
            raise TypeError('f is %s, not a function' % type(f))
        self.f = f

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

    def solve(self, time_points):
        """Compute u for t values in time_points list."""
        self.t = np.asarray(time_points)
        self.u = np.zeros(len(time_points))
        # Assume self.t[0] corresponds to self.U0
        self.u[0] = self.U0

        for k in range(len(self.t)-1):
            self.k = k
            self.u[k+1] = self.advance()
        return self.u, self.t

    def advance(self):
        """Advance the solution one time step."""
        # Load attributes into local variables to
        # obtain a formula that is as close as possible
        # to the mathematical notation.
        u, f, k, t = self.u, self.f, self.k, self.t

        dt = t[k+1] - t[k]
        u_new = u[k] + dt*f(u[k], t[k])
        return u_new

class RungeKutta4(ForwardEuler):
    def advance(self):
        u, f, k, t = self.u, self.f, self.k, self.t
        dt = t[k+1] - t[k]

        K1 = dt*f(u[k], t[k])
        K2 = dt*f(u[k]+0.5*K1,t[k]+0.5*dt)
        K3 = dt*f(u[k]+0.5*K2,t[k]+0.5*dt)
        K4 = dt*f(u[k]+K3,t[k]+dt)

        return u[k]+1.0/6*(K1+2*K2+2*K3+K4)

def f(u,t):
    return u

U0 = 1
T = 4
n = 10

time = np.linspace(0,T,n+1)

solver = RungeKutta4(f)
solver.set_initial_condition(U0)
u,t = solver.solve(time)

plt.plot(t,u)
plt.show()