import os
#from math import sqrt, pi
from shallow_water_h import Domain, Reflective_boundary
from Numeric import zeros, Float
#from config import g, epsilon
from analytic_dam import AnalyticDam


h0 = 5.0
h1 = 10.0

analytical_sol = AnalyticDam(h0, h1)

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):
    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 I"

L = 2000.0     # Length of channel (m)
N = 100        # Number of computational 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 height(x):
    y = zeros(len(x),Float)
    for i in range(len(x)):
        if x[i]<=1000.0:
            y[i] = h1
        else:
            y[i] = h0
    return y


domain.set_quantity('height', height)

domain.default_order = 2
domain.order = 2
domain.default_time_order = 1
domain.cfl = 0.8
#domain.beta = 1.0
domain.split = False
print "domain.order", domain.order
domain.limiter = "minmod"


domain.set_boundary({'exterior': Reflective_boundary(domain)})

X = domain.vertices
C = domain.centroids
plot1 = newLinePlot("Height")
plot2 = newLinePlot("Momentum")

import time
t0 = time.time()
yieldstep = 1.0
finaltime = 20.0
print "integral", domain.quantities['height'].get_integral()
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()
    print "integral", domain.quantities['height'].get_integral()
    if t > 0.0:
        HeightQ = domain.quantities['height'].vertex_values
        y , my = analytical_sol(X.flat,domain.time)
        linePlot(plot1,X,HeightQ,X,y)
        MomentumQ = domain.quantities['xmomentum'].vertex_values
        linePlot(plot2,X,MomentumQ,X,my)
    #raw_input('press_return')

    #pass

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

#C = domain.quantities['height'].centroid_values