"""Validation study of Merimbula lake using Pyvolution.
Example of shallow water wave equation applied to
Malpasset dam break simulation.

   Copyright 2004
   Christopher Zoppou, Stephen Roberts
   Australian National University
   
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.

Existence of file 'Hobart_mesh.tsh' is assumed.
"""

###############################
# Setup Path and import modules
import sys
from os import sep, path
sys.path.append('..'+sep+'pyvolution')

from shallow_water import Domain, Reflective_boundary, File_boundary,\
     Dirichlet_boundary, Transmissive_boundary, Constant_height
from pmesh2domain import pmesh_to_domain_instance
from util import inside_polygon, read_polygon

######################
# Domain
filename = 'Hobart_mesh.tsh'
yieldstep = 1
finaltime = 100
    
print 'Creating domain from', filename
domain = pmesh_to_domain_instance(filename, Domain)
print 'Number of triangles = ', len(domain)
print 'Extent = ', domain.get_extent()


domain.default_order = 1
domain.smooth = True

#domain.visualise = True


#------------------------------
# Boundary Conditions
tags = {}
tags['external'] = Reflective_boundary(domain) 
domain.set_boundary(tags)

#-----------------
#Initial condition
#   Throughout
domain.set_quantity('stage', 0.)

#    Within polygon region
#        Set elevation to bed plus depth
depth = 10.
#        Read polygon boundary
p0 = read_polygon('Hobart_recent_source_zone.xya')
#        Get all centroid coordinates
X = domain.get_centroid_coordinates()
#        Find triangle indices which are within polygon boundary only
indices = inside_polygon(X, p0)
#        Get the bed elevation at the centroid of every triangle in polygon region

zp = domain.get_quantity('elevation', location='centroids', indexes=indices)
#        Set the stage to bed elevation plus depth for all triangle vertices in the polygon region

domain.set_quantity('stage', zp+depth, location='centroids', indexes=indices) 


#Alternative way which also works (I used it for testing - OMN)
#

#z = domain.get_quantity('elevation')
#from copy import copy
#w = copy(z)
#for i in indices:
#    w[i, :] = z[i, :] + depth
#domain.set_quantity('stage', w)


#-------------------------------------
# Provide file name for storing output
domain.store = True
domain.format = 'sww'
domain.filename = 'Hobart_first_order'

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

#----------------------------
# Friction
domain.set_quantity('friction', 0.05)
         
######################
#Evolution
import time
t0 = time.time()
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()
    
print 'That took %.2f seconds' %(time.time()-t0)