"""Example of shallow water wave equation analytical solution of the
one-dimensional Thacker and Greenspan wave run-up treated as a two-dimensional solution.

   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, Reflective_boundary, Time_boundary
from math import sqrt, cos, sin, pi
from mesh_factory import strang_mesh

#Convenience functions
def imag(a):
    return a.imag

def real(a):
    return a.real


######################
# 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('Run-up.pt')
points, elements = strang_mesh('strang_7389.pt')
domain = Domain(points, elements)

domain.default_order = 2
domain.smooth = True

#Set a default tagging
epsilon = 1.0e-12

for id, face in domain.boundary:
    domain.boundary[(id,face)] = 'external'
    if domain.get_vertex_coordinate(id,(face+1)%3)[0] < -200.0+1.0e-10:
        domain.boundary[(id,face)] = 'left'




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

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

#Define the boundary condition
def stage_setup(x,t_star1):
    vh = 0
    alpha = 0.1
    eta = 0.1
    a = 1.5*sqrt(1.+0.9*eta)
    l_0 = 200.
    ii = complex(0,1)
    g = 9.81
    v_0 = sqrt(g*l_0*alpha)
    v1 = 0.

    sigma_max = 100.
    sigma_min = -100.
    for j in range (1,50):
        sigma0 = (sigma_max+sigma_min)/2.
        lambda_prime = 2./a*(t_star1/sqrt(l_0/alpha/g)+v1)
        sigma_prime = sigma0/a
        const = (1.-ii*lambda_prime)**2+sigma_prime**2

        v1 = 8.*eta/a*imag(1./const**(3./2.) \
            -3./4.*(1.-ii*lambda_prime)/const**(5./2.))

        x1 = -v1**2/2.-a**2*sigma_prime**2/16.+eta \
            *real(1.-2.*(5./4.-ii*lambda_prime) \
            /const**(3./2.)+3./2.*(1.-ii*lambda_prime)**2 \
            /const**(5./2.))

        neta1 = x1 + a*a*sigma_prime**2/16.

        v_star1 = v1*v_0
        x_star1 = x1*l_0
        neta_star1 = neta1*alpha*l_0
        stage = neta_star1
        z = stage - x_star1*alpha
        uh = z*v_star1

        if x_star1-x > 0:
            sigma_max = sigma0
        else:
            sigma_min = sigma0

        if abs(abs(sigma0)-100.) < 10:

#     solution does not converge because bed is dry
            stage = 0.
            uh = 0.
            z = 0.

    return [stage, uh, vh]

def boundary_stage(t):
    x = -200
    return stage_setup(x,t)


######################
#Initial condition
print 'Initial condition'
t_star1 = 0.0
slope = -0.1

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

domain.set_quantity('elevation', x_slope)

#Set the water depth
print 'Initial water depth'
def stage(x,y):
    z = x_slope(x,y)
    n = x.shape[0]
    w = 0*x
    for i in range(n):
        w[i], uh, vh = stage_setup(x[i],t_star1)
        h = w[i] - z[i]
        if h < 0:
            h = 0
            w[i] = z[i]
    return w

domain.set_quantity('stage', stage)


#####################
#Set up boundary conditions
Br = Reflective_boundary(domain)
Bw = Time_boundary(domain, boundary_stage)
domain.set_boundary({'left': Bw, 'external': Br})



#domain.visualise = True

######################
#Evolution
import time
t0 = time.time()
for t in domain.evolve(yieldstep = 1., finaltime = 100):
    domain.write_time()
    print boundary_stage(domain.time)

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