import numpy as np
%matplotlib inline

gamma = 1.4
M = 1.5

# Right state
rho_r = 1.0
u_r   = 0.0  # This must be zero.
p_r   = 1.0

# Left state
p_l   = 5.0
u_l   = 0.0    # This must be zero.

c_r = np.sqrt(gamma*p_r/rho_r)  # Right state sound speed
s = M * c_r                     # Shock speed
mu = 2*(M**2-1)/(gamma+1)/M

# Find state to right of discontinuity via Rankine-Hugoniot conditions
p   = p_r * ((2*M**2-1)*gamma+1)/(gamma+1)
u   = mu*c_r
rho_rs = rho_r*M/(M-mu)

# Check that the 1-wave is a rarefaction
if p_l<p:
    print "Warning: the states provided do not lead to the expected solution structure."

# Find remaining densities
rho_l = 4*gamma*p_l/((gamma-1)**2*u**2) * (1-(p/p_l)**((gamma-1)/(2*gamma)))**2
rho_ls= (p/p_l)**(1/gamma) * rho_l

from clawpack import pyclaw
from clawpack import riemann
from clawpack.riemann.euler_with_efix_1D_constants import density, momentum, energy, num_eqn

# Specify equations and boundary conditions
solver = pyclaw.ClawSolver1D(riemann.euler_with_efix_1D)
solver.bc_lower[0] = pyclaw.BC.extrap
solver.bc_upper[0] = pyclaw.BC.extrap

# Grid
x = pyclaw.Dimension('x',-1.,1.,1000)
domain = pyclaw.Domain([x])
state = pyclaw.State(domain,num_eqn)
state.problem_data['gamma'] = gamma
xc = state.grid.p_centers[0]

# Initial values
state.q[density,:] = rho_l*(xc<0) + rho_r*(xc>=0)
state.q[momentum,:] = rho_l*u_l*(xc<0) + rho_r*u_r*(xc>=0)
E_r = 0.5*rho_r * u_r**2 + p_r/(gamma-1)
E_l = 0.5*rho_l * u_l**2 + p_l/(gamma-1)
state.q[energy, :] = E_l*(xc<0) + E_r*(xc>=0)

claw = pyclaw.Controller()
claw.tfinal = 1.
claw.solution = pyclaw.Solution(state,domain)
claw.solver = solver
claw.num_output_times = 10
claw.keep_copy = True

claw.run()

from matplotlib import animation
from clawpack.visclaw.JSAnimation import IPython_display
import matplotlib.pyplot as plt

# Animated plot of the solution density
fig = plt.figure(figsize=[8,4])
ymax = max(rho_l, rho_ls, rho_rs, rho_r)
ax = plt.axes(xlim=(xc[0], xc[-1]), ylim=(0, ymax+0.5))
line, = ax.plot([], [], lw=1)

def fplot(i):
    frame = claw.frames[i]
    q = frame.q[0,:]
    line.set_data(xc, q)
    ax.set_title('Density at t= '+ str(frame.t))
    return line,

animation.FuncAnimation(fig, fplot, frames=len(claw.frames))