# --------------------/
%matplotlib inline
# --------------------/
import math
import numpy as np
import matplotlib.pyplot as plt

from pylab import *
from scipy import *
from ipywidgets import *

"""
Explanation: <h1> Lax–Wendroff Method

The scalar advection equation
\begin{equation}
u_t+au_x=0
\end{equation}
has the standard Lax–Wendroff method
\begin{equation}
U^{n+1}_j = U_j^n - \frac{ak}{2h}\left(U_{j+1}^n-U_{j-1}^n\right) + \frac{a^2k^2}{2h^2}\left(U^n_{j-1} -2 U^n-j + U^n_{j+1}\right).
\end{equation}

<h4> Essential Libraries </h4>
End of explanation
"""


# --------------------/
# domains and 
# parameters

c = 2.0  # velocity
a = -1.0
b = 1.0
T = 1.0

# --------------------/
# mesh and grid points

# nodes
n = 200

# tolerance
tol = 1e-5

h = (b - a) / (1.0 + n)    
k = 0.25 * h
m = int(round(T / k))
t = np.linspace(0, T, m)
x = np.linspace(a, b, n + 2)

tt, xx = meshgrid(t,x)
    
# --------------------/

if abs( k * m - T) > tol:
    print 'instability'

# --------------------/

# Courant-Friedrichs-Lewis
CFL = c * k / h  

# additional ghost cell
u = np.zeros(n + 3)
U = np.zeros((n + 2,m))

# true solution
f = lambda x: np.exp(-50*x**2)*np.cos(x)

# initial conditions u(x,t0)
u = f(x)

# periodic boundary conditions
u[0], u[-1] = u[-2], u[1]

# --------------------/
# Lax--Wendroff algorithm

for i in range(m):
    
    u[1:-1] = (u[1:-1] - 
               0.5 * CFL * (u[2:] - u[0:-2]) + 
               0.5 * CFL**2 * (u[0:-2] - 2*u[1:-1] + u[2:])
              )
    
    u[0], u[-1] = u[-2], u[1]
    
    U[:,i] = u
    

"""
Explanation: <h4> Lax–Wendroff Advection </h4>
End of explanation
"""


# --------------------/
# plots

def evolution(step):
    plt.figure(figsize=(5,5))
    plt.plot(x, U[:,step], lw=4, alpha=0.5, color='dodgerblue')
    plt.grid(color='lightgray', alpha=0.75)
    plt.xlim(x.min() - 0.125, x.max() + 0.125)
    plt.ylim(U[:,step].min() - 0.125, U[:,step].max() + 0.125)

# --------------------/
# interactive plot

time = widgets.IntSlider(min=1, max=m-1, description='step')
interact(evolution, step=time)

"""
Explanation: <h4> Plots
End of explanation
"""


