"""
    Quick plot of the Carrier-Greenspan outputs

"""
import anuga.utilities.plot_utils as util
from matplotlib import pyplot as pyplot
from carrier_greenspan_analytic import *

p_st = util.get_output('carrier_greenspan.sww')
p2_st=util.get_centroids(p_st)

v = p2_st.y[10]
v2=(p2_st.y==v)

W1, P1, Z1, H1, U1 = analytic_cg(p2_st.x[v2], p2_st.time[55], h0=5e2, L=5e4, a=1.0, Tp=900.0)
W2, P2, Z2, H2, U2 = analytic_cg(p2_st.x[v2], p2_st.time[56], h0=5e2, L=5e4, a=1.0, Tp=900.0)
W3, P3, Z3, H3, U3 = analytic_cg(p2_st.x[v2], p2_st.time[57], h0=5e2, L=5e4, a=1.0, Tp=900.0)
W4, P4, Z4, H4, U4 = analytic_cg(p2_st.x[v2], p2_st.time[58], h0=5e2, L=5e4, a=1.0, Tp=900.0)
W5, P5, Z5, H5, U5 = analytic_cg(p2_st.x[v2], p2_st.time[59], h0=5e2, L=5e4, a=1.0, Tp=900.0)
W6, P6, Z6, H6, U6 = analytic_cg(p2_st.x[v2], p2_st.time[60], h0=5e2, L=5e4, a=1.0, Tp=900.0)

#Plot the stages##############################################################
pyplot.clf()
pyplot.ion()
pyplot.plot(p2_st.x[v2], p2_st.stage[55,v2], 'b.', label='numerical') # 1*T/6
pyplot.plot(p2_st.x[v2], p2_st.stage[56,v2], 'b.')                    # 2*T/6
pyplot.plot(p2_st.x[v2], p2_st.stage[57,v2], 'b.')                     # 3*T/6
pyplot.plot(p2_st.x[v2], p2_st.stage[58,v2], 'b.')                     # 4*T/6
pyplot.plot(p2_st.x[v2], p2_st.stage[59,v2], 'b.')                     # 5*T/6
pyplot.plot(p2_st.x[v2], p2_st.stage[60,v2], 'b.')  # 6*T/6

pyplot.plot(p2_st.x[v2], W1,'r-', label='analytical')
pyplot.plot(p2_st.x[v2], W2,'r-')
pyplot.plot(p2_st.x[v2], W3,'r-')
pyplot.plot(p2_st.x[v2], W4,'r-')
pyplot.plot(p2_st.x[v2], W5,'r-')
pyplot.plot(p2_st.x[v2], W6,'r-')


pyplot.title('Stage at several instants in time')
pyplot.ylim(-5.0,5.0)
pyplot.legend(loc=2)
pyplot.savefig('stage_plot.png')


#Plot the momentums##########################################################
pyplot.clf()
pyplot.ion()
pyplot.plot(p2_st.x[v2], p2_st.xmom[55,v2], 'b.', label='numerical') # 1*T/6
pyplot.plot(p2_st.x[v2], p2_st.xmom[56,v2], 'b.')                    # 2*T/6
pyplot.plot(p2_st.x[v2], p2_st.xmom[57,v2], 'b.')                     # 3*T/6
pyplot.plot(p2_st.x[v2], p2_st.xmom[58,v2], 'b.')                     # 4*T/6
pyplot.plot(p2_st.x[v2], p2_st.xmom[59,v2], 'b.')                     # 5*T/6
pyplot.plot(p2_st.x[v2], p2_st.xmom[60,v2], 'b.')  # 6*T/6

pyplot.plot(p2_st.x[v2], P1,'r-', label='analytical')
pyplot.plot(p2_st.x[v2], P2,'r-')
pyplot.plot(p2_st.x[v2], P3,'r-')
pyplot.plot(p2_st.x[v2], P4,'r-')
pyplot.plot(p2_st.x[v2], P5,'r-')
pyplot.plot(p2_st.x[v2], P6,'r-')

pyplot.title('Xmomentum at several instants in time')
pyplot.legend(loc=4)
pyplot.savefig('xmom_plot.png')



#Plot the velocities#########################################################
pyplot.clf()
pyplot.ion()
pyplot.plot(p2_st.x[v2], p2_st.xvel[55,v2], 'b.', label='numerical') # 1*T/6
pyplot.plot(p2_st.x[v2], p2_st.xvel[56,v2], 'b.')                    # 2*T/6
pyplot.plot(p2_st.x[v2], p2_st.xvel[57,v2], 'b.')                     # 3*T/6
pyplot.plot(p2_st.x[v2], p2_st.xvel[58,v2], 'b.')                     # 4*T/6
pyplot.plot(p2_st.x[v2], p2_st.xvel[59,v2], 'b.')                     # 5*T/6
pyplot.plot(p2_st.x[v2], p2_st.xvel[60,v2], 'b.')  # 6*T/6

pyplot.plot(p2_st.x[v2], U1,'r-', label='analytical')
pyplot.plot(p2_st.x[v2], U2,'r-')
pyplot.plot(p2_st.x[v2], U3,'r-')
pyplot.plot(p2_st.x[v2], U4,'r-')
pyplot.plot(p2_st.x[v2], U5,'r-')
pyplot.plot(p2_st.x[v2], U6,'r-')

pyplot.title('Velocity at several instants in time')
pyplot.legend(loc=3)
pyplot.savefig('xvel_plot.png')