########################################################################
#
#   Remember to:   export QT_API=pyside
#
########################################################################

import sys
import burgers_sys as wv
import matplotlib
matplotlib.use('Qt4Agg')
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation

try:
    flux_arg = int(sys.argv[1])
except:
    print "Usage: ",sys.argv[0],"<flux> [ <slope> <idtype> <cfl> <boundary> <nx> ]"
    print "\n This code animates the advection equation.  The flux argument is"
    print "   flux     = 1 :   Godunov"
    print "              2 :   Upwind"
    print " "
    print "   slope    = 1 :   Constant"
    print "            = 2 :   Centered"
    print "            = 3 :   minmod"
    print "            = 4 :   MC"
    print "            = 5 :   Superbee"
    print " "
    print "   idtype   = 1 :  Gaussian"
    print "            = 2 :  Top hat"
    print "            = 3 :  Sine"
    print " "
    print "   boundary = 1 :  Outflow"
    print "              2 :  Periodic"
    sys.exit(2)

if len(sys.argv) > 2:
    slope_arg = int(sys.argv[2])
else:
    slope_arg = 1    # default value

if len(sys.argv) > 3:
    id_arg = int(sys.argv[3])
else:
    id_arg = 1    # default value

if len(sys.argv) > 4:
    cfl_arg = float(sys.argv[4])
else:
    cfl_arg = 0.5    # default value

if cfl_arg > 1.0 or cfl_arg <= 0.0:
    print "cfl is unphysical.  cfl = ",cfl_arg,". Die here."
    sys.exit(2)

if len(sys.argv) > 5:
    bc_arg = int(sys.argv[5])
else:
    bc_arg = 1

if bc_arg != 1 and bc_arg != 2:
    print "boundary must be 1 or 2.  boundary = ",bc_arg,". Die here."
    sys.exit(2)

if len(sys.argv) > 6:
    nx_arg = int(sys.argv[6])
else:
    nx_arg = 51


gxmin = -1.0
gxmax =  1.0

ng = 3
N=nx_arg+ng
w = wv.wave_system(gxmin,gxmax,N,add_ghost_zones=True)

# ode solvers: Huen, RK2, Euler
w.set_ode_type('Euler')

if id_arg == 1:
#w.set_u_n(wv.id_step,[1.0,0.2])
    w.set_u_n(wv.id_gaussian,[1.0,0.0,0.25])
elif id_arg == 2:
    w.set_u_n(wv.id_box,[1.0,0.2])
elif id_arg == 3:
    w.set_u_n(wv.id_sine,[0.5,0.5])

if slope_arg == 1:
    w.set_slope_type('Constant')
    slab = '+ Constant'
elif slope_arg == 2:
    w.set_slope_type('Centered')
    slab = '+ Centered'
elif slope_arg == 3:
    w.set_slope_type('Minmod')
    slab = '+ Minmod'
elif slope_arg == 4:
    w.set_slope_type('MC')
    slab = '+ MC'
elif slope_arg == 5:
    w.set_slope_type('Superbee')
    slab = '+ Superbee'
else:
    print "Unknown slope type " + slope_arg

w.set_boundary_condition(bc_arg)
w.set_flux_type(flux_arg)
w.set_cfl(cfl_arg)

if flux_arg == 1:
    lab = 'Godunov' + slab
elif flux_arg == 2:
    lab = 'Upwind'
elif flux_arg == 3:
    lab = 'Lax-Friedrichs' + slab
else:
    lab = 'Unknown'

#-----------------------------------------------------------------
# animation
#-----------------------------------------------------------------
fig = plt.figure()
ax = fig.add_subplot(111,xlim=(gxmin,gxmax), ylim=(-0.2,1.2), 
     autoscale_on=False)


line, = ax.plot(w.x(),w.u_n(),'o-',label=lab)
ax.legend()

def init():
    line.set_data([],[])
    return line,

def animate(i):
    w.step(1)
    line.set_data(w.x(),w.u_n())
    return line,
    

ani = animation.FuncAnimation(fig, animate, init_func=init, frames=5, blit=True )

plt.show()