#!/usr/bin/env python3

"""Ordinary Differential Equations"""

import numpy as np


def mk_phi_euler(f):
    """ Return a Phi computed for the Euler Method. """
    def phi_euler(x, y, h):
        return f(x[-1], y[-1])

    return phi_euler


def mk_phi_euler_mod(f):
    """ Return a Phi computed for the Modified Euler (Collatz) Method. """
    def phi_euler_mod(x, y, h):
        return f(x[-1] + 0.5*h, y[-1] + 0.5*h*f(x[-1], y[-1]))

    return phi_euler_mod


def mk_phi_heun(f):
    """ Return a Phi computed for the Heun Method. """
    def phi_heun(x, y, h):
        return (f(x[-1], y[-1]) + f(x[-1] + h, y[-1] + h * f(x[-1], y[-1])))/2
    
    return phi_heun

def mk_phi_rk4(f):
    """ Return a Phi computed for the 4th order Runge-Kutta Method. """
    def phi_rk4(x, y, h):
        k1 = f(x[-1], y[-1])
        k2 = f(x[-1] + 0.5*h, y[-1] + 0.5*k1)
        k3 = f(x[-1] + 0.5*h, y[-1] + 0.5*h*k2)
        k4 = f(x[-1] + h, y[-1] + h*k3)

        return (k1 + 2*k2 + 2*k3 + k4)/6

    return phi_rk4

def integrator(x, y_0, phi):
    x = np.asarray(x)
    if isinstance(y_0, (list,np.ndarray)):
        y_0 = np.asarray(y_0)
    else:
        y_0 = np.array([ y_0 ])

    N = len(x)
    M = len(y_0)

    y = np.zeros((N,M), dtype=np.float64)
    y[0] = y_0

    h = x[1]-x[0]
    for i, x_i in enumerate(x):
        # Skip the first iteration
        if not i:
            continue

        y[i] = y[i-1] + h*phi(x[:i], y[:i], h)

    return y



def test_integrator():
    np.random.seed(0)

    def ODEF(x, y):
        return y - x**2 + 1
    
    def ODEF_sol(x):
        return (x+1)**2 - 0.5*np.exp(x)


    test_func = ODEF
    exact_sol = ODEF_sol

    N = 2e1
    M = 1

    # Domain on which to do the integration
    x = np.linspace(0, 1, N)

    # Generate Initial Value
    y_0 = 0.5

    schemes = [ 
            ["Euler", mk_phi_euler],
            ["Collatz", mk_phi_euler_mod],
            ["Heun", mk_phi_heun],
            ["RK4", mk_phi_rk4],
            ]

    # Show Plot
    from matplotlib import pyplot
    pyplot.subplots()


    for name, func in schemes:
        pyplot.plot(x, integrator(x, y_0, func(test_func)), '--o', label=name)

    pyplot.plot(x, exact_sol(x), '-', label="Exact Solution")
    pyplot.xlabel("x")
    pyplot.ylabel("y")

    pyplot.legend()

    pyplot.show()

if __name__ == "__main__":
    test_integrator()