"""Simple water flow example using ANUGA
"""

#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------
from anuga.abstract_2d_finite_volumes.mesh_factory import *
from anuga.shallow_water import *
from anuga.shallow_water.shallow_water_domain import *
from math import *
from numpy import *
import numpy
from matplotlib import *
from pylab import *
import csv

#Parameters

filename = "WORKING-RIP-LAB_Expt-Geometry_Triangular_Mesh"
location_of_shore = 140    #the position along the y axis of the shorefront
sandbar = 1.2 #height of sandbar
sealevel = 0 #height of coast above sea level
steepness = 8000 #period of sandbar - larger number gives smoother slope - longer period 
halfchannelwidth = 5
bank_slope = 0.1
simulation_length = 1
timestep = 1

#------------------------------------------------------------------------------
# Setup computational domain
#------------------------------------------------------------------------------
length = 120
width = 170
seafloor_resolution = 60.0 # Resolution: Max area of triangles in the mesh for seafloor
feature_resolution = 1.0
beach_resolution = 10.0

sea_boundary_polygon = [[0,0],[length,0],[length,width],[0,width]]
feature_boundary_polygon = [[0,100],[length,100],[length,150],[0,150]]
beach_interior_polygon = [[0,150],[length,150],[length,width],[0,width]]

meshname = str(filename)+'.msh'

#interior regions
feature_regions = [[feature_boundary_polygon, feature_resolution],
                   [beach_interior_polygon, beach_resolution]]

create_mesh_from_regions(sea_boundary_polygon,
                         boundary_tags={'bottom': [0],
                                        'right' : [1],
                                        'top' :   [2],
                                        'left':   [3]},
                         maximum_triangle_area = seafloor_resolution,
                         filename = meshname,
                         interior_regions = feature_regions,
                         use_cache = False,
                         verbose = False)
domain = Domain(meshname, use_cache=True, verbose=True)
domain.set_name(filename)                  # Output name
print domain.statistics()

#------------------------------------------------------------------------------
# Setup initial conditions
#------------------------------------------------------------------------------

def topography(x,y):
    """Complex topography defined by a function of vectors x and y."""

    #general slope and buildings
    z=0.05*(y-(location_of_shore))
    
    N = len(x)
    for i in range(N):
        if y[i] < 25:
            z[i] = (0.2*(y[i]-25)) + 0.05*(y[i]-(location_of_shore))
    for i in range(N):
        if y[i]>150:
            z[i] = (0.1*(y[i]-150)) + 0.05*(y[i] - (location_of_shore))

    return z


def topography3(x,y):
    z=0*x
    
    N = len(x)
    for i in range(N):
        if -1*(bank_slope)*x[i] + 112 < y[i] < -1*(bank_slope)*x[i] + 124 and 0<x[i]<((length/2)-halfchannelwidth):
            z[i] += sandbar*((cos((y[i]-118)/steepness)))
    for i in range(N):
        if -1*(bank_slope)*(x[i]-(length/2)) + (-1*(bank_slope)*(length/2)+112) < y[i] < -1*(bank_slope)*(x[i]-(length/2)) + (-1*(bank_slope)*(length/2)+124) and ((length/2)+halfchannelwidth)<x[i]<183:
            z[i] += sandbar*(cos((y[i]-118)/steepness))
    return z

domain.set_quantity('elevation', topography)           # elevation is a function
domain.set_quantity('friction', 0.01)                  # Constant friction

domain.add_quantity('elevation', topography3)

domain.set_quantity('stage', 0)   # Dry initial condition

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


amplitude = 0.4 #amplitude of wave (wave height m)
period =  5     #wave period (sec)

def wave(t):

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

    if t < 30000000000:
        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': Br, 'right': Br, 'top': Bo, 'bottom': Bt})

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

list1 = [x for x in csv.reader(open('New_gauges.csv','r'),dialect='excel',delimiter=",") ]

print list1

u = numpy.zeros(len(list1), 'd')
v = numpy.zeros(len(list1), 'd')


for t in domain.evolve(yieldstep = (timestep), finaltime = (simulation_length)):
    print domain.timestepping_statistics()
    S = domain.get_quantity('stage').get_values(interpolation_points=list1)
    E = domain.get_quantity('elevation').get_values(interpolation_points=list1)
    depth = S-E

    uh = domain.get_quantity('xmomentum').get_values(interpolation_points=list1)
    vh = domain.get_quantity('ymomentum').get_values(interpolation_points=list1)
    u += uh/depth
    v += vh/depth

n_time_intervals = (simulation_length)/(timestep)
u_average = u / (n_time_intervals)
v_average = v / (n_time_intervals)

print "there were", n_time_intervals, "time steps"

print "sum y velocity", v
print "average y velocity", v_average
print "sum x velocity", u
print "average x velocity", u_average

x_output = file("x_velocity.txt", 'w')
y_output = file("y_velocity.txt", 'w')

print >> x_output, " "
print >> y_output, " "

print >> x_output, u_average
print >> y_output, v_average

X = [x[0] for x in csv.reader(open('New_gauges.csv','r'),dialect='excel',delimiter=",") ]
Y = [x[1] for x in csv.reader(open('New_gauges.csv','r'),dialect='excel',delimiter=",") ]
U = u_average.tolist()
V = v_average.tolist()

print "U = ", U
print "U has type", type(U)


from matplotlib import *
from pylab import *

figure()
quiver(X,Y,U,V)
show()