"""Example of shallow water wave equation analytical solution of the
circular hydraulic jump experimental data treated as a two-dimensional solution.

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

#-------------------------------
# Set up path and module imports
import sys
from os import sep
sys.path.append('..'+sep+'pyvolution')

from shallow_water import Domain, Dirichlet_Discharge_boundary
from shallow_water import Transmissive_Momentum_Set_Stage_boundary, Dirichlet_boundary
from math import pi, sqrt
from mesh_factory import strang_mesh


#---------
# Geometry
bed = ([.519, .519, .519, .519, .5192, .5194, .5196, .520, .5207, .5215, .5233, .5233])
distance = ([.08, .10, .11, .16, .21, .26, .31, .36, .41, .46, .50, .52])
n_bed = 12

#---------
# Case A.4
Q = 9.985/1000.0
wh0 = Q/(2.0*pi*0.1)
stage0 = bed[2] + 0.005
wh1 = -Q/(2.0*pi*0.503)
stage1 = 0.562
Manning = 0.009

#------------------
# Set up the 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('circular.pt')
domain = Domain(points, elements)

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

#----------------------
# Set a default tagging


for id, face in domain.boundary:
    domain.boundary[(id,face)] = 'outer'
    point = domain.get_vertex_coordinate(id,(face+1)%3)
    radius2 = point[0]*point[0] + point[1]*point[1]
    typical_outer = (id,face)
    if radius2 < 0.1:
        domain.boundary[(id,face)] = 'inner'
        typical_inner = (id,face)


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

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

#---------------------------
# Function for bed-elevation
def bed_z(x,y):
    n = x.shape[0]
    z = 0*x
    for i in range(n):
        r = sqrt(x[i]*x[i]+y[i]*y[i])
        for j in range(n_bed-1):
            if distance[j] <= r:
                if distance[j+1] > r:
                    z[i] = bed[j] + (bed[j+1] - bed[j])/(distance[j+1] - distance[j])*(r - distance[j])
    return z

domain.set_quantity('elevation', bed_z)

#---------
# Friction
domain.set_quantity('friction', Manning)

#---------------------------------
# Function for initial water elevation
# (stage)
def level(x,y):
    z = bed_z(x,y)
    n = x.shape[0]
    stage = 0*x
    for i in range(n):
        stage[i] = stage0
    return stage


#def outflow_stage(t):
#    return [stage1, 0 , 0]


domain.set_quantity('stage', level)

#---------------------------
# Set up boundary conditions
DD_BC_INNER = Dirichlet_Discharge_boundary(domain, stage0, wh0)
DD_BC_OUTER = Dirichlet_Discharge_boundary(domain, stage1, wh1)

domain.set_boundary({'inner': DD_BC_INNER, 'outer': DD_BC_OUTER})

#------------------
# Order of accuracy
domain.default_order = 1
domain.CFL = 0.75
#domain.beta_w = 0.5
#domain.beta_h = 0.2
domain.smooth = True


# domain.initialise_visualiser()
# domain.visualiser.coloring['stage'] = True
# domain.visualiser.scale_z['stage'] = 2.0
# domain.visualiser.scale_z['elevation'] = 0.05


#domain.initialise_visualiser()
#    #domain.visualiser.coloring['stage'] = True
#domain.visualiser.scale_z['stage'] = 2.0
#domain.visualiser.scale_z['elevation'] = 0.05
#
#
#from realtime_visualisation_new import Visualiser
##vxmom = Visualiser(domain,title='xmomentum',scale_z=10.0)
##vymom = Visualiser(domain,title='ymomentum',scale_z=10.0)

stage = domain.quantities['stage']

#----------
# Evolution
import time

f = open("circular_hydraulic_jump_true.txt","w")

t0 = time.time()
for t in domain.evolve(yieldstep = .1, finaltime = 3.0):
    domain.write_time()

    exp = '(xmomentum**2 + ymomentum**2)**0.5'
    radial_momentum = domain.create_quantity_from_expression(exp)

    print 'outer stage      ', stage.get_values(location='vertices',
                                            indices=[typical_outer[0]])
    print '      radial mom ', \
          radial_momentum.get_values(location='centroids',
                                     indices=[typical_outer[0]])

    print 'inner stage      ', stage.get_values(location='centroids',
                                            indices=[typical_inner[0]])
    print '      radial mom ', \
          radial_momentum.get_values(location='centroids',
                                     indices=[typical_inner[0]])

#    f.write('time = %25.15e wall clock time %g \n' % (domain.time, time.time()))
    f.write('%10.3f %25.15e %25.15e %25.15e %25.15e \n' % (domain.time, inner_stage, inner_radial_mom, outer_stage, outer_radial_mom))

    f.write('time = %25.15e wall clock time %g \n' % (domain.time, time.time()))
    f.write('%g \n' % stage.get_values(location='centroids',
                                   indices=[typical_outer[0]])[0])

f.close()

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