import os
from math import sqrt, pi, sin, cos
from shallow_water_domain import *
from Numeric import allclose, array, zeros, ones, Float, take, sqrt
from config import g, epsilon
"""
def newLinePlot(title='Simple Plot'):
    import Gnuplot
    g = Gnuplot.Gnuplot(persist=0)
    g.title(title)
    g('set data style linespoints') 
    g.xlabel('x')
    g.ylabel('y')
    return g

def linePlot(g,x1,y1,x2,y2,x3,y3):
    import Gnuplot
    plot1 = Gnuplot.PlotItems.Data(x1.flat,y1.flat,with="lines 2")
    plot2 = Gnuplot.PlotItems.Data(x2.flat,y2.flat,with="lines 3")
    plot3 = Gnuplot.PlotItems.Data(x3.flat,y3.flat,with="linespoints 1")
    g.plot(plot1,plot2,plot3)
"""

def analytic_cannal(C,t):
    N = len(C)
    u = zeros(N,Float)    ## water velocity
    h = zeros(N,Float)    ## water depth
    x = C
    g = 9.81


    ## Define Basin Bathymetry
    z_b = zeros(N,Float) ## elevation of basin
    z = zeros(N,Float)   ## elevation of water surface
    z_infty = 10.0       ## max equilibrium water depth at lowest point.
    L_x = 2500.0         ## width of channel

    A0 = 0.5*L_x                  ## determines amplitudes of oscillations
    omega = sqrt(2*g*z_infty)/L_x ## angular frequency of osccilation

    #x1 = A0*cos(omega*t)-L_x # left shoreline
    #x2 = A0*cos(omega*t)+L_x # right shoreline
    
    for i in range(N):
        z_b[i] = z_infty*(x[i]**2/L_x**2)
        u[i] = -A0*omega*sin(omega*t)
        z[i] = z_infty+2*A0*z_infty/L_x*cos(omega*t)*(x[i]/L_x-0.5*A0/(L_x)*cos(omega*t))

    h = z-z_b
    return u,h,z,z_b


#plot2 = newLinePlot("Momentum")  

L_x = 2500.0     # Length of channel (m)
N = 100         # Number of compuational cells
cell_len = 4*L_x/N   # Origin = 0.0

points = zeros(N+1,Float)
for i in range(N+1):
        points[i] = -2*L_x +i*cell_len

domain = Domain(points)

domain.order = 2
domain.set_timestepping_method('rk2')
domain.cfl = 1.0
domain.beta = 1.0
#domain.limiter = "minmod_kurganov"
#domain.limiter = "vanleer"
#domain.limiter = "vanalbada"
#domain.limiter = "superbee"
#domain.limiter = "steve_minmod"
domain.limiter = "pyvolution"

def stage(x):

    z_infty = 10.0       ## max equilibrium water depth at lowest point.
    L_x = 2500.0         ## width of channel
    A0 = 0.5*L_x                  ## determines amplitudes of oscillations
    omega = sqrt(2*g*z_infty)/L_x ## angular frequency of osccilation
    t=0.0
    y = zeros(len(x),Float)
    for i in range(len(x)):
        #xy[i] = z_infty+2*A0*z_infty/L_x*cos(omega*t)*(x[i]/L_x-0.5*A0/(L_x)*cos(omega*t))
        y[i] = 12.0
    return y

def elevation(x):
    N = len(x)
    z_infty = 10.0
    z = zeros(N,Float)
    L_x = 2500.0
    A0 = 0.5*L_x
    omega = sqrt(2*g*z_infty)/L_x
    for i in range(N):
        z[i] = z_infty*(x[i]**2/L_x**2)
    return z

def height(x):
    z_infty = 10.0       ## max equilibrium water depth at lowest point.
    L_x = 2500.0         ## width of channel
    A0 = 0.5*L_x                  ## determines amplitudes of oscillations
    omega = sqrt(2*g*z_infty)/L_x ## angular frequency of osccilation
    t=0.0
    y = zeros(len(x),Float)
    for i in range(len(x)):
        y[i] = max(z_infty+2*A0*z_infty/L_x*cos(omega*t)*(x[i]/L_x-0.5*A0/(L_x)*cos(omega*t))-z_infty*(x[i]**2/L_x**2),0.0)
    return y

    
domain.set_quantity('stage', stage)
domain.set_quantity('elevation',elevation)
#domain.set_quantity('height',height)
domain.set_boundary({'exterior': Reflective_boundary(domain)})
 
import time
t0 = time.time()
#finaltime = 1122.0*0.75
yieldstep = finaltime = 10.0
#yieldstep = 10.0
#plot1 = newLinePlot("Stage")
#plot2 = newLinePlot("Xmomentum")
C = domain.centroids
X = domain.vertices
StageQ = domain.quantities['stage'].vertex_values
XmomQ = domain.quantities['xmomentum'].vertex_values
#Bed = domain.quantities['elevation'].vertex_values
import time
t0 = time.time()
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    #pass
    domain.write_time()
    u,h,z,z_b = analytic_cannal(X.flat,t)
    #linePlot(plot1,X,z,X,z_b,X,StageQ)
    #linePlot(plot2,X,u*h,X,u*h,X,XmomQ)
    HeightQ = domain.quantities['stage'].centroid_values-domain.quantities['elevation'].centroid_values
    u,hc,z,z_b = analytic_cannal(C,t)
    #for k in range(N):
    #    if hc[k] < 0.0:
    #        hc[k] = 0.0
    error = 1.0/(N)*sum(abs(hc-HeightQ))
    print 'Error measured at centroids', error
    # #raw_input()         
    # #from Gnuplot import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
    from pylab import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
    # #import Gnuplot  
    X = domain.vertices
    u,h,z,z_b = analytic_cannal(X.flat,domain.time)
    StageQ = domain.quantities['stage'].vertex_values
    XmomQ = domain.quantities['xmomentum'].vertex_values
    hold(False)
    plot1 = subplot(211)
    plot(X,h,X,StageQ,X,z_b)
    #plot1.set_ylim([4,11])
    #title('?????????')
    xlabel('Position')
    ylabel('Stage')
    legend(('Analytical Solution', 'Numerical Solution', 'Bed'),
           'upper right', shadow=True)
    plot2 = subplot(212)
    plot(X,u*h,X,XmomQ)
    #plot2.set_ylim([-1,25])
    #title('?????????????????????????')
    xlabel('Position')
    ylabel('Xmomentum')
    show()