#!/usr/bin/env python
"""A rectangular domain partitioner and associated communication
functionality for solving PDEs in (1D,2D) using FDM
written in the python language

The global solution domain is assumed to be of rectangular shape,
where the number of cells in each direction is stored in nx, ny, nz

The numerical scheme is fully explicit

Authors
-------

 * Xing Cai
 * Min Ragan-Kelley

"""
from __future__ import print_function
import time

from numpy import zeros, ascontiguousarray, frombuffer
try:
    from mpi4py import MPI
except ImportError:
    pass
else:
    mpi = MPI.COMM_WORLD

class RectPartitioner:
    """
    Responsible for a rectangular partitioning of a global domain,
    which is expressed as the numbers of cells in the different
    spatial directions. The partitioning info is expressed as an
    array of integers, each indicating the number of subdomains in
    one spatial direction.
    """

    def __init__(self, my_id=-1, num_procs=-1, \
                 global_num_cells=[], num_parts=[]):
        self.nsd = 0
        self.my_id = my_id
        self.num_procs = num_procs
        self.redim (global_num_cells, num_parts)

    def redim (self, global_num_cells, num_parts):
        nsd_ = len(global_num_cells)
#        print("Inside the redim function, nsd=%d" %nsd_)

        if nsd_<1 | nsd_>3 | nsd_!=len(num_parts):
            print('The input global_num_cells is not ok!')
            return

        self.nsd = nsd_
        self.global_num_cells = global_num_cells
        self.num_parts = num_parts

    def prepare_communication (self):
        """
        Find the subdomain rank (tuple) for each processor and
        determine the neighbor info.
        """
        
        nsd_ = self.nsd
        if nsd_<1:
            print('Number of space dimensions is %d, nothing to do' %nsd_)
            return
        
        self.subd_rank = [-1,-1,-1]
        self.subd_lo_ix = [-1,-1,-1]
        self.subd_hi_ix = [-1,-1,-1]
        self.lower_neighbors = [-1,-1,-1]
        self.upper_neighbors = [-1,-1,-1]

        num_procs = self.num_procs
        my_id = self.my_id

        num_subds = 1
        for i in range(nsd_):
            num_subds = num_subds*self.num_parts[i]
        if my_id==0:
            print("# subds=", num_subds)
            # should check num_subds againt num_procs

        offsets = [1, 0, 0]

        # find the subdomain rank
        self.subd_rank[0] = my_id%self.num_parts[0]
        if nsd_>=2:
            offsets[1] = self.num_parts[0]
            self.subd_rank[1] = my_id/offsets[1]
        if nsd_==3:
            offsets[1] = self.num_parts[0]
            offsets[2] = self.num_parts[0]*self.num_parts[1]
            self.subd_rank[1] = (my_id%offsets[2])/self.num_parts[0]
            self.subd_rank[2] = my_id/offsets[2]
                    
        print("my_id=%d, subd_rank: "%my_id, self.subd_rank)
        if my_id==0:
            print("offsets=", offsets)

        # find the neighbor ids
        for i in range(nsd_):
            rank = self.subd_rank[i]
            if rank>0:
                self.lower_neighbors[i] = my_id-offsets[i]
            if rank<self.num_parts[i]-1:
                self.upper_neighbors[i] = my_id+offsets[i]
                
            k = self.global_num_cells[i]/self.num_parts[i]
            m = self.global_num_cells[i]%self.num_parts[i]
            
            ix = rank*k+max(0,rank+m-self.num_parts[i])
            self.subd_lo_ix[i] = ix
            
            ix = ix+k
            if rank>=(self.num_parts[i]-m):
                ix = ix+1  # load balancing
            if rank<self.num_parts[i]-1:
                ix = ix+1  # one cell of overlap
            self.subd_hi_ix[i] = ix

        print("subd_rank:",self.subd_rank,\
              "lower_neig:", self.lower_neighbors, \
              "upper_neig:", self.upper_neighbors)
        print("subd_rank:",self.subd_rank,"subd_lo_ix:", self.subd_lo_ix, \
              "subd_hi_ix:", self.subd_hi_ix)

        
class RectPartitioner1D(RectPartitioner):
    """
    Subclass of RectPartitioner, for 1D problems
    """
    def prepare_communication (self):
        """
        Prepare the buffers to be used for later communications
        """
        
        RectPartitioner.prepare_communication (self)
        
        if self.lower_neighbors[0]>=0:
            self.in_lower_buffers = [zeros(1, float)]
            self.out_lower_buffers = [zeros(1, float)]
        if self.upper_neighbors[0]>=0:
            self.in_upper_buffers = [zeros(1, float)]
            self.out_upper_buffers = [zeros(1, float)]

    def get_num_loc_cells(self):
        return [self.subd_hi_ix[0]-self.subd_lo_ix[0]]


