"""Example of shallow water wave equation analytical solution
consists of a flat water surface profile in a parabolic basin
with linear friction. The analytical solution was derived by Sampson in 2002.

   Copyright 2004
   Christopher Zoppou, Stephen Roberts, Ole Nielsen, Duncan Gray
   Geoscience Australia
   
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
from os import sep
sys.path.append('..'+sep+'pyvolution')

from shallow_water import Domain, Dirichlet_boundary, gravity, linear_friction
from math import sqrt, cos, sin, pi, exp
from mesh_factory import strang_mesh
from quantity import Quantity


######################
# Domain
# Strang_domain will search through the file and test to see if there are
# two or three entries. Two entries are for points and three for triangles.


points, elements = strang_mesh('yoon_circle.pt')
domain = Domain(points, elements) 

domain.default_order = 2
domain.smooth = True

domain.quantities['linear_friction'] = Quantity(domain)

#Reconstruct forcing terms with linear friction instead og manning friction
domain.forcing_terms = []
domain.forcing_terms.append(gravity)
domain.forcing_terms.append(linear_friction)

print domain.forcing_terms

# Provide file name for storing output
domain.store = True
domain.format = 'sww'
domain.set_name('sampson_strang_second_order')

print 'Number of triangles = ', len(domain)


#Reduction operation for get_vertex_values              
from anuga.pyvolution.util import mean
domain.reduction = mean 
#domain.reduction = min  #Looks better near steep slopes


######################
#Initial condition
#
print 'Initial condition'
t = 0.0
h0 = 10.
a = 3000.
g = 9.81
tau =0.001
B = 5
A = 0
p = sqrt(8*g*h0/a/a)
s = sqrt(p*p-tau*tau)/2
t = 0.


#Set bed-elevation and friction
def x_slope(x,y):
    n = x.shape[0]
    z = 0*x
    for i in range(n):
        z[i] = h0 - h0*(1.0 -x[i]*x[i]/a/a - y[i]*y[i]/a/a)
    return z

domain.set_quantity('elevation', x_slope)
domain.set_quantity('linear_friction', tau)


#Set the water stage
def stage(x,y):
    z = x_slope(x,y)
    n = x.shape[0]
    h = 0*x
    for i in range(n):
        h[i] = h0-B*B*exp(-tau*t)/2/g-1/g*(exp(-tau*t/2)*(B*s*cos(s*t) \
               +tau*B/2*sin(s*t)))*x[i] \
          -1/g*(exp(-tau*t/2)*(B*s*sin(s*t)-tau*B/2*cos(s*t)))*y[i]          
    if h[i] < z[i]:
        h[i] = z[i]
    return h   

domain.set_quantity('stage', stage)


############
#Boundary
domain.set_boundary({'exterior': Dirichlet_boundary([0.0, 0.0, 0.0])})
    

######################
#Evolution
import time
t0 = time.time()
for t in domain.evolve(yieldstep = 10.0, finaltime = 5000):
    domain.write_time()

print 'That took %.2f seconds' %(time.time()-t0)