%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from ipywidgets import StaticInteract, RangeWidget
from clawpack import riemann

from clawpack.riemann import riemann_tools   
import SW_exact_Riemann_solver
import SW_phase

grav = 9.81

q_l = np.array((1.,3.))
q_r = np.array((1.,-3.))

ex_states, ex_speeds, reval = SW_exact_Riemann_solver.exact_riemann_solution(q_l ,q_r, grav)

print "States in solution:"
print ex_states

plot_function = riemann_tools.make_plot_function(ex_states, ex_speeds, reval)
StaticInteract(plot_function, t=RangeWidget(0,.9,.1))

solver = riemann.shallow_1D_py.shallow_roe_1D
num_eqn = riemann.shallow_1D_py.num_eqn

problem_data = {}
problem_data['grav'] = grav
problem_data['efix'] = False

print "Roe solver solution to SW equations:"
states, s, roe_eval = riemann_tools.riemann_solution(num_eqn,solver,q_l,q_r,problem_data=problem_data)

#riemann_tools.plot_phase(states)

plt.figure()
SW_phase.phase_plane_curves(q_l[0],q_l[1],'qleft',1)
SW_phase.phase_plane_curves(q_r[0],q_r[1],'qright',2)
plt.plot(states[0,:],states[1,:],'k--o')
plt.xlim(0,4)
plt.ylim(-6,6)

plot_function = riemann_tools.make_plot_function(states,s,roe_eval)
StaticInteract(plot_function, t=RangeWidget(0,.9,.1))

plot_function = riemann_tools.make_plot_function([ex_states,states],
                                                 [ex_speeds,s],
                                                 [reval,roe_eval],
                                                 ['Exact','Roe'])
plt.figure()
SW_phase.phase_plane_curves(q_l[0],q_l[1],'qleft',1)
SW_phase.phase_plane_curves(q_r[0],q_r[1],'qright',2)
plt.plot(states[0,:],states[1,:],'k--o')
plt.xlim(0,4)
plt.ylim(-6,6)

StaticInteract(plot_function, t=RangeWidget(0,.9,.1))

q_l = np.array((3.,2.))
q_r = np.array((1.,1.))

ex_states, ex_speeds, reval = SW_exact_Riemann_solver.exact_riemann_solution(q_l ,q_r, grav)

states, s, roe_eval = riemann_tools.riemann_solution(num_eqn,solver,q_l,q_r,problem_data=problem_data)


plot_function = riemann_tools.make_plot_function([ex_states,states],
                                                 [ex_speeds,s],
                                                 [reval,roe_eval],
                                                 ['Exact','Roe'])
plt.figure()
SW_phase.phase_plane_curves(q_l[0],q_l[1],'qleft',1)
SW_phase.phase_plane_curves(q_r[0],q_r[1],'qright',2)
plt.plot(states[0,:],states[1,:],'k--o')
plt.xlim(0,4)
plt.ylim(-6,8)

StaticInteract(plot_function, t=RangeWidget(0,.9,.1))

q_l = np.array((0.5, -1.5))
q_r = np.array((1.5, -4.5))

ex_states, ex_speeds, reval = SW_exact_Riemann_solver.exact_riemann_solution(q_l ,q_r, grav)

states, s, roe_eval = riemann_tools.riemann_solution(num_eqn,solver,q_l,q_r,problem_data=problem_data)


plot_function = riemann_tools.make_plot_function([ex_states,states],
                                                 [ex_speeds,s],
                                                 [reval,roe_eval],
                                                 ['Exact','Roe'])
plt.figure()
SW_phase.phase_plane_curves(q_l[0],q_l[1],'qleft',1)
SW_phase.phase_plane_curves(q_r[0],q_r[1],'qright',2)
plt.plot(states[0,:],states[1,:],'k--o')
plt.xlim(0,4)
plt.ylim(-6,8)

StaticInteract(plot_function, t=RangeWidget(0,.9,.1))