#!/usr/bin/env python3

"""Ordinary Differential Equations"""

import numpy as np

# Integrations Schemes #
########################

# Single Step
#------------
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

# Multi Step
#-----------
def mk_phi_AB3(f, phi_short = None):
    steps = 3
    def phi_AB3(x, y, h):
        if len(x) < steps:
            if phi_short is None:
                raise ValueError("This function needs at least {} steps, {} given.".format(steps, len(x)))
            else:
                return phi_short(x, y, h)
        else:
            return ( 23*f(x[-1], y[-1]) - 16*f(x[-2], y[-2]) + 5*f(x[-3], y[-3]) )/12

    return phi_AB3

def mk_phi_AB4(f, phi_short = None):
    steps = 4
    def phi_AB4(x, y, h):
        if len(x) < steps:
            if phi_short is None:
                raise ValueError("This function needs at least {} steps, {} given.".format(steps, len(x)))
            else:
                return phi_short(x, y, h)
        else:
            return ( 55*f(x[-1], y[-1]) - 59*f(x[-2], y[-2]) + 37*f(x[-3], y[-3]) -9*f(x[-4],y[-4]))/24

    return phi_AB4


# Integrator #
##############
def integrator(x, y_0, phi, y_i = None ):
    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

    j = 0
    if y_i is not None:
        for i, _ in enumerate(y_i):
            y[i+1] = y_i[i]
        j = i

    h = x[1]-x[0]
    for i in range(j+1, N):
        y[i] = y[i-1] + h*phi(x[:i], y[:i], h)

    return y


# Math Test Functions #
#######################
def ODEF(x, y):
    return y - x**2 + 1

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


# Testing #
###########
def test_integrator_singlestep():

    test_func = ODEF
    exact_sol = ODEF_sol

    N = 2e1

    # 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()

def test_integrator_multistep():
    test_func = ODEF
    exact_sol = ODEF_sol

    N = 2e1

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

    # Calculate Initial Values using RK4
    y_0 = 0.5
    y_init = integrator(x[:4], y_0, mk_phi_rk4(test_func))

    #Name, func, steps
    multi_schemes = [ 
            ["AB3", mk_phi_AB3, 3],
            ["AB4", mk_phi_AB4, 4],
        ]

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

    # Exact Solution
    pyplot.plot(x, exact_sol(x), '-', label="Exact Solution")

    # Plot the schemes
    for name, func, steps in multi_schemes:
        try:
            pyplot.plot(x, integrator(x, y_0, func(test_func), y_init[:steps]), '--o', label=name)
        except Exception as e:
            print(e)
            pass

    pyplot.xlabel("x")
    pyplot.ylabel("y")

    pyplot.legend()

    pyplot.show()

def plot_integration_errors():
    test_func = ODEF
    exact_sol = ODEF_sol

    # Domain on which to do the integration
    x = np.arange(0, 2.001, step=0.02)

    # Calculate Initial Values using RK4
    y_0 = 0.5
    y_init = integrator(x[:4], y_0, mk_phi_rk4(test_func))

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

        ]
    multi_schemes = [ 
            ["AB3", mk_phi_AB3, 3],
            ["AB4", mk_phi_AB4, 4],
        ]

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

    # Pre calculate the exact solution
    exact = exact_sol(x)
    exact = np.reshape(exact, (-1,1))

    # Plot Single Step Schemes
    for name, func in single_schemes:
        try:
            pyplot.plot(x, np.abs(exact - integrator(x, y_0, func(test_func))), '--o', label=name)
        except Exception as e:
            print(e)
            pass

    # Plot Multi Step Schemes
    for name, func, steps in multi_schemes:
        try:
            pyplot.plot(x, np.abs(exact - integrator(x, y_0, func(test_func), y_init[:steps])), '--o', label=name)
        except Exception as e:
            print(e)
            pass

    pyplot.xlabel("x")
    pyplot.ylabel("absolute error $|\\bar{y} - y|$")
    pyplot.title("absolute error between an approach and the exact solution")

    pyplot.legend()

    pyplot.show()


def accuracy_per_stepsize( debug=True ):
    stepsizes = np.asarray([1e-1, 1e-2, 1e-3, 1e-5])

    # define the problem to be solved
    test_func = lambda x,y: x
    exact_sol = lambda x: np.exp(x)

    # set domain and intial value
    x_min = 0
    x_max = 2
    y_0 = 1

    exact_end = exact_sol(x_max)

    # define the schemes to use
    schemes = [
            ["Euler", mk_phi_euler, 1],
            ["Collatz", mk_phi_euler_mod, 1],
            ["Heun", mk_phi_heun, 1],
            ["RK4", mk_phi_rk4, 1],
            ["AB3", mk_phi_AB3, 3],
            ["AB4", mk_phi_AB4, 4],
        ]

    # get max steps needed for the multistep schemes
    max_steps = 0
    for _, _, steps in schemes:
        if steps > max_steps:
            max_steps = steps

    # pre calculate for multistep integrations
    y_init = np.zeros((len(stepsizes),1))# Note the dimensionality of y_0
    max_steps += 1
    for i, h in enumerate(stepsizes):
        x = x_min + np.linspace(0, h * max_steps, max_steps, True)
        # Note the dimensionality of y_0
        # TODO: fix up this ugliness
        z = integrator(x[:4], y_0, mk_phi_rk4(test_func))
        y_init[i] = z[1,:]

    # plot schemes
    from matplotlib import pyplot
    pyplot.subplots()
    for name, scheme, steps in schemes:
        diffs = np.zeros(len(stepsizes))
        if debug:
            print("Calculating {}".format(name))

        for i, h in enumerate(stepsizes):
            N = int(abs(x_max-x_min)/h)
            x = np.linspace(x_min, x_max, N, True)
            y = integrator(x, y_0, scheme(test_func), y_init[:steps])
            diffs[i] = np.abs(exact_end - [-1])
        
        pyplot.plot(stepsizes, diffs, '--o', label=name)

    pyplot.xlabel('Stepsize $h$')
    pyplot.ylabel('Absolute Error')
    pyplot.title('Absolute error at the end of integration')

    pyplot.legend()

    pyplot.show()


def pendulum_problem():
    pass

if __name__ == "__main__":
    np.random.seed(0)
    test_integrator_singlestep()
    #test_integrator_multistep()

    #plot_integration_errors()
    #accuracy_per_stepsize()