class RectPartitioner2D(RectPartitioner):
    """
    Subclass of RectPartitioner, for 2D problems
    """
    def prepare_communication (self):
        """
        Prepare the buffers to be used for later communications
        """
        
        RectPartitioner.prepare_communication (self)
        
        self.in_lower_buffers = [[], []]
        self.out_lower_buffers = [[], []]
        self.in_upper_buffers = [[], []]
        self.out_upper_buffers = [[], []]

        size1 = self.subd_hi_ix[1]-self.subd_lo_ix[1]+1
        if self.lower_neighbors[0]>=0:
            self.in_lower_buffers[0] = zeros(size1, float)
            self.out_lower_buffers[0] = zeros(size1, float)
        if self.upper_neighbors[0]>=0:
            self.in_upper_buffers[0] = zeros(size1, float)
            self.out_upper_buffers[0] = zeros(size1, float)

        size0 = self.subd_hi_ix[0]-self.subd_lo_ix[0]+1
        if self.lower_neighbors[1]>=0:
            self.in_lower_buffers[1] = zeros(size0, float)
            self.out_lower_buffers[1] = zeros(size0, float)
        if self.upper_neighbors[1]>=0:
            self.in_upper_buffers[1] = zeros(size0, float)
            self.out_upper_buffers[1] = zeros(size0, float)

    def get_num_loc_cells(self):
        return [self.subd_hi_ix[0]-self.subd_lo_ix[0],\
                self.subd_hi_ix[1]-self.subd_lo_ix[1]]


class MPIRectPartitioner2D(RectPartitioner2D):
    """
    Subclass of RectPartitioner2D, which uses MPI via mpi4py for communication
    """
    
    def __init__(self, my_id=-1, num_procs=-1,
                 global_num_cells=[], num_parts=[],
                 slice_copy=True):
        RectPartitioner.__init__(self, my_id, num_procs,
                                 global_num_cells, num_parts)
        self.slice_copy = slice_copy
        
    def update_internal_boundary (self, solution_array):
        nsd_ = self.nsd
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]

        # communicate in the x-direction first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            mpi.Isend(self.out_lower_buffers[0], lower_x_neigh)
            
        if upper_x_neigh>-1:
            mpi.Recv(self.in_upper_buffers[0], upper_x_neigh)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            mpi.Isend(self.out_upper_buffers[0], upper_x_neigh)

        if lower_x_neigh>-1:
            mpi.Recv(self.in_lower_buffers[0], lower_x_neigh)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]

        # communicate in the y-direction afterwards
        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            mpi.Isend(self.out_lower_buffers[1], lower_y_neigh)
            
        if upper_y_neigh>-1:
            mpi.Recv(self.in_upper_buffers[1], upper_y_neigh)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            mpi.Isend(self.out_upper_buffers[1], upper_y_neigh)
            
        if lower_y_neigh>-1:
            mpi.Recv(self.in_lower_buffers[1], lower_y_neigh)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]

