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

import sys
import advection_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> [ <idtype> <cfl> <boundary> <nx> ]"
    print "\n This code animates the advection equation.  The flux argument is"
    print "   flux     = 1 :   Upwind"
    print "              2 :   Centered"
    print "              3 :   Lax-Friedrichs"
    print "              4 :   Lax-Wendroff"
    print "              5 :   Beam-Warming"
    print " "
    print "   idtype   = 1 :  Gaussian"
    print "            = 2 :  Top hat"
    print " "
    print "   boundary = 1 :  Outflow"
    print "              2 :  Periodic"
    sys.exit(2)

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

if len(sys.argv) > 3:
    cfl_arg = float(sys.argv[3])
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) > 4:
    bc_arg = int(sys.argv[4])
else:
    bc_arg = 2

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) > 5:
    nx_arg = int(sys.argv[5])
else:
    nx_arg = 51



ng = 3
N=nx_arg+ng
w = wv.wave_system(-1,1,N,add_ghost_zones=True)

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])
else:
    w.set_u_n(wv.id_box,[1.0,0.2])

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

if flux_arg == 1:
    lab = 'Upwind'
elif flux_arg == 2:
    lab = 'Centered'
elif flux_arg == 3:
    lab = 'Lax-Friedrichs'
elif flux_arg == 4:
    lab = 'Lax-Wendroff'
elif flux_arg == 5:
    lab = 'Beam-Warming'

#-----------------------------------------------------------------
# animation
#-----------------------------------------------------------------
fig = plt.figure()
ax = fig.add_subplot(111,xlim=(-1,1), 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()