"""Exer E.22. Class implementation of the RungeKutta4 method,
as a subclass of the ForwardEuler_v1 class. The code is
based on the file ForwardEuler.py from the book, but with
some test code and class versions removed for simplicity."""

import numpy as np

class ForwardEuler_v1:
    """
    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)
    n: number of time steps in the simulation
    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, U0, T, n):
        if not callable(f):
            raise TypeError('f is %s, not a function' % type(f))

        self.f, self.U0, self.T, self.n = f, U0, T, n
        self.dt = T/float(n)
        self.u = np.zeros(self.n+1)
        self.t = np.zeros(self.n+1)

    def solve(self):
        """Compute solution for 0 <= t <= T."""
        self.u[0] = float(self.U0)
        self.t[0] = float(0)

        for k in range(self.n):
            self.k = k
            self.t[k+1] = self.t[k] + self.dt
            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, dt, f, k, t = \
           self.u, self.dt, self.f, self.k, self.t

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

class RK4(ForwardEuler_v1):
    def __init__(self, f, U0, T, n):
        ForwardEuler_v1.__init__(self, f, U0, T, n)

    def advance(self):
        u, dt, f, k, t = \
           self.u, self.dt, self.f, self.k, self.t
        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)
        u_new = u[k]+1.0/6*(K1+2*K2+2*K3+K4)
        #u_new = u[k] + dt*f(u[k], t[k])
        return u_new



def f(u, t):
    return u



solver = RK4(f,U0=1,T=4,n=20)
u,t = solver.solve()

import matplotlib.pyplot as plt
plt.plot(t,u)
plt.xlabel('t')
plt.ylabel('u')
plt.title("Solution of the ODE u'=u, u(0)=1")
plt.show()