# Import some libraries
import numpy as np
from scipy.integrate import quad

# Set advection speed (c in PDE)
c = 10.0

# Set spatial domain 
xmin = 0
xmax = 15

# Define q0(x), initial function to be advected
def q0(x):
    qq = np.exp(-(x-2.5)*(x-2.5))
    return qq

# Create a function for the REA algorithm (parameter: number of cells and final time)
def REA(ncells,tfinal):
    
    # Create grid edges, space step and init time
    x = np.linspace(xmin,xmax,ncells+1) # grid edges
    dx = x[1] - x[0]
    time = 0
    
    # Define time step size obeying CFL condition
    dt = dx/(2.0*c)    
    
    # Initialize solution q(x,t): q
    q = np.zeros(ncells)
    
    # Calculate cell averages at time 0 (using q0 function)
    Qavg = np.zeros(ncells)
    for i in range(ncells):
        # Integrate q0 function in the grid intervals
        integral = quad(q0, x[i], x[i+1]) 
        Qavg[i] = integral[0]/dx
    
    # Initialiaze REA algorithm 
    while (time <= tfinal):
        # RECONSTRUCT STEP
        q = 1.0*Qavg
        
        # EVOLVE and AVERAGE STEP (move by dt and average)
        time = time + dt
        for i in range(ncells):
            if (c>=0):
                Qavg[i] = (c*dt*q[i-1] + (dx-c*dt)*q[i])/dx
            else:
                Qavg[i] = (c*dt*q[i+1] + (dx-c*dt)*q[i])/dx
        
    # If final time close adjust dt to match desired final time
    if (time > tfinal):
        dt = tfinal - (time - dt)
        for i in range(ncells):
            Qavg[i] = (c*dt*q[i-1] + (dx-c*dt)*q[i])/dx
        q = Qavg
    
    return (x,q)

# Function to calculate exact solution for comparison
def q_exact(tfinal):
    # Calculate dense x
    xx = np.linspace(xmin,xmax,1000)
    # Obtain exact solution
    qexact = q0(xx-c*tfinal)
    return (xx,qexact)

#Import plotting libraries
%matplotlib inline
import matplotlib.pyplot as plt
from pylab import *
from IPython.html.widgets import interact
from ipywidgets import StaticInteract, RangeWidget

# Create plotting function for REA and exact solution
def plot_REA(ncells, tfinal):
    # Create Figure to plot
    fig = plt.figure()
    ax = fig.add_subplot(111) 
    
    # Call REA algorithm and exact solution
    x, qREA = REA(ncells,tfinal)
    xx, qexact = q_exact(tfinal)
    
    # Transform REA into piecewise function for pretty plotting
    qREA_pp = 0*xx
    for i in range(len(x)-1):
        for j in range(len(xx)):
            if (xx[j] >= x[i] and xx[j] < x[i+1]): 
                qREA_pp[j] = qREA[i]
    
    # Plot solutions
    ax.plot(xx,qREA_pp,'-b',linewidth=2)
    ax.plot(xx,qexact,'-g',linewidth=2)
    ax.axis([min(xx),max(xx),0,1.2])
    return fig

# Create animated solution
StaticInteract(plot_REA, ncells=RangeWidget(60,760,100), tfinal=RangeWidget(0,0.8,0.05))