#!/usr/bin/env python3

import numpy as np

def pdeHyperbolic(a, x, t, f, g, dtype=np.float64):
    """ Solve a Hyperbolic Partial Differential using finite differences. """
    
    m = len(x) # Amount of objects to track
    n = len(t) # Length of Time Vector

    # Determine stepsizes
    h = x[0] - x[1]
    k = t[0] - t[1]

    λ_sq = (a*k/h)**2

    # Create array to hold the solution
    w = np.zeros((n,m), dtype=dtype)

    # Create finite difference matrix 
    A = np.diag(m*[2*(1 - λ_sq)], k=0) + np.diag((m-1)*[λ_sq], k=-1) + np.diag((m-1)*[λ_sq], k=1)

    # Initialise first two timesteps
    w[0] = f(x, t[0])
    w[1] = A@w[0]/2 + k*g(x, t[0])
   
    # Calculate for following timesteps
    for j in range(2, n-1):
        w[j] = A@w[j-1] - w[j-2]

    return w


def test_pdeHyperbolic_case1(x_steps=1e2, t_steps=1e2, max_x=1, max_t=1):
    
    a = 1 # from the Schroedinger Equation

    # Setup spatial and time grids
    x = np.linspace(0, max_x, x_steps)
    t = np.linspace(0, max_t, t_steps)

    # Boundary conditions
    def f(x,t):
        return np.sin(2*np.pi*x)

    def g(x,t):
        return 2*np.pi*np.sin(2*np.pi*x)    

    # Solve it
    sol = pdeHyperbolic(a, x, t, f, g)

    # Plot it with the exact solution
    exact_f = lambda x,t: np.sin(2*np.pi*x)*(np.cos(2*np.pi*t) + np.sin(2*np.pi*t))

    plot_animation(x, sol, func=exact_f)


def test_pdeHyperbolic_case2(x_steps=2e2, t_steps=4e2, max_x=1, max_t=1):
    
    a = 1 # from the Schroedinger Equation

    # Setup spatial and time grids
    x = np.linspace(0, max_x, x_steps)
    t = np.linspace(0, max_t, t_steps)

    # Boundary conditions
    def f(x,t):
            return 2*(x < 0.5) -1

    def g(x,t):
            return 0

    # Solve it
    sol = pdeHyperbolic(a, x, t, f, g)

    plot_animation(x, sol)



def plot_animation(x, sol, func=None, interval = 10):
    from matplotlib import pyplot
    from matplotlib import animation as anim
    fig, _ = pyplot.subplots()
    
    def animate(i):
        pyplot.clf()
        pyplot.grid()
        pyplot.ylim(-1.5, 1.5)
        pyplot.title('t = {}/{}'.format(i, len(sol)))
        #if func is not None:
        #    pyplot.plot(x, func(x, sol[i]), label='exact')
        pyplot.plot(x, sol[i,:], label="iter")

    frames = np.arange(0, len(sol))

    try:
        myAnim = anim.FuncAnimation(fig, animate, frames, interval = interval )

        pyplot.legend()
        pyplot.show()
    except AttributeError:
        # This final error is fugly
        pass



if __name__ == "__main__":
    test_pdeHyperbolic_case1()
    #test_pdeHyperbolic_case2()