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

   Copyright 2004
   Christopher Zoppou, Stephen Roberts, Ole Nielsen, Duncan Gray
   Geoscience Australia

   Modified by Sudi Mungkasi, ANU, 2011
   
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
import anuga
from math import sqrt, cos, sin, pi
from numpy import asarray
from time import localtime, strftime, gmtime

#-------------------------------------------------------------------------------
# Copy scripts to time stamped output directory and capture screen
# output to file
#-------------------------------------------------------------------------------
time = strftime('%Y%m%d_%H%M%S',localtime())

output_dir = 'transient_carrier_greenspan_'+time
output_file = 'transient_carrier_greenspan'

#anuga.copy_code_files(output_dir,__file__)
#start_screen_catcher(output_dir+'_')


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

def real(a):
    return a.real


#------------------------------------------------------------------------------
# Setup domain
#------------------------------------------------------------------------------
dx = 10.
dy = dx
L = 100.
W = 10*dx

# structured mesh
points, vertices, boundary = anuga.rectangular_cross(int(L/dx), int(W/dy), L, W, (0.0, -W/2))

domain = anuga.Domain(points, vertices, boundary) 

domain.set_name(output_file)                
domain.set_datadir(output_dir) 

#------------------------------------------------------------------------------
# Setup Algorithm
#------------------------------------------------------------------------------
domain.set_default_order(2)
domain.set_timestepping_method(2)
domain.set_beta(0.7)
domain.set_CFL(0.6)


#-----------------------
#Define the boundary condition
#-----------------------
def stage_setup(x,t):
    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/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
#-------------------------------------------------------------------------------
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
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)


#-----------------------------------------------------------------------------
# Setup boundary conditions
#------------------------------------------------------------------------------
Br = anuga.Reflective_boundary(domain)
Bt = anuga.Time_boundary(domain, boundary_stage)

# Associate boundary tags with boundary objects
domain.set_boundary({'left': Bt, 'right': Br, 'top': Br, 'bottom': Br})

##visualize = True
##if visualize:
##    from anuga.visualiser import RealtimeVisualiser
##    vis = RealtimeVisualiser(domain)
##    vis.render_quantity_height("elevation", zScale=3.0, offset = 0.01, dynamic=False)
##    vis.render_quantity_height("stage", zScale = 3.0, dynamic=True, opacity = 0.6, wireframe=False)
##    #vis.colour_height_quantity('stage', (lambda q:q['stage'], 1.0, 2.0))
##    vis.colour_height_quantity('stage', (0.4, 0.6, 0.4))
##    vis.start()
##    time.sleep(2.0)

#domain.visualise = True


#------------------------------------------------------------------------------
# Evolve system through time
#------------------------------------------------------------------------------
import time
t0 = time.time()
for t in domain.evolve(yieldstep = 1., finaltime = 1000):
    domain.write_time()
    #print boundary_stage(domain.time)
    #if visualize: vis.update()
    
#if visualize: vis.evolveFinished()    
print 'That took %.2f seconds' %(time.time()-t0)