import numpy as np
from numpy.linalg import norm, solve
import matplotlib.pyplot as plt

plt.rcParams['axes.grid'] = True
plt.ion()

def explicit_euler(y0, t0, T, f, Nmax):
    ys = [y0]
    ts = [t0]
    dt = (T - t0)/Nmax
    while(ts[-1] < T):
        t, y = ts[-1], ys[-1]  
        ys.append(y + dt*f(t, y))
        ts.append(t + dt)
    return (np.array(ts), np.array(ys))
# end of explicit_euler

def implicit_euler(y0, t0, T, f, Nmax):
    ys = [y0]
    ts = [t0]
    dt = (T - t0)/Nmax
    while(ts[-1] < T):
        t, y = ts[-1], ys[-1]  
        ys.append(y + dt*f(t, y))
        ts.append(t + dt)
    return (np.array(ts), np.array(ys))
# end of implicit_euler

def heun(y0, t0, T, f, Nmax):
    ys = [y0]
    ts = [t0]
    dt = (T - t0)/Nmax
    while(ts[-1] < T):
        t, y = ts[-1], ys[-1]
        k1 = f(t,y)
        k2 = f(t+dt, y+dt*k1)
        ys.append(y + 0.5*dt*(k1+k2))
        ts.append(t + dt)
    return (np.array(ts), np.array(ys))
# end of heun

class ExplicitRungeKutta:
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

    def __call__(self, y0, t0, T, f, Nmax):
        # Extract Butcher table
        a, b, c = self.a, self.b, self.c
        
        # Stages
        s = len(b)
        ks = [np.zeros_like(y0, dtype=np.double) for s in range(s)]

        # Start time-stepping
        ys = [y0]
        ts = [t0]
        dt = (T - t0)/Nmax
        
        while(ts[-1] < T):
            t, y = ts[-1], ys[-1]
            
            # Compute stages derivatives k_j
            for j in range(s):
                t_j = t + c[j]*dt
                dY_j = np.zeros_like(y, dtype=np.double)
                for l in range(j):
                    dY_j += dt*a[j,l]*ks[l]

                ks[j] = f(t_j, y + dY_j)
                
            # Compute next time-step
            dy = np.zeros_like(y, dtype=np.double)
            for j in range(s):
                dy += dt*b[j]*ks[j]
            
            ys.append(y + dy)
            ts.append(t + dt)
            
        return (np.array(ts), np.array(ys))
# end of class ExplicitRungeKutta

def example1():
    t0, T = 0, 1
    y0 = 1
    lam = 1
    Nmax = 4

    # rhs of IVP
    f = lambda t,y: lam*y

    # Compute numerical solution using Euler
    ts, ys_eul = explicit_euler(y0, t0, T, f, Nmax)

    # Exact solution to compare against
    y_ex = lambda t: y0*np.exp(lam*(t-t0))
    ys_ex = y_ex(ts)

    # Plot it
    plt.plot(ts, ys_ex)
    plt.plot(ts, ys_eul, 'ro-')
    plt.legend(["$y_{ex}$", "y" ])
# end of example1

def example_lotka_volterra():
    # Define rhs
    def lotka_volterra(t, y):
        # Set parameters
        alpha, beta, delta, gamma = 2, 1, 0.5, 1  
        # Define rhs of ODE
        dy = np.array([alpha*y[0]-beta*y[0]*y[1], 
                    delta*y[0]*y[1]-gamma*y[1]])
        return dy

    t0, T = 0, 20            # Integration interval
    y0 = np.array([2, 0.5])  # Initital values

    # Solve the equation
    tau = 0.002
    Nmax = int(20/tau)
    print("Nmax = {:4}".format(Nmax))
    ts, ys_eul = explicit_euler(y0, t0, T, lotka_volterra, Nmax)

    # Plot results
    plt.plot(ts, ys_eul)
    plt.xlabel('t')
    plt.legend(['$y_0(t)$ - Euler', '$y_1(t)$ - Euler'],
            loc="upper right" )
