""" 
Small test of directional flow aka momentum jet
"""

#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------
from anuga.abstract_2d_finite_volumes.mesh_factory import rectangular_cross
from anuga.shallow_water import Domain
from anuga.shallow_water.shallow_water_domain import Reflective_boundary
from anuga.shallow_water.shallow_water_domain import Dirichlet_boundary
from anuga.shallow_water.shallow_water_domain import Inflow, General_forcing
from anuga.culvert_flows.culvert_polygons import create_culvert_polygons
from anuga.utilities.polygon import plot_polygons

from math import pi, sqrt, cos, sin

#------------------------------------------------------------------------------
# Setup computational domain
#------------------------------------------------------------------------------



length = 40.
width = 5.

#dx = dy = 1           # Resolution: Length of subdivisions on both axes
#dx = dy = .5           # Resolution: Length of subdivisions on both axes
dx = dy = .25           # Resolution: Length of subdivisions on both axes
#dx = dy = .1           # Resolution: Length of subdivisions on both axes

# OLE.... How do I refine the resolution... in the area where I have the Culvert Opening ???......
#  Can I refine in a X & Y Range ???
points, vertices, boundary = rectangular_cross(int(length/dx), int(width/dy),
                                               len1=length, len2=width)
domain = Domain(points, vertices, boundary)   
domain.set_name('jet_example')                 # Output name
domain.set_default_order(2)
domain.H0 = 0.01
domain.tight_slope_limiters = True
domain.set_minimum_storable_height(0.02)
print 'Size', len(domain)


#------------------------------------------------------------------------------
# Setup initial conditions
#------------------------------------------------------------------------------
domain.set_quantity('elevation', 0.0)
domain.set_quantity('friction', 0.01)         # Constant friction 
domain.set_quantity('stage',
                    expression='elevation')   # Dry initial condition


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

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

#------------------------------------------------------------------------------
# Setup Application of  specialised forcing terms
#------------------------------------------------------------------------------

# This tests straight Inflow in m^3/s
center=(6.0, 2.0)
r = 1.2
fixed_flow = Inflow(domain,
                    rate=6.00,
                    center=center,
                    radius=r)
domain.forcing_terms.append(fixed_flow)


# This tests momentum update (6, 3) in m^2/s/s
xmom_forcing = General_forcing(domain,
                               'xmomentum',
                               rate=6.00,
                               center=center,
                               radius=r)
ymom_forcing = General_forcing(domain,
                               'ymomentum',
                               rate=3.00,
                               center=center,
                               radius=r)

domain.forcing_terms.append(xmom_forcing)
domain.forcing_terms.append(ymom_forcing)                               



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



for t in domain.evolve(yieldstep = 0.1, finaltime = 20):

    if t > 6:
        # Swing forcing rate around with time
        xmom_forcing.rate = 5*cos(t/5*2*pi)    
        ymom_forcing.rate = 5*sin(t/5*2*pi)
    
    print domain.timestepping_statistics()