import sys
from os import sep
sys.path.append('..'+sep+'pyvolution')

#========================================================================
from config import g, epsilon
from Numeric import allclose, array, zeros, ones, Float
from parallel_meshes import parallel_rectangle

from advection import Domain
from parallel_advection import Parallel_Domain
from parallel_advection import Transmissive_boundary, Dirichlet_boundary

import pypar

numprocs = pypar.size()
myid = pypar.rank()
processor_name = pypar.Get_processor_name()



#N = 50
N = 50
M = 50

points, vertices, boundary, full_send_dict, ghost_recv_dict =  \
    parallel_rectangle(N, M, len1_g=1.0)

#Create advection domain with direction (1,-1)
domain = Parallel_Domain(points, vertices, boundary,
                         full_send_dict, ghost_recv_dict, velocity=[1.0, 0.0])

# Initial condition is zero by default

#turn on the visualisation
rect = [0.0, 0.0, 1.0, 1.0]
domain.initialise_visualiser(rect=rect)

#Boundaries
T = Transmissive_boundary(domain)
D = Dirichlet_boundary(array([1.0]))

domain.default_order = 2


domain.set_boundary( {'left': T, 'right': T, 'bottom': T, 'top': T} )
domain.check_integrity()

class Set_Stage:
    """Set an initial condition with constant water height, for x<x0
    """

    def __init__(self, x0=0.25, x1=0.5, h=1.0):
        self.x0 = x0
        self.x1 = x1
        self.h  = h

    def __call__(self, x, y):
        return self.h*((x>self.x0)&(x<self.x1))

domain.set_quantity('stage', Set_Stage(0.2,0.4,1.0))

if myid == 0:
    import time
    t0 = time.time()

#Check that the boundary value gets propagated to all elements
for t in domain.evolve(yieldstep = 0.1, finaltime = 3.0):
    if myid == 0:
        domain.write_time()

if myid == 0:
    print 'That took %.2f seconds' %(time.time()-t0)
    print 'Communication time %.2f seconds'%domain.communication_time
    print 'Reduction Communication time %.2f seconds'%domain.communication_reduce_time