#!/usr/bin/env python
#########################################################
#
#  Main file for parallel mesh testing.
#
#  This is a modification of the run_parallel_advection.py
# file.
#
#
#  Authors: Linda Stals, Steve Roberts and Matthew Hardy,
# June 2005
#
#
#
#########################################################
import sys
import pypar    # The Python-MPI interface
import time


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

from Numeric import array
# pmesh

#from shallow_water import Domain

from shallow_water import Domain
from parallel_shallow_water import Parallel_Domain


# mesh partition routines
from parallel_meshes import parallel_rectangle


from pmesh_divide import pmesh_divide, pmesh_divide_steve

# read in the processor information

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

M = 20
N = M*numprocs

#N = M = 250

if myid == 0:
    print 'N == %d' %N

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



domain = Parallel_Domain(points, vertices, boundary,
                         full_send_dict  = full_send_dict,
                         ghost_recv_dict = ghost_recv_dict)

print 'number of triangles = ', domain.number_of_elements


rect = [ 0.0, 0.0, 1.0*numprocs, 1.0]
try:
    domain.initialise_visualiser(rect=rect)
    domain.visualiser.qcolor['stage'] = (0.0, 0.0, 0.8)
    domain.visualiser.scale_z['stage'] = 1.0
    domain.visualiser.scale_z['elevation'] = 0.05
except:
    print 'No visualiser'




domain.default_order = 2

#Boundaries
from parallel_shallow_water import Transmissive_boundary, Reflective_boundary

T = Transmissive_boundary(domain)
R = Reflective_boundary(domain)


domain.set_boundary( {'left': R, 'right': R, 'bottom': R, 'top': R} )
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, y0=0.0, y1=1.0, h=1.0):
        self.x0 = x0
        self.x1 = x1
        self.y0 = y0
        self.y1 = y1
        self.h  = h

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

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

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

yieldstep = 0.005
finaltime = 1.0

#Check that the boundary value gets propagated to all elements
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    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
    print 'Broadcast time %.2f seconds'%domain.communication_broadcast_time


pypar.finalize()