"""Simple water flow example using ANUGA

Water driven towards a continental slope and into a harbour
environment. Set up as per:
P.L.-F Liu 
Effects of the Continental Shelf on Harbor Resonance, in
Tsunamis - Their Science and Engineering,
edited by K. Iida and T. Iwasaki, 303-314. 
1983, Terra Scientific Publishing Company, Tokyo.

"""

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

from anuga.pmesh.mesh_interface import create_mesh_from_regions
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 anuga.shallow_water import Transmissive_boundary

waveheight = 1.0
harbour_length = 2000.0 
shelf_length = 10000.0
bottom_length = 10000.0
harbour_width = 250.0
shelf_depth = -100. 
harbour_depth = -10.
if harbour_depth > shelf_depth: print 'WARNING: depths not consistent'
#------------------------------------------------------------------------------
# Setup computational domain
#------------------------------------------------------------------------------

east = 10000.+harbour_length+shelf_length+bottom_length
west = 0.
south = 0.
north = 5000.
polygon = [[east,north],[west,north],[west,south],[east,south]]
meshname = 'harbour_test.msh'

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

domain = Domain(meshname,use_cache=True,verbose = True)
domain.set_name('harbour')                  # Output to harbour.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
#------------------------------------------------------------------------------
# if can't do this, then can set up x = arange(west,east,xstep)
# and y = arange(south,north,ystep) and use Geospatial data to set elevation

def topography(x,y):
    from Numeric import zeros, Float
    z = zeros(len(x), Float)
    
    xcentre = (west+east)/2.
    ycentre = (south+north)/2.
    ycentreup = ycentre+harbour_width/2.
    ycentredown = ycentre-harbour_width/2.
    for i, xi in enumerate(x):
        yi = y[i]
        if xi >= xcentre: z[i] = shelf_depth
        if xi > west and xi <= xcentre-harbour_length: 
            if yi < ycentredown: z[i] = 0.0
            if yi > ycentreup: z[i] = 0.0
            if yi > ycentredown and yi < ycentreup: z[i] = harbour_depth
        if xi > xcentre-harbour_length and xi < xcentre: z[i] = harbour_depth
        
    return z

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 = 100.0):
    domain.write_time()

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

# Gauge line
def gauge_line(west,east,north,south):
    from Numeric import arange
    gaugex = arange(east,west,-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 anuga.abstract_2d_finite_volumes.util import file_function

f = file_function('harbour.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
ion()

maxw = []
minw = []
maxd = []
count = 0
for k, g in enumerate(gauges):
    stage = []
    elevation = []
    depths = []
    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)
    maxw.append(max(stage))
    minw.append(min(stage))
    maxd.append(max(depths))

filename = 'point'+str(waveheight)+'.csv'
fid = open(filename,'w')    
s = 'Waveheight,\n'
fid.write(s)

figure(1)
plot(gaugex[:loc],maxw[:loc],'g-')
xlabel('x')
ylabel('stage')
title('Maximum stage for gauge line')
savefig('max_stage')

close('all')