"""Example of shallow water wave equation.

Flat bed with rotational wind stress

"""

######################
# Module imports 
#
from mesh_factory import rectangular
from shallow_water import Domain, Reflective_boundary, Wind_stress

#Create basic mesh
N = 80
len = 100
points, vertices, boundary = rectangular(N, N, len, len)

#Create shallow water domain
domain = Domain(points, vertices, boundary)
domain.smooth = True
domain.visualise = False
domain.store = True
domain.default_order=2
domain.set_name('wind_rotation')

#Set initial conditions
domain.set_quantity('elevation', 0.0)
domain.set_quantity('level', 2.0)
domain.set_quantity('friction', 0.0)


#Variable windfield implemented using functions
def speed(x,y,t):
    """Large speeds halfway between center and edges
    Low speeds at center and edges
    """
    
    from math import sqrt, exp, cos, pi
    
    c = (len1/2, len2/2)
    r = sqrt((x - c[0])**2 + (y - c[1])**2)

    r /= max(len1,len2)

    factor = exp( -(r-0.15)**2 )
        
    #print x,y,factor

    return 4000 * factor * (cos(t*2*pi/150) + 2)


def phi(x,y,t):
    """Rotating field
    """
    from math import sqrt, atan, pi

    c = (len1/2, len2/2)
    r = sqrt((x - c[0])**2 + (y - c[1])**2)    
    
    xx = (x - c[0])/len1
    yy = (y - c[1])/len2

    angle = atan(yy/xx)

    if xx < 0:
        angle+=pi #atan in ]-pi/2; pi/2[

    #Take normal direction
    angle -= pi/2
    
    #Ensure positive radians    
    if angle < 0:
        angle += 2*pi

    return angle/pi*180

    
domain.forcing_terms.append( Wind_stress(speed, phi) )


#Add lateral wind gusts bearing 25 deg
def gust(x,y,t): 
    from math import sin, pi

    tt = sin(2*pi*t/200)

    if tt > 0.9:
        return 6000*tt
    else:
        return 0
    
domain.forcing_terms.append(Wind_stress(gust, 25))


######################
# Boundary conditions
Br = Reflective_boundary(domain)
domain.set_boundary({'left': Br, 'right': Br, 'top': Br, 'bottom': Br})

######################
#Evolution
for t in domain.evolve(yieldstep = 0.5, finaltime = 1000):
    domain.write_time()

print 'Done'