import os
from math import sqrt, pi, sin, cos
from shallow_water_vel_domain import *
from numpy import allclose, array, zeros, ones, numpy.float, take, sqrt
from config import g, epsilon

def analytic_cannal(C,t):
    N = len(C)
    u = zeros(N,numpy.float)    ## water velocity
    h = zeros(N,numpy.float)    ## water depth
    x = C
    g = 9.81
    ## Define Basin Bathymetry
    z_b = zeros(N,numpy.float) ## elevation of basin
    z = zeros(N,numpy.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
    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
    T = 2.0*pi/omega
    return u,h,z,z_b, T


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

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

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),numpy.float)
    for i in range(len(x)):
        y[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,numpy.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),numpy.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 = Domain(points)
domain.order = 2
domain.set_timestepping_method('euler')
domain.set_CFL(1.0)
domain.beta = 1.0
domain.set_limiter("vanleer")

    
domain.set_quantity('stage', stage)
domain.set_quantity('elevation',elevation)
domain.set_boundary({'exterior': Reflective_boundary(domain)})
 
X = domain.vertices
u,h,z,z_b,T = analytic_cannal(X.flat,domain.time)
print 'T = ',T

yieldstep = finaltime = 10.0 #T/2.0
StageQ = domain.quantities['stage'].vertex_values
XmomQ = domain.quantities['xmomentum'].vertex_values

import time
t0 = time.time()
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()
    print "integral", domain.quantities['stage'].get_integral()
    if t>0.0:
        
        x = X.flat
        print 'domain.time=',domain.time
        
        w_V = domain.quantities['stage'].vertex_values.flat
        uh_V = domain.quantities['xmomentum'].vertex_values.flat
        u_V = domain.quantities['velocity'].vertex_values.flat
        h_V = domain.quantities['height'].vertex_values.flat
        
        u,h,z,z_b,T = analytic_cannal(x,domain.time)
        w = z
        for k in range(len(h)):
            if h[k] < 0.0:
                h[k] = 0.0
                w[k] = z_b[k]
        
        from pylab import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
        hold(False)

        #print 'size X',X.shape
        #print 'size w',w.shape
        signh = (h_V>0)

        plot1 = subplot(311)
        #plot(x,w,  x,w_V,  x,z_b)
        plot(x,signh)
        plot1.set_xlim([-6000,6000])
        xlabel('Position')
        ylabel('Stage')
        legend(('Analytic Solution', 'numpyal Solution', 'Bed'),
               'upper center', shadow=True)

        
        plot2 = subplot(312)
        plot(x,u*h,x,uh_V)
        #plot2.set_ylim([-1,25])
        xlabel('Position')
        ylabel('Xmomentum')

        print u_V
        
        plot3 = subplot(313)
        plot(x,u, x,u_V)
        #plot2.set_ylim([-1,25])
        xlabel('Position')
        ylabel('Velocity')
        show()
raw_input("Press the return key!")