import os
from math import sqrt
from sww_domain_simplified import *
from Numeric import Float
from numpy import zeros
from sf_parameters import *

N = int(N) # number of cells
print "number of cells=",N
boundary = {(0,0):'left', (N-1,1): 'right'}
domain = Domain(points,boundary)
domain.order = 2
domain.set_timestepping_method('rk2')
domain.cfl = 1.0
domain.limiter = "minmod"

def stage(x):
    y=zeros(len(x), Float)
    for i in range(len(x)):
        if x[i] < 11.666:
            y[i] = 0.4125
        else:
            y[i] = 0.33
    return y
domain.set_quantity('stage',stage)

def elevation(x):
    z_b = zeros(len(x),Float)
    for i in range(len(x)):
        if (x[i] >= 8.0) & (x[i] <= 12.0):
            z_b[i] = 0.2 - 0.05*(x[i]-10.0)**2.0
        else:
            z_b[i] = 0.0 
    return z_b
domain.set_quantity('elevation',elevation)

### ================ Define the boundary function =========================
# ['stage', 'xmomentum', 'elevation', 'height', 'velocity']
D_left = Dirichlet_boundary([0.4125,  0.18,  0.0,  0.4125,  0.18/0.4125])
D_right = Dirichlet_boundary([0.33,  0.18,  0.0,  0.33,  0.18/0.33])
domain.set_boundary({'left':D_left,'right':D_right})
### ================  End of the definition of boundary function ===========

X=domain.vertices
C=domain.centroids
import time
yieldstep=finaltime=0.1
t0=time.time()
i=1

while finaltime < 0.101:
    for t in domain.evolve(yieldstep=yieldstep, finaltime=finaltime):
        domain.write_time()
    finaltime = finaltime + 10.0 
if t>0.0:
    N = float(N)
    StageC = domain.quantities['stage'].centroid_values
    XmomC = domain.quantities['xmomentum'].centroid_values
    VelC = domain.quantities['velocity'].centroid_values

    X = domain.vertices
    StageQ = domain.quantities['stage'].vertex_values.flat
    XmomQ = domain.quantities['xmomentum'].vertex_values.flat
    VelQ = domain.quantities['velocity'].vertex_values.flat
    BedQ = domain.quantities['elevation'].vertex_values.flat

    SD = domain.shock_detector

    from pylab import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
    hold(False)

    plot1 = subplot(211)
    plot(X,StageQ, X,BedQ)
    #plot1.set_ylim([-1,11])
    #plot1.set_xlim([0.0,2000.0])
    legend(('Numerical solution', 'Bed elevation'),
            'upper left', shadow=False)    
    #xlabel('Position')
    ylabel('Stage')
    

    plot2 = subplot(212)
    plot(points,SD)
    #plot2.set_xlim([0.0,2000.0])
    xlabel('Position')
    ylabel('Smoothness indicator')
    
    print 'That took %.2f seconds'%(time.time()-t0)
    show()
        
    #filename = "%s%04i%s" %("dam_", i, ".eps")
    #savefig(filename)
    #finaltime = finaltime + 0.25
    #print "finaltime=", finaltime
    #i = i + 1
    #print "The domain.limiter is", domain.limiter
    #print 'That took %.2f seconds'%(time.time()-t0)
    #print '============================================================================='