# end of example_lotka_volterra
    

def exercise2a():
    t0, T = 0, 1
    y0 = 1
    lam = 1
    Nmax = 4

    # rhs of IVP
    f = lambda t,y: lam*y

    # Compute numerical solution using Euler and Heun
    ts, ys_eul = explicit_euler(y0, t0, T, f, Nmax)
    ts, ys_heun = heun(y0, t0, T, f, Nmax)

    # Exact solution to compare against
    y_ex = lambda t: y0*np.exp(lam*(t-t0))
    ys_ex = y_ex(ts)

    # Plot it
    plt.plot(ts, ys_ex)
    plt.plot(ts, ys_eul, 'ro-')
    plt.plot(ts, ys_heun, 'b+-')
    plt.legend(["$y_{ex}$", "$y$ Euler", "$y$ Heun" ])
# end of exercise2a

def exercise2b():
    Nmax_list = [4, 8, 16, 32, 64, 128]
    error_study(y0, t0, T, f, Nmax_list, heun, y_ex)
# end of exercise2b

def exercise1():
    def error_study(y0, t0, T, f, Nmax_list, solver, y_ex):
        max_errs = []
        for Nmax in Nmax_list:
            ts, ys = solver(y0, t0, T, f, Nmax)
            ys_ex = y_ex(ts)
            errors = ys - ys_ex
            max_errs.append(np.abs(errors).max())
            print("For Nmax = {:3}, max ||y(t_i) - y_i||= {:.3e}".format(Nmax,max_errs[-1]))
        print("The computed error reduction rates are")
        max_errs = np.array(max_errs)
        print(max_errs[:-1]/max_errs[1:])

    Nmax_list = [4, 8, 16, 32, 64, 128]
    error_study(y0, t0, T, f, Nmax_list, explicit_euler, y_ex)
# end of exercise1


def compute_eoc(y0, t0, T, f, Nmax_list, solver, y_ex):
    errs = [ ]
    for Nmax in Nmax_list:
        ts, ys = solver(y0, t0, T, f, Nmax)
        ys_ex = y_ex(ts)
        errs.append(np.abs(ys - ys_ex).max())
        print("For Nmax = {:3}, max ||y(t_i) - y_i||= {:.3e}".format(Nmax,errs[-1]))

    errs = np.array(errs)
    Nmax_list = np.array(Nmax_list)
    dts = (T-t0)/Nmax_list

    eocs = np.log(errs[1:]/errs[:-1])/np.log(dts[1:]/dts[:-1])

    # Insert inf at beginning of eoc such that errs and eoc have same length 
    eocs = np.insert(eocs, 0, np.Inf)
    return errs, eocs
# end of compute_eoc

def exercise_eoc_study():
    t0, T = 0, 1
    y0 = 1
    lam = 1

    # rhs of IVP
    f = lambda t,y: lam*y

    # Exact solution to compare against
    y_ex = lambda t: y0*np.exp(lam*(t-t0))

    Nmax_list = [4, 8, 16, 32, 64, 128]

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, explicit_euler, y_ex)
    print(eocs)

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, heun, y_ex)
    print(eocs)
# end of exercise_eoc_study


def exercise_van_der_pol():
    # Define the ODE
    def f(t, y):
        mu = 2
        dy = np.array([y[1],
                    mu*(1-y[0]**2)*y[1]-y[0] ])
        return dy

    # Set initial time, stop time and initial value
    t0, T - 0, 20
    y0 = np.array([2,0])

    # Solve the equation using Euler and plot
    tau = 0.1
    Nmax = int(20/tau)
    print("Nmax = {:4}".format(Nmax))
    ts, ys_eul  = explicit_euler(y0, t0, T, f, Nmax)
    plt.plot(ts,ys_eul);

    # Solve the equation using Heun
    tau = 0.1
    Nmax = int(20/tau)
    print("Nmax = {:4}".format(Nmax))
    ts, ys_heun = heun(y0, t0, T, f, Nmax)
    plt.plot(ts,ys_heun);

    plt.xlabel('x')
    plt.title('Van der Pols ligning')
    plt.legend(['y1 - Euler','y2 - Euler', 'y1 - Heun','y2 - Heun'],loc='upper right');
