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

def analytical_sol(C,t):
    
    #t  = 0.0     # time (s)
    g  = 9.81    # 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)) 

        if ( 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

#def newLinePlot(title='Simple Plot'):
#   import Gnuplot
#    gg = Gnuplot.Gnuplot(persist=0)
#    gg.terminal(postscript)
#    gg.title(title)
#    gg('set data style linespoints') 
#    gg.xlabel('x')
#    gg.ylabel('y')
#    return gg

#def linePlot(gg,x1,y1,x2,y2):
#    import Gnuplot
#    plot1 = Gnuplot.PlotItems.Data(x1.flat,y1.flat,with="linespoints")
#    plot2 = Gnuplot.PlotItems.Data(x2.flat,y2.flat, with="lines 3")
#    g.plot(plot1,plot2)



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

L = 2000.0     # Length of channel (m)
N = 800        # Number of compuational cells
cell_len = L/N # Origin = 0.0

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

domain = Domain(points)

def stage(x):
    h1 = 10.0
    h0 = 0.0
    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 = 20.0
yieldstep = 1.0
L = 2000.0     # Length of channel (m)
k = 0

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.default_order = 2
domain.default_time_order = 1
domain.cfl = 1.0
domain.limiter = "vanleer"

t0 = time.time()

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

import profile, pstats
FN = 'profile.dat'

profile.run(s, FN)

    #print 'That took %.2f seconds' %(time.time()-t0)

S = pstats.Stats(FN)
s = S.sort_stats('cumulative').print_stats(30)

print s    
        
N = float(N)
StageC = domain.quantities['stage'].centroid_values
XmomC = domain.quantities['xmomentum'].centroid_values
C = domain.centroids
h, uh = analytical_sol(C,domain.time)
h_error = 1.0/(N)*sum(abs(h-StageC))
uh_error = 1.0/(N)*sum(abs(uh-XmomC))
print "h_error %.10f" %(h_error)
print "uh_error %.10f"% (uh_error)

print 'That took %.2f seconds' %(time.time()-t0)