import os
from math import sqrt, pi
from shallow_water_1d import *
from Numeric import allclose, array, zeros, ones, Float, take, sqrt
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=1)
    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)
    #g.plot(Gnuplot.PlotItems.Data(x1.flat,y1.flat),with="linespoints")
    #g.plot(Gnuplot.PlotItems.Data(x2.flat,y2.flat), with="lines")

debug = False

print "TEST 1D-SOLUTION I"

L = 2000.0     # Length of channel (m)
N = 100        # 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):
    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('stage', stage)

domain.order = 2
domain.default_order = 2
domain.cfl = 0.8
#domain.beta = 0.0
print "domain.order", domain.order

if (debug == True):
    area = domain.areas
    for i in range(len(area)):
        if area != 20:
            print "Cell Areas are Wrong"
            
    L = domain.quantities['stage'].vertex_values
    print "Initial Stage"
    print L

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

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

import time
t0 = time.time()
yieldstep = 1.0
finaltime = 50.0
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()
    if t > 0.0:
        StageQ = domain.quantities['stage'].vertex_values
        y , my = analytical_sol(X.flat,domain.time)
        linePlot(plot1,X,StageQ,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['stage'].centroid_values

if (debug == True):
    L = domain.quantities['stage'].vertex_values
    print "Final Stage Vertex Values"
    print L
    print "Final Stage Centroid Values"
    print C

#f = file('test_solution_I.out', 'w')
#for i in range(N):
#    f.write(str(C[i]))
#    f.write("\n")
#f.close