#!/usr/bin/env python
###
#########################################################
#
#  Main file for parallel mesh testing. Runs a shallow
# water simulation using the merimbula mesh
#
#
#
# *) The (new) files that have been added to manage the
# grid partitioning are
#    +) pmesh_divide_metis.py: subdivide a pmesh
#    +) build_submesh.py: build the submeshes on the host
# processor.
#    +) build_local.py: build the GA mesh datastructure
# on each processor.
#    +) build_commun.py: handle the communication between
# the host and processors
#
#  Authors: Linda Stals, Steve Roberts and Matthew Hardy,
# June 2005
#
#
#
#########################################################
import sys
from pypar_dist import pypar   # The Python-MPI interface
import time

from os import sep

# ADD directory ..../anuga/inundation to PYTHONPATH instead
#sys.path.append('..'+sep+'pyvolution')
#sys.path.append('..'+sep+'parallel')

# Numeric arrays

from Numeric import array, zeros, Float

# pmesh

from pyvolution.shallow_water import Domain
from parallel.parallel_shallow_water import Parallel_Domain
from pyvolution.pmesh2domain import pmesh_to_domain_instance

# Mesh partition routines

from parallel.pmesh_divide import pmesh_divide_metis
from parallel.build_submesh import build_submesh
from parallel.build_local   import build_local_mesh
from parallel.build_commun  import send_submesh, rec_submesh, extract_hostmesh

###############################
# Read in processor information
###############################

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

############################
# Set the initial conditions
############################

rect = zeros( 4, Float) # Buffer for results

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))

#######################
# Partition the mesh
#######################

if myid == 0:

    # Read in the test files

    filename = 'parallel/merimbula_10785_1.tsh'

    # Build the whole mesh
    
    mesh_full = pmesh_to_domain_instance(filename, Domain)

    # Define the domain boundaries for visualisation

    rect = array(mesh_full.xy_extent, Float)

    # Initialise the wave

    mesh_full.set_quantity('stage', Set_Stage(756000.0,756500.0,2.0))

    # Subdivide the mesh
    
    nodes, triangles, boundary, triangles_per_proc, quantities = \
         pmesh_divide_metis(mesh_full, numprocs)

    # Build the mesh that should be assigned to each processor,
    # this includes ghost nodes and the communicaiton pattern
    
    submesh = build_submesh(nodes, triangles, boundary,\
                            quantities, triangles_per_proc)

    # Send the mesh partition to the appropriate processor

    for p in range(1, numprocs):
      send_submesh(submesh, triangles_per_proc, p)

    # Build the local mesh for processor 0
    
    points, vertices, boundary, quantities, ghost_recv_dict, full_send_dict = \
              extract_hostmesh(submesh, triangles_per_proc)

else:
    
    # Read in the mesh partition that belongs to this
    # processor (note that the information is in the
    # correct form for the GA data structure)

    points, vertices, boundary, quantities, ghost_recv_dict, full_send_dict \
            = rec_submesh(0)


###########################################
# Start the computations on each subpartion
###########################################

# The visualiser needs to know the size of the whole domain

pypar.broadcast(rect,0)

# Build the domain for this processor

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

# Visualise the domain

try:
    domain.initialise_visualiser(rect=rect)
    domain.visualiser.scale_z['stage'] = 0.2
    domain.visualiser.scale_z['elevation'] = 0.05
except:
    print 'No visualiser'


domain.default_order = 1

# Define the boundaries, including the ghost boundary

from parallel.parallel_shallow_water import Transmissive_boundary, Reflective_boundary

T = Transmissive_boundary(domain)
R = Reflective_boundary(domain)
domain.set_boundary( {'outflow': R, 'inflow': R, 'inner':R, 'exterior': R, \
                      'open':R, 'ghost':None} )


# Set the initial quantities

domain.set_quantity('stage', quantities['stage'])
domain.set_quantity('elevation', quantities['elevation'])

domain.store = False

# Set the number of time steps, as well as the start and end time

t0 = time.time()
yieldstep = 1
finaltime = 90


# Start the evolve calculations

for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    if myid == 0:
        domain.write_time()

# Print some timing statistics

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