from IPython.core.display import HTML
css_file = './custom.css'
HTML(open(css_file, "r").read())

gamma = 1.4
gamma1 = gamma - 1.

def step_Euler_radial(solver,state,dt):
    """
    Geometric source terms for Euler equations with radial symmetry.
    Integrated using a 2-stage, 2nd-order Runge-Kutta method.
    This is a Clawpack-style source term routine.
    """
    
    dt2 = dt/2.
    ndim = 2

    aux=state.aux
    q = state.q

    rad = aux[0,:,:]

    rho = q[0,:,:]
    u   = q[1,:,:]/rho
    v   = q[2,:,:]/rho
    press  = gamma1 * (q[3,:,:] - 0.5*rho*(u**2 + v**2))

    qstar = np.empty(q.shape)

    qstar[0,:,:] = q[0,:,:] - dt2*(ndim-1)/rad * q[2,:,:]
    qstar[1,:,:] = q[1,:,:] - dt2*(ndim-1)/rad * rho*u*v
    qstar[2,:,:] = q[2,:,:] - dt2*(ndim-1)/rad * rho*v*v
    qstar[3,:,:] = q[3,:,:] - dt2*(ndim-1)/rad * v * (q[3,:,:] + press)

    rho = qstar[0,:,:]
    u   = qstar[1,:,:]/rho
    v   = qstar[2,:,:]/rho
    press  = gamma1 * (qstar[3,:,:] - 0.5*rho*(u**2 + v**2))

    q[0,:,:] = q[0,:,:] - dt*(ndim-1)/rad * qstar[2,:,:]
    q[1,:,:] = q[1,:,:] - dt*(ndim-1)/rad * rho*u*v
    q[2,:,:] = q[2,:,:] - dt*(ndim-1)/rad * rho*v*v
    q[3,:,:] = q[3,:,:] - dt*(ndim-1)/rad * v * (qstar[3,:,:] + press)

def qinit(state,z0=0.5,r0=0.,bubble_radius=0.2,rhoin=0.1,pinf=5.):
    grid = state.grid

    # Background state
    rhoout = 1.
    pout   = 1.
    
    # Bubble state
    pin    = 1.

    # Shock state
    zshock = 0.2
    rinf = (gamma1 + pinf*(gamma+1.))/ ((gamma+1.) + gamma1*pinf)
    vinf = 1./np.sqrt(gamma) * (pinf - 1.) / np.sqrt(0.5*((gamma+1.)/gamma) * pinf+0.5*gamma1/gamma)
    einf = 0.5*rinf*vinf**2 + pinf/gamma1
    
    z = grid.z.centers
    r = grid.r.centers
    R,Z = np.meshgrid(r,z)
    rad = np.sqrt((Z-z0)**2 + (R-r0)**2)

    state.q[0,:,:] = rinf*(Z<zshock) + rhoin*(rad<=bubble_radius) + rhoout*(rad>bubble_radius)*(Z>=zshock)
    state.q[1,:,:] = rinf*vinf*(Z<zshock)
    state.q[2,:,:] = 0.
    state.q[3,:,:] = einf*(Z<zshock) + (pin*(rad<=bubble_radius) + pout*(rad>bubble_radius)*(Z>=zshock))/gamma1
    state.q[4,:,:] = 1.*(rad<=bubble_radius)

def shockbc(state,dim,t,qbc,num_ghost):
    """
    Incoming shock at left boundary.
    """
    p = 5.
    rho = (gamma1 + p*(gamma+1.))/ ((gamma+1.) + gamma1*p)
    v = 1./np.sqrt(gamma) * (p - 1.) / np.sqrt(0.5*((gamma+1.)/gamma) * p+0.5*gamma1/gamma)
    e = 0.5*rho*v**2 + p/gamma1

    for i in xrange(num_ghost):
        qbc[0,i,...] = rho
        qbc[1,i,...] = rho*v
        qbc[2,i,...] = 0.
        qbc[3,i,...] = e
        qbc[4,i,...] = 0.

from clawpack import riemann
from clawpack import pyclaw
import numpy as np

solver = pyclaw.ClawSolver2D(riemann.euler_5wave_2D)
solver.step_source=step_Euler_radial

# Boundary conditions
solver.bc_lower[0]=pyclaw.BC.custom
solver.bc_upper[0]=pyclaw.BC.extrap
solver.bc_lower[1]=pyclaw.BC.wall
solver.bc_upper[1]=pyclaw.BC.extrap
solver.user_bc_lower=shockbc

# We must also set boundary conditions for the auxiliary variable (r)
# But in this case the values don't matter
solver.aux_bc_lower[0]=pyclaw.BC.extrap
solver.aux_bc_upper[0]=pyclaw.BC.extrap
solver.aux_bc_lower[1]=pyclaw.BC.extrap
solver.aux_bc_upper[1]=pyclaw.BC.extrap

# Initialize domain
mx=320; my=80
z = pyclaw.Dimension('z',0.0,2.0,mx)
r = pyclaw.Dimension('r',0.0,0.5,my)
domain = pyclaw.Domain([z,r])
num_aux=1
state = pyclaw.State(domain,solver.num_eqn,num_aux)

state.problem_data['gamma']= gamma
state.problem_data['gamma1']= gamma1

qinit(state)

rc = state.grid.r.centers
for j,rcoord in enumerate(rc):
    state.aux[0,:,j] = rcoord

solver.cfl_max = 0.5
solver.cfl_desired = 0.45
solver.source_split = 1        # Use first-order splitting for the source term

claw = pyclaw.Controller()
claw.keep_copy = True
claw.output_format = None
claw.tfinal = 0.6
claw.solution = pyclaw.Solution(state,domain)
claw.solver = solver
claw.num_output_times = 60

claw.run()

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

fig = plt.figure(figsize=[8,4])

frame = claw.frames[0]
rho = frame.q[0,:,:]

x, y = frame.state.grid.c_centers    

# This essentially does a pcolor plot, but it returns the appropriate object
# for use in animation.  See http://matplotlib.org/examples/pylab_examples/pcolor_demo.html.
# Note that it's necessary to transpose the data array because of the way imshow works.
im = plt.imshow(rho.T, cmap='Greys', vmin=rho.min(), vmax=rho.max(),
           extent=[x.min(), x.max(), y.min(), y.max()],
           interpolation='nearest', origin='lower')

def fplot(frame_number):
    frame = claw.frames[frame_number]
    rho = frame.q[0,:,:]
    im.set_data(rho.T)
    return im,

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

from clawpack.pyclaw import examples

examples.