import os
import time
from shallow_water_domain_avalanche import *
from Numeric import array, zeros, Float, sqrt
from config import g, epsilon
from numpy import sin, cos, tan
from scipy import linspace
from parameters import *

def analytical_sol(X,t):
    N = len(X)
    u = zeros(N,Float)
    h = zeros(N,Float)
    w = zeros(N,Float)
    z = zeros(N,Float)
    mom = zeros(N,Float)
    for i in range(N):
        # Calculate Analytical Solution at time t > 0
        if X[i] <= -2.0*c0*t + 0.5*m*t**2.0:
            u[i] = 0.0
            h[i] = 0.0
        elif X[i] <= c0*t + 0.5*m*t**2.0:
            u[i] = 2.0/3.0 * (X[i]/t - c0 + m*t)
            h[i] = 1.0/(9.0*g) * (X[i]/t + 2.0*c0 - 0.5*m*t)**2.0
        else:
            u[i] = m*t
            h[i] = h_0
        z[i] = bed_slope*X[i]
        w[i] = h[i] + z[i]
        mom[i] = u[i]*h[i]
    return u,h,w,z,mom

def height(X):
    N = len(X)
    y = zeros(N,Float)
    for i in range(N):
        if X[i]<=0.0:
            y[i] = 0.0
        else:
            y[i] = h_0
    return y

def stage(X):
    N = len(X)
    w = zeros(N,Float)
    for i in range(N):
        if X[i]<=0.0:
            w[i] = bed_slope*X[i]
        else:
            w[i] = bed_slope*X[i] + h_0
    return w        

def elevation(X):
    N = len(X)
    y=zeros(N, Float)
    for i in range(N):
        y[i] = bed_slope*X[i]
    return y
"""
#=========================================================================#
#The following values are set in parameters.py
h_0 = 10.0         # depth upstream. Note that the depth downstream is 0.0
L = 2000.0         # length of stream/domain
n = 100             # number of cells
cell_len = L/n     # length of each cell
points = zeros(n+1, Float)
for i in range (n+1):
    points[i] = i*cell_len - 0.5*L

bed_slope = 0.005      # bottom slope, positive if it is increasing bottom.
c0 = sqrt(g*h_0)   # sound speed
m = -1.0*g*bed_slope   # auxiliary variable
#==========================================================================#
"""
boundary = { (0,0): 'left',(n-1,1): 'right'}
domain = Domain(points,boundary)
domain.order = 2
domain.set_timestepping_method('rk2')
domain.set_CFL(1.0)
domain.beta = 1.0
domain.set_limiter("minmod")

def f_right(t):
    z_r = bed_slope*(0.5*L)
    h_r = h_0 #+ bed_slope*cell_len
    w_r = z_r + h_r
    u_r = m*t   
    #['stage', 'xmomentum', 'elevation', 'height', 'velocity']
    return [w_r,  u_r*h_r,  z_r,  h_r,  u_r]

T_right = Time_boundary(domain,f_right)
#T_right = Transmissive_boundary(domain)
#D_right = Dirichlet_boundary([bed_slope*(0.5*L)+h_0,  (m*domain.time)*h_0,  bed_slope*(0.5*L),  h_0,  m*domain.time])
D_left = Dirichlet_boundary([-1.0*bed_slope*(0.5*L),  0.0,  -1.0*bed_slope*(0.5*L),  0.0,  0.0])
domain.set_boundary({'left':D_left,'right':T_right})
domain.set_quantity('stage',stage)
domain.set_quantity('elevation',elevation)

X = domain.vertices
C = domain.centroids

yieldstep = finaltime = 15.0
t0 = time.time()

while finaltime <= 15.1:
    for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
        domain.write_time()
        
    #The following is for computing the error and plotting the result
    Mom = domain.quantities['xmomentum']
    Height = domain.quantities['height']
    Stage = domain.quantities['stage']
    Velocity = domain.quantities['velocity']
    Elevation = domain.quantities['elevation']

    #The following is for computing the error
    Uc,Hc,Wc,Zc,Mc = analytical_sol(C,domain.time)
    HeightC = Height.centroid_values
    MomC = Mom.centroid_values
    StageC = Stage.centroid_values
    VelC = Velocity.centroid_values
    ElevationC = Elevation.centroid_values

    print "number of cells=",n
    W_error = (1.0/n)*sum(abs(Wc-StageC))
    M_error = (1.0/n)*sum(abs(Mc-MomC))
    U_error = (1.0/n)*sum(abs(Uc-VelC))
    print "stage_error %.10f" %(W_error)
    print "momentum_error %.10f"%(M_error)
    print "velocity_error %.10f" %(U_error)


    #The following is for plotting the result.
    Uv,Hv,Wv,Zv,Mv = analytical_sol(X.flat,domain.time)
    HeightV = Height.vertex_values
    MomV = Mom.vertex_values
    StageV = Stage.vertex_values
    VelV = Velocity.vertex_values
    ElevationV = Elevation.vertex_values
    from pylab import clf,plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
    hold(False)
    clf()

    plot1 = subplot(311)
    plot(X.flat,Wv,'b-', X.flat,StageV.flat,'k--', X.flat,ElevationV.flat,'k:')
    #plot(X.flat,Wv, X.flat,StageV.flat, X.flat,ElevationV.flat)    
    #xlabel('Position')
    ylabel('Stage')
    #plot1.set_ylim([-1.0,21.0])
    #plot1.set_xlim([-480.0,-420.0])#([-9.0e-3,9.0e-3])
    #legend(('analytical solution', 'numerical solution', 'discretized bed'), 'upper left', shadow=False)
    
    plot2 = subplot(312)
    plot(X.flat,Mv,'b-', X.flat,MomV.flat,'k--')
    #xlabel('Position')
    ylabel('Momentum')
    #plot2.set_xlim([-300.0,300.0])
    plot2.set_ylim([-310.0,10.0])#([-90.0,10.0])
    #legend(('analytical solution', 'numerical solution'), 'lower right', shadow=False)    

    plot3 = subplot(313)
    plot(X.flat,Uv,'b-', X.flat,VelV.flat,'k--')
    xlabel('Position')
    ylabel('Velocity')
    #plot3.set_xlim([-300.0,300.0])
    plot3.set_ylim([-45.0,5.0])#([-30.0,5.0])
    #legend(('analytical solution', 'numerical solution'), 'lower right', shadow=False)
    
    finaltime = finaltime + 10.0

    file = "A-case3-"
    file += str(n)
    file += ".eps"
    savefig(file)      
    
    show()