"""Simple water flow example using ANUGA

Water driven up a linear slope and time varying boundary,
similar to a beach environment
"""


#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------

from pmesh.mesh_interface import create_mesh_from_regions
from pyvolution.shallow_water import Domain
from pyvolution.shallow_water import Reflective_boundary
from pyvolution.shallow_water import Dirichlet_boundary
from pyvolution.shallow_water import Time_boundary
from pyvolution.shallow_water import Transmissive_boundary


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

waveheight = 1800.0
depth_east_edge = -4000.
timeend = 1000.0
west = 0.
east = 80000.
south = 0.
north = 5000.
polygon = [[east,north],[west,north],[west,south],[east,south]]
meshname = 'test.msh'

create_mesh_from_regions(polygon,
                         boundary_tags={'top': [0],
                                        'left': [1],
                                        'bottom': [2],
                                        'right': [3]},
                         maximum_triangle_area=200000,
                         filename=meshname)
                         #interior_regions=interior_regions)

domain = Domain(meshname,use_cache=True,verbose = True)
domain.set_name('test')                     # Output to test.sww
domain.set_datadir('.')                     # Use current directory for output
domain.set_quantities_to_be_stored(['stage',# Store all conserved quantities
                                    'xmomentum',
                                    'ymomentum'])   


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

def topography(x,y):
    slope = depth_east_edge/((east-west)/2.)
    c = -(west+east)/2.*slope
    return slope*x+c                         # linear bed slope

domain.set_quantity('elevation', topography) # Use function for elevation
domain.set_quantity('friction', 0. )         # Constant friction 
domain.set_quantity('stage', 0.0)            # Constant initial stage

#------------------------------------------------------------------------------
# Setup boundary conditions
#------------------------------------------------------------------------------

from math import sin, pi
Br = Reflective_boundary(domain)      # Solid reflective wall
Bt = Transmissive_boundary(domain)    # Continue all values on boundary 
Bd = Dirichlet_boundary([0.,0.,0.])   # Constant boundary values
Bw = Time_boundary(domain=domain,     # Time dependent boundary  
                   f=lambda t: [((10<t<20)*waveheight), 0.0, 0.0])

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


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

stagestep = []
for t in domain.evolve(yieldstep = 10, finaltime = timeend):
    domain.write_time()

#-----------------------------------------------------------------------------
# Interrogate solution
#-----------------------------------------------------------------------------

# Gauge line
def gauge_line(west,east,north,south):
    from Numeric import arange
    gaugex = arange(west,(west+east)/2.,1000.)
    gauges = []
    gaugey = []
    for x in gaugex:
        y = (north+south)/2.
        gaugey.append(y)
        gauges.append([x,y])
    return gauges, gaugex, gaugey

gauges, gaugex, gaugey = gauge_line(west,east,north,south)

from pyvolution.util import file_function

f = file_function('test.sww',
                  quantities = ['stage', 'elevation', 'xmomentum', 'ymomentum'],
                  interpolation_points = gauges,
                  verbose = True,
                  use_cache = True)

from pylab import plot, xlabel, ylabel, title, ion, close, savefig, figure, axis
ion()

maxw = []
minw = []
maxd = []
count = 0
for k, g in enumerate(gauges):
    stage = []
    elevation = []
    depths = []
    bed = topography(g[0],g[1])
    for i, t in enumerate(f.get_time()):
        w = f(t, point_id = k)[0]
        elev = f(t, point_id = k)[1]
        depth = w-elev
        stage.append(w)
        elevation.append(elev)
        depths.append(elev)

    if max(stage)-bed <= 0.2:
        count+=1
        posx=g[0]
        loc = k  
    maxw.append(max(stage))
    minw.append(min(stage))
    maxd.append(max(depths))
    

filename = 'maxrunup'+str(waveheight)+str(depth_east_edge)+'.csv'
fid = open(filename,'w')    
s = 'Depth at eastern boundary,Waveheight,Runup distance,Runup height\n'
fid.write(s)
runup_distance = posx 
runup_height = topography(runup_distance,(north+south)/2.)
print 'run up height:   ', runup_height 
if runup_distance < (east+west)/2.:
    runup_distance_from_coast = (east+west)/2.-runup_distance
else:
    runup_distance_from_coast = runup_distance-(east+west)/2.
print 'run up distance from coastline: ', runup_distance_from_coast
s = '%.2f,%.2f,%.2f,%.2f\n' %(depth_east_edge,waveheight,runup_distance_from_coast,runup_height)
fid.write(s)

figure(1)
plot(gaugex,maxw,'g+',gaugex,topography(gaugex,(north+south)/2.),'r-')
xlabel('x')
ylabel('stage')
title('Maximum stage for gauge line')
#axis([33000, 47000, -1000, 3000])
savefig('max_stage')

close('all')