# end of exercise_van_der_pol

def exercise_rk2():
    # Define Butcher table for improved Euler
    a = np.array([[0,   0],
                [0.5, 0]])
    b = np.array([0, 1])
    c = np.array([0, 0.5])

    # Create a new Runge Kutta solver
    rk2 = Explicit_Runge_Kutta(a, b, c)

    t0, T = 0, 1
    y0 = 1
    lam = 1
    Nmax = 10

    # rhs of IVP
    f = lambda t,y: lam*y

    # the solver can be simply called as before, namely as function:
    ts, ys = rk2(y0, t0, T, f, Nmax)

    plt.figure()
    plt.plot(ts, ys, "c--o", label="$y_{\mathrm{heun}}$")

    # Exact solution to compare against
    y_ex = lambda t: y0*np.exp(lam*(t-t0))

    # Plot the exact solution (will appear in the plot above)
    plt.plot(ts, y_ex(ts), "m-", label="$y_{\mathrm{ex}}$")
    plt.legend()

    # Run an EOC test
    Nmax_list = [4, 8, 16, 32, 64, 128]

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, rk2, y_ex)
    print(errs)
    print(eocs)

    # Do a pretty print of the tables using panda

    import pandas as pd
    from IPython.display import display

    table = pd.DataFrame({'Error': errs, 'EOC' : eocs})
    display(table)
# end of exercise_rk2:


def exercise_rkm3_heun():
    # Define Butcher table for Heun's 3rd order method
    a = np.array([[0, 0, 0],
                [1.0/3.0, 0, 0],
                [0, 2.0/3.0, 0]])

    b = np.array([1.0/4.0, 0, 3.0/4.0])

    c = np.array([0, 
                1.0/3.0, 
                2.0/3.0])

    # Define rkm
    rk3_heun = ExplicitRungeKutta(a, b, c)

    # Time interval
    t0, T = 0, 1

    # Inital data
    y0 = -0.5

    # rhs of IVP
    f = lambda t,y: np.exp(y)*(1+t)

    # Exact solution to compare against
    y_ex = lambda t: -np.log(np.exp(0.5) - t - t**2/2)

    Nmax = 100

    ts, ys = rk3_heun(y0, t0, T, f, Nmax)

    plt.plot(ts, ys)
    plt.plot(ts, y_ex(ts))

    plt.legend(["y", "y_ex"])

    # EOC test
    Nmax_list = [4, 8, 16, 32, 64, 128]

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, rk3_heun, y_ex)
    print(errs)
    print(eocs)
# end of exercise_rkm3_heun

def exercise_rkm3_kutta():
    # Define Butcher table for Kutta's method
    a = np.array([[0, 0, 0],
                [1.0/2.0, 0, 0],
                [-1, 2.0, 0]])

    b = np.array([1.0/6.0, 2.0/3, 1.0/6])

    c = np.array([0, 
                1.0/2.0, 
                1.0])

    # Define rkm
    rk3_kutta = ExplicitRungeKutta(a, b, c)

    # Time interval
    t0, T = 0, 1

    # Inital data
    y0 = -0.5

    # rhs of IVP
    f = lambda t,y: np.exp(y)*(1+t)

    # Exact solution to compare against
    y_ex = lambda t: -np.log(np.exp(0.5) - t - t**2/2)

    Nmax = 100

    ts, ys = rk3_kutta(y0, t0, T, f, Nmax)

    plt.plot(ts, ys)
    plt.plot(ts, y_ex(ts))

    plt.legend(["y", "y_ex"])
    # EOC test
    Nmax_list = [4, 8, 16, 32, 64, 128]

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, rk3_kutta, y_ex)
    print(errs)
    print(eocs)
