"""Example of shallow water wave equation analytical solution
consists of a symmetrical converging channel with friction and bed slope.

Specific methods pertaining to the 2D shallow water equation
are imported from shallow_water
for use with the generic finite volume framework

   Copyright 2004
   Christopher Zoppou, Stephen Roberts, Ole Nielsen, Duncan Gray
   Geoscience Australia, ANU
   
Specific methods pertaining to the 2D shallow water equation
are imported from shallow_water
for use with the generic finite volume framework

Conserved quantities are h, uh and vh stored as elements 0, 1 and 2 in the
numerical vector named conserved_quantities.
"""

#-----------------------------
# Module imports 
import sys, string
from os import sep
sys.path.append('..'+sep+'pyvolution')

from shallow_water import Transmissive_boundary, Reflective_boundary, \
     Dirichlet_boundary
from shallow_water import Domain
from pmesh2domain import pmesh_to_domain_instance
from Numeric import zeros, Float 


#------------------------------
# Domain 
filename = 'MacDonald_77541.tsh'
print 'Creating domain from', filename
domain = pmesh_to_domain_instance(filename, Domain)
print 'Number of triangles = ', len(domain)

domain.default_order = 1
domain.smooth = True

#-------------------------------------
# Provide file name for storing output
domain.store = True     #Store for visualisation purposes
domain.format = 'sww'   #Native netcdf visualisation format
domain.filename = 'MacDonald_77541'

#-------------
#Bed Elevation
fid = open('MacDonald_bed.xya')
lines = fid.readlines()
n_bed = len(lines)
z_bed = zeros(n_bed, Float)
x_bed = zeros(n_bed, Float)

for line in range(n_bed):
    value = string.split(lines[line])
    x_bed[line] = float(value[0])
    z_bed[line] = float(value[1])
    
#-----------------
#Set bed-elevation
def x_slope(x,y):
    n = x.shape[0]
    z = 0*x
    for i in range(n):
        for j in range(n_bed-1):
            if x[i] >= x_bed[j] and x[i] < x_bed[j+1]:
                z[i] = z_bed[j] + (z_bed[j+1] - z_bed[j])/(x_bed[j+1] - x_bed[j])*(x[i] - x_bed[j])
            return z

domain.set_quantity('elevation', x_slope)
    
#---------------------------------------------------------
#Decide which quantities are to be stored at each timestep
domain.quantities_to_be_stored = ['stage', 'xmomentum', 'ymomentum']

#Reduction operation for get_vertex_values              
from util import mean
domain.reduction = mean 

#--------------------
# Boundary Conditions
upstream_depth = 5.035 - 4.393
downstream_depth = 1.5 - 0
discharge = 20
tags = {}
tags['Upstream'] = Dirichlet_boundary([upstream_depth, 2, 0.0])
tags['external'] = Reflective_boundary(domain) 
tags['Downstream'] = Dirichlet_boundary([downstream_depth, 2, 0.0])
domain.set_boundary(tags)

#---------
# friction
domain.set_quantity('friction', 0.02)

#-----------------
#Initial condition
domain.set_quantity('elevation', 0.0)
domain.set_quantity('stage', 0.2)
    
#-----------------
#Evolution
import time
t0 = time.time()
for t in domain.evolve(yieldstep = 5, finaltime = 500):
    domain.write_time()
    
print 'That took %.2f seconds' %(time.time()-t0)