"""Simple water flow example using ANUGA

Water flowing down a channel with more complex topography
"""

#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------
from anuga.abstract_2d_finite_volumes.mesh_factory import rectangular_cross
from anuga.shallow_water import Domain
from anuga.shallow_water import Reflective_boundary
from anuga.shallow_water import Dirichlet_boundary
from anuga.shallow_water import Time_boundary

from math import cos, sin, pi, sqrt

#------------------------------------------------------------------------------
# Setup computational domain
#------------------------------------------------------------------------------
length = 40.
width = 5.
#dx = dy = .05           # Resolution: Length of subdivisions on both axes
dx = dy = .5           # Resolution: Length of subdivisions on both axes

points, vertices, boundary = rectangular_cross(int(length/dx), int(width/dy),
                                               len1=length, len2=width)
domain = Domain(points, vertices, boundary)   
domain.set_name('channel3')                  # Output name
print domain.statistics()


#------------------------------------------------------------------------------
# Setup initial conditions
#------------------------------------------------------------------------------
def topography(x,y):
    """Complex topography defined by a function of vectors x and y
    """

    # Geometric parameters
    gradient = 0.02 
    weir_height = 0.5
    obstruction_height = 0.0
    pole_height = 0.0

    # Upward slope    
    z = gradient*x

    N = len(x)
    for i in range(N):
    
        # Weir
        if 10 < x[i] < 12:
            z[i] += weir_height

        # Obstruction
        if 27 < x[i] < 29 and y[i] > 3:
            z[i] += obstruction_height        
        
        # Pole
        if (x[i] - 34)**2 + (y[i] - 2)**2 < 0.4**2:
            z[i] += pole_height

        



    return z


domain.set_quantity('elevation', topography)  # Use function for elevation
domain.set_quantity('friction', 0.01)         # Constant friction 
domain.set_quantity('stage',
                    expression='elevation')   # Dry initial condition


#------------------------------------------------------------------------------
# Setup boundary conditions
#------------------------------------------------------------------------------
Bi = Dirichlet_boundary([1.0, 0, 0])          # Inflow
Br = Reflective_boundary(domain)              # Solid reflective wall
Bo = Dirichlet_boundary([-1, 0, 0])           # Outflow

def wave(t):

    A = 1.5 # Amplitude [m] (Wave height)
    T = 30  # Wave period [s]

    if t < 120:
        return [A*sin(2*pi*t/T) + 1, 0, 0] 
    else:
        return [0.0, 0, 0]

Bt = Time_boundary(domain, f=wave)

domain.set_boundary({'left': Bt, 'right': Bo, 'top': Br, 'bottom': Br})

        

#------------------------------------------------------------------------------
# Evolve system through time
#------------------------------------------------------------------------------

washed_away = False

for t in domain.evolve(yieldstep = 0.1, finaltime = 120.0):
    print domain.timestepping_statistics()

    S = domain.get_quantity('stage').get_values(interpolation_points=[[40, 2.5]])
    E = domain.get_quantity('elevation').get_values(interpolation_points=[[40, 2.5]])
    
    depth = S-E

    print '-------------'
    print 'depth', depth
    print '-------------'    
    
    #print 'elevation', E
    
    if depth > 1.63/3:

        print 'Water level will sweep Fiona away',
        if washed_away is True:
            print ' - AGAIN!'
        else:
            print

        if washed_away is False:

            washaway_time = t
            print 'Getting momentum'            
            uh = domain.get_quantity('xmomentum').get_values(interpolation_points=[[40, 2.5]])
            vh = domain.get_quantity('ymomentum').get_values(interpolation_points=[[40, 2.5]])

            print 'Computing velocity'            
            u = uh/depth
            v = vh/depth
            print 'velocity: (%f, %f' %(u,v)
            washaway_velocity = sqrt(u*u + v*v)            
            print 'done'
            
        washed_away = True
        #raw_input('press key')

if washed_away is True:
    print 'Fiona got swept away at time = %.1f sec with a velocity of %.2f m/s'\
          %(washaway_time, washaway_velocity)