# end of exercise_rkm3_kutta

def exercise_rkm4():
    # Define Butcher table for classical Runge-Kutta method
    a = np.array([[0, 0, 0, 0],
                [1.0/2.0, 0, 0, 0],
                [0, 1.0/2.0, 0, 0],
                [0, 0, 1, 0]])

    b = np.array([1/6, 1/3, 1/3, 1/6])

    c = np.array([0, 
                1.0/2.0, 
                1.0/2.0,
                1.0])

    # Define rkm
    rk4 = ExplicitRungeKutta(a, b, c)

    # Time interval
    t0, T = 0, 1

    # Inital data
    y0 = -0.5

    # rhs of IVP
    f = lambda t,y: np.exp(y)*(1+t)

    # Exact solution to compare against
    y_ex = lambda t: -np.log(np.exp(0.5) - t - t**2/2)

    Nmax = 100

    ts, ys = rk4(y0, t0, T, f, Nmax)

    plt.plot(ts, ys)
    plt.plot(ts, y_ex(ts))

    plt.legend(["y", "y_ex"])
    # EOC test
    Nmax_list = [4, 8, 16, 32, 64, 128, 256, 512, 1024]

    errs, eocs = compute_eoc(y0, t0, T, f, Nmax_list, rk4, y_ex)
    print(errs)
    print(eocs)
# end of exercise_rkm4
 

def exercise_sir():
    class SIR:
        def __init__(self, beta, gamma):
            self.beta = beta
            self.gamma = gamma

        def __call__(self, t, y):
            return np.array([-beta*y[0]*y[1],
                            beta*y[0]*y[1] - gamma*y[1],
                            gamma*y[1]])
    # Model parameters
    beta = 10./(40*8*24)
    gamma = 3./(15*24)
    sir = SIR(beta, gamma)

    # Define Butcher table for classical Runge-Kutta method
    a = np.array([[0, 0, 0, 0],
                [1.0/2.0, 0, 0, 0],
                [0, 1.0/2.0, 0, 0],
                [0, 0, 1, 0]])

    b = np.array([1/6, 1/3, 1/3, 1/6])

    c = np.array([0, 
                1.0/2.0, 
                1.0/2.0,
                1.0])
    # Define RKM
    rkm = ExplicitRungeKutta(a, b, c)

    # Initial data
    t0 = 0
    y0 = np.array([50, 10, 0])

    # Simulate
    dt = 0.1 # 6 min
    D = 30   # Simulate for D days
    Nmax = int(D*24/dt) # Corresponding no of hours
    T = Nmax*dt # End time in minutes

    ts, ys = rkm(y0, t0, T, sir, Nmax)
    plt.plot(ts, ys)

    # Plot total sum of individuals
    plt.plot(ts, ys[:,0]+ys[:,1]+ys[:,2])

    plt.legend(["S", "I", "R", "S+I+R"])
# end of exercise_sir

def example_stiff_problem():
    plt.rcParams['figure.figsize'] = (16.0, 12.0)
    t0, T = 0, 1
    y0 = 1
    lams = [-10, -50, -250]

    fig, axes = plt.subplots(3,3)
    fig.tight_layout(pad=3.0)

    for i in range(len(lams)):
        lam = lams[i]
        tau_l = 2/abs(lam)
        taus = [0.1*tau_l, tau_l, 1.1*tau_l]

        # rhs of IVP
        f = lambda t,y: lam*y

        # Exact solution to compare against
        y_ex = lambda t: y0*np.exp(lam*(t-t0))

        # Compute solution for different time step size
        for j in range(len(taus)):
            tau = taus[j]
            Nmax = int(1/tau)
            ts, ys = explicit_euler(y0, t0, T, f, Nmax)
            ys_ex = y_ex(ts)
            axes[i,j].set_title(f"$\lambda = {lam}$, $\\tau = {tau:0.2e}$")
            axes[i,j].plot(ts, ys, "ro-")
            axes[i,j].plot(ts, ys_ex)
            axes[i,j].legend(["$y_{\mathrm{FE}}$", "$y_{\mathrm{ex}}$"])
