%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
plt.rcParams['figure.figsize'] = 10, 10

#!/usr/bin/env python
# encoding: utf-8

"""
Solve the Euler equations of compressible fluid dynamics.
"""
from clawpack import pyclaw
from clawpack import riemann

Nx=200
Ny=200

solver = pyclaw.ClawSolver2D(riemann.euler_4wave_2D)
solver.all_bcs = pyclaw.BC.extrap

domain = pyclaw.Domain([0.,0.],[1.,1.],[Nx,Ny])
solution = pyclaw.Solution(solver.num_eqn,domain)
gamma = 1.4
solution.problem_data['gamma']  = gamma
solver.dimensional_split = False
solver.transverse_waves = 2

# Set initial data
xx,yy = domain.grid.p_centers
l = xx<0.8; r = xx>=0.8; b = yy<0.8; t = yy>=0.8
solution.q[0,...] = 1.5*r*t + 0.532258064516129*l*t + 0.137992831541219*l*b + 0.532258064516129*r*b
u = 0.*r*t + 1.206045378311055*l*t + 1.206045378311055*l*b + 0.*r*b
v = 0.*r*t + 0.*l*t + 1.206045378311055*l*b + 1.206045378311055*r*b
p = 1.5*r*t + 0.3*l*t + 0.029032258064516*l*b + 0.3*r*b
solution.q[1,...] = solution.q[0,...] * u
solution.q[2,...] = solution.q[0,...] * v
solution.q[3,...] = 0.5*solution.q[0,...]*(u**2+v**2) + p/(gamma-1.)

def display_q(q):
    x = np.linspace(0, 1, Nx)
    y = np.linspace(0, 1, Ny)
    x,y = np.meshgrid(x,y)
    z = q[0,:,:]
    cont = ax.contourf(x, y, z, 500)
    return cont

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(0, 1))
display_q(solution.q)

claw = pyclaw.Controller()
claw.keep_copy=True
claw.tfinal = 0.8
claw.num_output_times = 10
claw.solution = solution
claw.solver = solver

status = claw.run()

VIDEO_TAG = """<video controls>
 <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
 Your browser does not support the video tag.
</video>"""

def anim_to_html(anim):
    if not hasattr(anim, '_encoded_video'):
        with open('animation.mp4','w') as f:
            anim.save(f.name, fps=1, extra_args=['-vcodec', 'libx264'])
            video = open(f.name, "rb").read()
        anim._encoded_video = video.encode("base64")
    
    return VIDEO_TAG.format(anim._encoded_video)

from IPython.display import HTML

def display_animation(anim):
    plt.close(anim._fig)
    return HTML(anim_to_html(anim))

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(0, 1))

def animate(i): 
    display_q(claw.frames[i].q)
    
anim = animation.FuncAnimation(fig, animate, frames=len(claw.frames), blit=False)
print "created animation"
display_animation(anim)