"""Example of shallow water wave equation.

This is called Netherlands because it shows a dam with a gap in it and
stylised housed behind it and below the water surface.

"""

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

from anuga.shallow_water import Domain
from anuga.shallow_water import Reflective_boundary, Dirichlet_boundary
from anuga.abstract_2d_finite_volumes.mesh_factory import rectangular_cross
import os

#from anuga.visualiser import RealtimeVisualiser
#import rpdb
#rpdb.set_active()


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

N = 150 # size = 45000
N = 130 # size = 33800
N = 600 # Size = 720000
N = 150

points, elements, boundary = rectangular_cross(N, N)
domain = Domain(points, elements, boundary, use_inscribed_circle=True)

domain.check_integrity()

import sys
base = os.path.basename(sys.argv[0])
domain.simulation_name = 'netherlands'
domain.set_name(os.path.splitext(__file__)[0])
domain.set_timestepping_method('rk2')
domain.set_default_order(2)
domain.set_store_vertices_uniquely(True) # Store as internally represented
domain.tight_slope_limiters = True
print domain.statistics()


# Setup order and all the beta's for the limiters (these should become defaults

#domain.beta_w      = 1.0
#domain.beta_w_dry  = 0.2
#domain.beta_uh     = 1.0
#domain.beta_uh_dry = 0.2
#domain.beta_vh     = 1.0
#domain.beta_vh_dry = 0.2

#domain.alpha_balance = 100.0



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

class Weir:
    """Set a bathymetry for simple weir with a hole.
    x,y are assumed to be in the unit square
    """

    def __init__(self, stage):
        self.inflow_stage = stage

    def __call__(self, x, y):
        from Numeric import zeros, Float

        N = len(x)
        assert N == len(y)

        z = zeros(N, Float)
        for i in range(N):
            z[i] = -x[i]/20  # General slope

            # Flattish bit to the left
            if x[i] <= 0.3:
                #z[i] = -x[i]/5
                z[i] = -x[i]/20


            # Weir
            if x[i] > 0.3 and x[i] < 0.4:
                z[i] = -x[i]/20+1.2

            # Dip
            #if x[i] > 0.6 and x[i] < 0.9:
            #    z[i] = -x[i]/20-0.5  #-y[i]/5

            # Hole in weir
            #if x[i] > 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
            if x[i] > 0.3 and x[i] < 0.4 and y[i] > 0.4 and y[i] < 0.6:
                #z[i] = -x[i]/5
                z[i] = -x[i]/20

            # Poles
            #if x[i] > 0.65 and x[i] < 0.8 and y[i] > 0.55 and y[i] < 0.65 or\
            #       x[i] > 0.75 and x[i] < 0.9 and y[i] > 0.35 and y[i] < 0.45:
            #    z[i] = -x[i]/20+0.4

            if (x[i] - 0.72)**2 + (y[i] - 0.6)**2 < 0.05**2:# or\
                   #x[i] > 0.75 and x[i] < 0.9 and y[i] > 0.35 and y[i] < 0.45:
                z[i] = -x[i]/20+0.4


            # Wall
            if x[i] > 0.995:
                z[i] = -x[i]/20+0.3

        return z/2


inflow_stage = 0.5
manning = 0.0

domain.set_quantity('elevation', Weir(inflow_stage))
domain.set_quantity('friction', manning)
domain.set_quantity('stage', expression='elevation + 0.0')


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

Br = Reflective_boundary(domain)
Bd = Dirichlet_boundary([inflow_stage, 0.0, 0.0]) # Constant inflow
domain.set_boundary({'left': Bd, 'right': Br, 'bottom': Br, 'top': Br})

#-------------------------------------------------------------------------------
# Copy scripts to time stamped output directory and capture screen
# output to file
#-------------------------------------------------------------------------------
import time
time = time.strftime('%Y%m%d_%H%M%S',time.localtime())

output_dir = 'dam_break_'+time
output_file = 'dam_break'

#copy_code_files(output_dir,__file__)
#start_screen_catcher(output_dir+'_')


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


if N <= 150:
    # Initialise real-time visualiser 

    #pass
    visualise = True
    if visualise:
        from anuga.visualiser import RealtimeVisualiser
        vis = RealtimeVisualiser(domain)
        vis.render_quantity_height("elevation", offset=0.01, dynamic=False)
        vis.render_quantity_height("stage", dynamic=True)
        vis.colour_height_quantity('stage', (0.2, 0.2, 0.8))
        vis.start()
        import time
        time.sleep(2.0)
    
    



import time
t0 = time.time()

for t in domain.evolve(yieldstep = 0.05, finaltime = 5.0):
    print domain.timestepping_statistics()
    #print domain.quantities['stage'].get_values(location='centroids',
                                                #indices=[0])
    if visualise:
        vis.update()						

        

if visualise: vis.evolveFinished()

print 'That took %.2f seconds' %(time.time()-t0)

#vis.join()