# end of example_stiff_problem

def example_stiff_problem_BE():
    plt.rcParams['figure.figsize'] = (16.0, 12.0)
    t0, T = 0, 1
    y0 = 1
    lams = [-10, -50, -250]

    fig, axes = plt.subplots(3,3)
    fig.tight_layout(pad=3.0)

    for i in range(len(lams)):
        lam = lams[i]
        tau_l = 2/abs(lam)
        taus = [0.1*tau_l, tau_l, 1.1*tau_l]

        # rhs of IVP
        f = lambda t,y: lam*y

        # Exact solution to compare against
        y_ex = lambda t: y0*np.exp(lam*(t-t0))

        # Compute solution for different time step size
        for j in range(len(taus)):
            tau = taus[j]
            Nmax = int(1/tau)
            ts, ys = explicit_euler(y0, t0, T, f, Nmax)
            ys_ex = y_ex(ts)
            axes[i,j].set_title(f"$\lambda = {lam}$, $\\tau = {tau:0.2e}$")
            axes[i,j].plot(ts, ys, "ro-")
            axes[i,j].plot(ts, ys_ex)
            axes[i,j].legend(["$y_{\mathrm{FE}}$", "$y_{\mathrm{ex}}$"])
# end of example_stiff_problem_BE

class EmbeddedExplicitRungeKutta:
    def __init__(self, a, b, c,  bhat=None, order=None):
        self.a = a
        self.b = b
        self.c = c
        self.bhat = bhat
        self.order = order


    def __call__(self, y0, t0, T, f, Nmax, tol=1e-3, store_rejected=False):
        # Extract Butcher table
        a, b, c, bhat, order = self.a, self.b, self.c, self.bhat, self.order

        # Some parameters controlling the time-step choice 
        # Machine precision
        eps = 1e-15
        fac = 0.8
        facmax = 5.0
        facmin = 0.1
        err  = 0
        
        # Stages
        s = len(b)
        ks = [np.zeros_like(y0, dtype=np.double) for s in range(s)]

        # Start time-stepping
        ys = [y0]
        ts = [t0]
        
        # Store rejected time-steps
        ts_rej = []
        ys_rej = []
        dt = (T - t0)/Nmax
        # Counting steps
        N = 0
        N_rej = 0
        
        while(ts[-1] < T and N < Nmax):
            t, y = ts[-1], ys[-1]
            N += 1
            
#             print("---------------------------------------------------")
#             print(f"Compute tentative step N = {N} starting from t = {t} with dt = {dt}")
            
            # Compute stages derivatives k_j
            for j in range(s):
                t_j = t + c[j]*dt
                dY_j = np.zeros_like(y, dtype=np.double)
                for l in range(j):
                    dY_j += a[j,l]*ks[l]

                ks[j] = f(t_j, y + dt*dY_j)
                
            # Compute next time-step
            dy = np.zeros_like(y, dtype=np.double)
            for j in range(s):
                dy += b[j]*ks[j]

            if bhat is None:
                ys.append(y + dt*dy)
                ts.append(t + dt)
            else:
                dyhat = np.zeros_like(y, dtype=np.double)
                for j in range(s):
                    dyhat += bhat[j]*ks[j]

                # Error estimate
#                 err = max(dt*norm(dy - dyhat), norm(y)*eps)
                err = dt*norm(dy - dyhat)

                # Accept time-step
