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


print "Radial Dam Break Test"

# Define functions for initial quantities

def initial_area(x):
    y=zeros(len(x),Float)
    for i in range (len(x)):
         if x[i]<=40:
            y[i]=10*(x[i]*pi*2)
         else:
            y[i]=1*(x[i]*pi*2)
    return y
        

def width(x):
    return 2*pi*x
 
import time

# Set final time and yield time for simulation
finaltime = 2.0
yieldstep = finaltime

# Length of channel (m)
L = 100.0   
# Define the number of cells
number_of_cells = [100]

# Define cells for finite volume and their size
N = int(number_of_cells[0])     
print "Evaluating domain with %d cells" %N
cell_len = L/N # Origin = 0.0
points = zeros(N+1,Float)

# Define the centroid points
for j in range(N+1):
    points[j] = j*cell_len

# Create domain with centroid points as defined above        
domain = Domain(points)

# Set initial values of quantities - default to zero
domain.set_quantity('area', initial_area)
domain.set_quantity('width',width)


# Set boundry type, order, timestepping method and limiter
domain.set_boundary({'exterior':Reflective_boundary(domain)})
domain.order = 2
domain.set_timestepping_method('rk2')
domain.set_CFL(1.0)
domain.set_limiter("vanleer")
#domain.h0=0.0001

# Start timer
t0 = time.time()

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

N = float(N)
HeightC = domain.quantities['height'].centroid_values
DischargeC = domain.quantities['discharge'].centroid_values
C = domain.centroids
print 'That took %.2f seconds' %(time.time()-t0)
X = domain.vertices
HeightQ = domain.quantities['height'].vertex_values
VelocityQ = domain.quantities['velocity'].vertex_values
x = X.flat
z = domain.quantities['elevation'].vertex_values.flat
stage=HeightQ.flat+z

## ################# REPEAT ABOVE WITH MORE CELLS
## # Set final time and yield time for simulation
## finaltime2 = 2.0
## yieldstep2 = finaltime2

## # Length of channel (m)
## L2 = 100.0   
## # Define the number of cells
## number_of_cells2 = [100]

## # Define cells for finite volume and their size
## N2 = int(number_of_cells2[0])     
## print "Evaluating domain with %d cells" %N2
## cell_len2 = L2/N2 # Origin = 0.0
## points2 = zeros(N2+1,Float)

## # Define the centroid points
## for j in range(N2+1):
##     points2[j] = j*cell_len2

## # Create domain with centroid points as defined above        
## domain2 = Domain(points2)

## # Set initial values of quantities - default to zero
## domain2.set_quantity('area', initial_area)
## domain2.set_quantity('width',width)


## # Set boundry type, order, timestepping method and limiter
## domain2.set_boundary({'exterior':Reflective_boundary(domain2)})
## domain2.order = 2
## domain2.set_timestepping_method('rk2')
## domain2.set_CFL(1.0)
## domain2.set_limiter("vanleer")
## #domain.h0=0.0001

## # Start timer
## t02 = time.time()

## for t in domain2.evolve(yieldstep = yieldstep2, finaltime = finaltime2):
##     domain2.write_time()

## N2 = float(N2)
## HeightC2 = domain2.quantities['height'].centroid_values
## DischargeC2 = domain2.quantities['discharge'].centroid_values
## C2 = domain2.centroids
## print 'That took %.2f seconds' %(time.time()-t02)
## X2 = domain2.vertices
## HeightQ2 = domain2.quantities['height'].vertex_values
## VelocityQ2 = domain2.quantities['velocity'].vertex_values
## x2 = X2.flat
## z2 = domain2.quantities['elevation'].vertex_values.flat
## stage2=HeightQ2.flat+z2

## ################# REPEAT ABOVE

from pylab import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot
import pickle
f=open('highresdam.txt','r')
## Data=[[x,stage],[x,VelocityQ.flat]]
## pickle.dump(Data,f)
## f.close()
[[x2,stage2],[x3,Velocity2]]=pickle.load(f)


hold(False)
plot1 = subplot(211)

plot(x2,stage2,x,stage,'.')
 
plot1.set_ylim([-1,11])
xlabel('Position')
ylabel('Stage')
legend(('Analytical Solution', 'Numerical Solution'),
           'upper right', shadow=True)
plot2 = subplot(212)
plot(x3,Velocity2,x,VelocityQ.flat,'.')
plot2.set_ylim([-10,10])
    
xlabel('Position')
ylabel('Velocity')
   
show()