"""Example of shallow water wave equation analytical solution
consists of a parabolic profile in a parabolic basin. Analytical
solutiuon to this problem was derived by Carrier and Greenspan
and used by Yoon and Chou.

   Copyright 2005
   Christopher Zoppou, Stephen Roberts, Ole Nielsen
   ANU, Geoscience Australia
   
"""

#---------------
# Module imports 
import sys
from os import sep
sys.path.append('..'+sep+'pyvolution')
from shallow_water import Domain, Dirichlet_boundary
from math import sqrt, cos, sin, pi
from mesh_factory import strang_mesh
from util import inside_polygon
from Numeric import asarray
from least_squares import Interpolation

#-------------------------------
# Set up the domain of triangles
# 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) 

#----------------
# Order of scheme
domain.default_order = 2

domain.smooth = True

#-------------------------------------
# Provide file name for storing output
domain.store = False
domain.format = 'sww'
domain.filename = 'yoon_mesh_second_order.2'
print 'Number of triangles = ', len(domain)

#----------------------------------------------------------
# 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 #domain.reduction = min  #Looks better near steep slopes


#------------------
# Initial condition
print 'Initial condition'
t = 0.0
D0 = 1.
L = 2500.
R0 = 2000.
g = 9.81

A = (L**4 - R0**4)/(L**4 + R0**4)
omega = 2./L*sqrt(2.*g*D0)
T = pi/omega

#------------------
# Set bed elevation
def x_slope(x,y):
    n = x.shape[0]
    z = 0*x
    for i in range(n):
        r = sqrt(x[i]*x[i] + y[i]*y[i])
        z[i] = -D0*(1.-r*r/L/L)
    return z
domain.set_quantity('elevation', x_slope)

#----------------------------
# Set the initial water level
def level(x,y):
    z = x_slope(x,y)
    n = x.shape[0]
    h = 0*x
    for i in range(n):
        r = sqrt(x[i]*x[i] + y[i]*y[i])
        h[i] = D0*((sqrt(1-A*A))/(1.-A*cos(omega*t))
                -1.-r*r/L/L*((1.-A*A)/((1.-A*cos(omega*t))**2)-1.))
    if h[i] < z[i]:
        h[i] = z[i]
    return h   
domain.set_quantity('stage', level)


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

#---------------------------------------------
# Find triangle that contains the point points
# and print to file
points = [0.,0.]
for n in range(len(domain.triangles)):
    t = domain.triangles[n]
    tri = [domain.coordinates[t[0]],domain.coordinates[t[1]],domain.coordinates[t[2]]]

    if inside_polygon(points,tri):
        print 'Point is within triangle with vertices '+'%s'%tri
        n_point = n
t = domain.triangles[n_point]
tri = [domain.coordinates[t[0]],domain.coordinates[t[1]],domain.coordinates[t[2]]]

filename=domain.filename
file = open(filename,'w')
    
#----------
# Evolution
import time
t0 = time.time()
for t in domain.evolve(yieldstep = 20.0, finaltime = 3000 ):
    domain.write_time()

    tri_array = asarray(tri)
    t_array = asarray([[0,1,2]])
    interp = Interpolation(tri_array,t_array,[points])

    stage = domain.get_quantity('stage').get_values()[n_point]
    xmomentum = domain.get_quantity('xmomentum').get_values()[n_point]
    ymomentum = domain.get_quantity('ymomentum').get_values()[n_point]

    interp_stage = interp.interpolate([[stage[0]],[stage[1]],[stage[2]]])
    interp_xmomentum = interp.interpolate([[xmomentum[0]],[xmomentum[1]],[xmomentum[2]]])
    interp_ymomentum = interp.interpolate([[ymomentum[0]],[ymomentum[1]],[ymomentum[2]]])
    
    file.write( '%10.6f   %10.6f  %10.6f   %10.6f\n'%(t,interp_stage[0][0],interp_xmomentum[0][0],interp_ymomentum[0][0]) )

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