#                 if True:
                if err <= tol:
                    ys.append(y + dt*dyhat)
                    ts.append(t + dt)
                else:
                    print(f"Step is rejected at t = {t} with err = {err}")
                    N_rej += 1
                    ys_rej.append(y + dt*dyhat)
                    ts_rej.append(t + dt)
                
                # New step size
                dt = min(dt*min(facmax, max(facmin, fac*(tol/err)**(1/(order)))),abs(T-t))
#                 dt = 0.8*(tol/err)**(1/(order))*dt
#                 print(f"New dt = {dt}")
#                 print(f"Error is err = {err}")
        
        print(f"Finishing time-stepping reaching t = {ts[-1]} with final time T = {T}")
        print(f"Used {N} steps out of {Nmax} with {N_rej} being rejected")
          
        if store_rejected:   
            return (np.array(ts), np.array(ys), np.array(ts_rej), np.array(ys_rej))
        else:
            return (np.array(ts), np.array(ys))

# end of class EmbeddedExplicitRungeKutta


# Define a number of Butcher tables 
butcher_tables = {
    "explicit_euler":
    {
        "a": np.array([[0]]),
        "b": np.array([1]),
        "c": np.array([0]),
        "order": 1
    },
    # Also known as explicit midpoint
    "rk2":
    {
        "a": np.array([[0, 0],
                       [0.5, 0]]),
        "b": np.array([0, 1]),
        "c": np.array([0,
                       0.5]),
        "order": 2
    },
    # Also known as explicit trapezoidal rule
    "heun":
    {
        "a": np.array([[0, 0],
                       [1, 0]]),
        "b": np.array([0.5, 0.5]),
        "c": np.array([0,
                       1]),
        "order": 2
    },
    "rkm3_heun":
    {
        # Define Butcher
        "a": np.array([[0, 0, 0],
                       [1.0/3.0, 0, 0],
                       [0, 2.0/3.0, 0]]),
        "b": np.array([1.0/4.0, 0, 3.0/4.0]),
        "c": np.array([0,
                       1.0/3.0,
                       2.0/3.0]),
        "order": 3
    },
    "rkm3_kutta":
    {
        "a": np.array([[0, 0, 0],
                       [1.0/2.0, 0, 0],
                       [-1, 2.0, 0]]),
        "b": np.array([1.0/6.0, 2.0/3, 1.0/6]),
        "c": np.array([0,
                       1.0/2.0,
                       1.0]),
        "order": 3
    },
    "rkm4":
    {
        "a": np.array([[0, 0, 0, 0],
                       [1.0/2.0, 0, 0, 0],
                       [0, 1.0/2.0, 0, 0],
                       [0, 0, 1, 0]]),
        "b": np.array([1/6, 1/3, 1/3, 1/6]),
        "c": np.array([0,
                       1.0/2.0,
                       1.0/2.0,
                       1.0]),
        "order": 4
    }
}

embedded_rkm_tables = {
    # Also known as explicit trapezoidal rule
    "euler_heun":
    {
        "a": np.array([[0, 0],
                       [1., 0]]),
        "b"   : np.array([1., 0]),
        "bhat": np.array([0.5, 0.5]),
        "c": np.array([0,
                       1]),
        "order": 2
    },
    "fehlberg" :
     {
        "a": np.array([[0.0, 0,   0,   0,   0],
                       [1/2, 0,   0,   0,   0],
                       [0, 1/2,   0,   0,   0],
                       [0,   0,   1,   0,   0],
                       [1/6, 1/3, 1/3, 1/6, 0],
                        ]),
        "b"   : np.array([1/6, 1/3, 1/3, 0, 1/6]),
        "bhat": np.array([1/6, 1/3, 1/3, 1/6, 0]),
        "c": np.array([0.,
                       1/2,
                       1/2,
                       1,
                       1]),
        "order": 4
    },
}


if __name__ == "__main__":
    # Test code
    print("Testing ODE module")