import os
from math import sqrt, pi
from shallow_water_vel_domain import *
from Numeric import allclose, array, zeros, ones, Float, take, sqrt
from config import g, epsilon


h1 = 10.0
h0 = 0.0

def analytical_sol(C,t):
    
    #t  = 0.0     # time (s)
    # gravity (m/s^2)
    #h1 = 10.0    # depth upstream (m)
    #h0 = 0.0     # depth downstream (m)
    L = 2000.0   # length of stream/domain (m)
    n = len(C)    # number of cells

    u = zeros(n,Float)
    h = zeros(n,Float)
    x = C-3*L/4.0
    

    for i in range(n):
        # Calculate Analytical Solution at time t > 0
        u3 = 2.0/3.0*(sqrt(g*h1)+x[i]/t) 
        h3 = 4.0/(9.0*g)*(sqrt(g*h1)-x[i]/(2.0*t))*(sqrt(g*h1)-x[i]/(2.0*t))
        u3_ = 2.0/3.0*((x[i]+L/2.0)/t-sqrt(g*h1))
        h3_ = 1.0/(9.0*g)*((x[i]+L/2.0)/t+2*sqrt(g*h1))*((x[i]+L/2.0)/t+2*sqrt(g*h1))

        if ( x[i] <= -1*L/2.0+2*(-sqrt(g*h1)*t)):
            u[i] = 0.0
            h[i] = h0
        elif ( x[i] <= -1*L/2.0-(-sqrt(g*h1)*t)):
            u[i] = u3_
            h[i] = h3_

        elif ( x[i] <= -t*sqrt(g*h1) ):
            u[i] = 0.0 
            h[i] = h1 
        elif ( x[i] <= 2.0*t*sqrt(g*h1) ):
            u[i] = u3 
            h[i] = h3 
        else:
            u[i] = 0.0 
            h[i] = h0 

    return h , u*h, u

print "TEST 1D-SOLUTION III -- DRY BED"

def stage(x):
    y = zeros(len(x),Float)
    for i in range(len(x)):
        if x[i]<=L/4.0:
            y[i] = h0
        elif x[i]<=3*L/4.0:
            y[i] = h1
        else:
            y[i] = h0
    return y


import time

finaltime = 2.0
yieldstep = 0.1
L = 2000.0     # Length of channel (m)

k = 0

N = 800
print "Evaluating domain with %d cells" %N
cell_len = L/N # Origin = 0.0
points = zeros(N+1,Float)

for j in range(N+1):
    points[j] = j*cell_len
        
domain = Domain(points)
    
domain.set_quantity('stage', stage)
domain.set_boundary({'exterior': Reflective_boundary(domain)})
domain.order = 2
domain.set_timestepping_method('euler')
domain.set_CFL(1.0)
domain.set_limiter("vanleer")
#domain.h0=0.0001

t0 = time.time()

for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()

print 'end'