class ZMQRectPartitioner2D(RectPartitioner2D):
    """
    Subclass of RectPartitioner2D, which uses 0MQ via pyzmq for communication
    The first two arguments must be `comm`, an EngineCommunicator object,
    and `addrs`, a dict of connection information for other EngineCommunicator
    objects.
    """

    def __init__(self, comm, addrs, my_id=-1, num_procs=-1,
                 global_num_cells=[], num_parts=[],
                 slice_copy=True):
        RectPartitioner.__init__(self, my_id, num_procs,
                                 global_num_cells, num_parts)
        self.slice_copy = slice_copy
        self.comm = comm # an Engine
        self.addrs = addrs
    
    def prepare_communication(self):
        RectPartitioner2D.prepare_communication(self)
        # connect west/south to east/north
        west_id,south_id = self.lower_neighbors[:2]
        west = self.addrs.get(west_id, None)
        south = self.addrs.get(south_id, None)
        self.comm.connect(south, west)
    
    def update_internal_boundary_x_y (self, solution_array):
        """update the inner boundary with the same send/recv pattern as the MPIPartitioner"""
        nsd_ = self.nsd
        dtype = solution_array.dtype
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]
        trackers = []
        flags = dict(copy=False, track=False)
        # communicate in the x-direction first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            t = self.comm.west.send(self.out_lower_buffers[0], **flags)
            trackers.append(t)
            
        if upper_x_neigh>-1:
            msg = self.comm.east.recv(copy=False)
            self.in_upper_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            t = self.comm.east.send(self.out_upper_buffers[0], **flags)
            trackers.append(t)
            

        if lower_x_neigh>-1:
            msg = self.comm.west.recv(copy=False)
            self.in_lower_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]
        
        # communicate in the y-direction afterwards
        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            t = self.comm.south.send(self.out_lower_buffers[1], **flags)
            trackers.append(t)
            
            
        if upper_y_neigh>-1:
            msg = self.comm.north.recv(copy=False)
            self.in_upper_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            t = self.comm.north.send(self.out_upper_buffers[1], **flags)
            trackers.append(t)
            
        if lower_y_neigh>-1:
            msg = self.comm.south.recv(copy=False)
            self.in_lower_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]
        
        # wait for sends to complete:
        if flags['track']:
            for t in trackers:
                t.wait()

    def update_internal_boundary_send_recv (self, solution_array):
        """update the inner boundary, sending first, then recving"""
        nsd_ = self.nsd
        dtype = solution_array.dtype
        if nsd_!=len(self.in_lower_buffers) | nsd_!=len(self.out_lower_buffers):
            print("Buffers for communicating with lower neighbors not ready")
            return
        if nsd_!=len(self.in_upper_buffers) | nsd_!=len(self.out_upper_buffers):
            print("Buffers for communicating with upper neighbors not ready")
            return

        loc_nx = self.subd_hi_ix[0]-self.subd_lo_ix[0]
        loc_ny = self.subd_hi_ix[1]-self.subd_lo_ix[1]

        lower_x_neigh = self.lower_neighbors[0]
        upper_x_neigh = self.upper_neighbors[0]
        lower_y_neigh = self.lower_neighbors[1]
        upper_y_neigh = self.upper_neighbors[1]
        trackers = []
        flags = dict(copy=False, track=False)
        
        # send in all directions first
        if lower_x_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[0] = ascontiguousarray(solution_array[1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_lower_buffers[0][i] = solution_array[1,i]
            t = self.comm.west.send(self.out_lower_buffers[0], **flags)
            trackers.append(t)
        
        if lower_y_neigh>-1:
            if self.slice_copy:
                self.out_lower_buffers[1] = ascontiguousarray(solution_array[:,1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_lower_buffers[1][i] = solution_array[i,1]
            t = self.comm.south.send(self.out_lower_buffers[1], **flags)
            trackers.append(t)
            
        if upper_x_neigh>-1:
            if self.slice_copy:
                self.out_upper_buffers[0] = ascontiguousarray(solution_array[loc_nx-1,:])
            else:
                for i in xrange(0,loc_ny+1):
                    self.out_upper_buffers[0][i] = solution_array[loc_nx-1,i]
            t = self.comm.east.send(self.out_upper_buffers[0], **flags)
            trackers.append(t)
        
        if upper_y_neigh>-1:
            if self.slice_copy:
                self.out_upper_buffers[1] = ascontiguousarray(solution_array[:,loc_ny-1])
            else:
                for i in xrange(0,loc_nx+1):
                    self.out_upper_buffers[1][i] = solution_array[i,loc_ny-1]
            t = self.comm.north.send(self.out_upper_buffers[1], **flags)
            trackers.append(t)
        
        
        # now start receiving
        if upper_x_neigh>-1:
            msg = self.comm.east.recv(copy=False)
            self.in_upper_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[loc_nx,:] = self.in_upper_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[loc_nx,i] = self.in_upper_buffers[0][i]

        if lower_x_neigh>-1:
            msg = self.comm.west.recv(copy=False)
            self.in_lower_buffers[0] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[0,:] = self.in_lower_buffers[0]
            else:
                for i in xrange(0,loc_ny+1):
                    solution_array[0,i] = self.in_lower_buffers[0][i]
        
        if upper_y_neigh>-1:
            msg = self.comm.north.recv(copy=False)
            self.in_upper_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,loc_ny] = self.in_upper_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,loc_ny] = self.in_upper_buffers[1][i]
            
        if lower_y_neigh>-1:
            msg = self.comm.south.recv(copy=False)
            self.in_lower_buffers[1] = frombuffer(msg, dtype=dtype)
            if self.slice_copy:
                solution_array[:,0] = self.in_lower_buffers[1]
            else:
                for i in xrange(0,loc_nx+1):
                    solution_array[i,0] = self.in_lower_buffers[1][i]
        
        # wait for sends to complete:
        if flags['track']:
            for t in trackers:
                t.wait()
    
    # use send/recv pattern instead of x/y sweeps
    update_internal_boundary = update_internal_boundary_send_recv