#include "mpi_array_distrib.h"

#ifdef _PARALLEL
#define MASTER 0

    Distributor3D::Distributor3D(int nx, int ny, int nz, int ndx, int ndy, int ndz, int nb, int ndom, int dID)
    {
        NX = nx; NY = ny; NZ = nz; ndomX = ndx; ndomY = ndy; ndomZ = ndz; 
        nbuf = nb; nDomains = ndom; domainID = dID;

        if (nDomains != ndomX*ndomY*ndomZ) FATAL("nDomains != ndomX*ndomY*ndomZ !");
        if (domainID >= nDomains) FATAL("domainID >= nDomains !");

        if ( NX % ndomX != 0 ) FATAL("NX not a multiple of ndomX !");
        if ( NY % ndomY != 0 ) FATAL("NY not a multiple of ndomY !");
        if ( NZ % ndomZ != 0 ) FATAL("NZ not a multiple of ndomZ !");

        /* domainID = idx*ndomY*ndomZ + idy*ndomZ + idz */
        idZ = (domainID % ndomZ);
        idY = (((domainID-idZ)/ndomZ) % ndomY);
        idX = ((domainID - idZ)/ndomZ - idY)/ndomY;
        if (domainID != (idX*ndomY*ndomZ + idY*ndomZ + idZ)) FATAL("domainID doesn't match idx,idy,idz !");

        if ((ndomX>1)&&(ndomY==1)&&(ndomZ==1)) {
            splitDim = 0;
        } else if ((ndomX==1)&&(ndomY>1)&&(ndomZ==1)) {
            splitDim = 1;
        } else if ((ndomX==1)&&(ndomY==1)&&(ndomZ>1)) {
            splitDim = 2;
        } else {
            splitDim = -1;
        }

        if (ndomX==1) {
           in0 = ig0 = 0;
           in1 = ig1 = NX;
        } else {
           in0 = round(((NX*1.0)/ndomX)*(idX));
           in1 = round(((NX*1.0)/ndomX)*(idX+1));
           ig0 = in0 - nbuf;
           ig1 = in1 + nbuf;
        }
        if (ndomY==1) {
           jn0 = jg0 = 0;
           jn1 = jg1 = NY;
        } else {
           jn0 = round(((NY*1.0)/ndomY)*(idY));
           jn1 = round(((NY*1.0)/ndomY)*(idY+1));
           jg0 = jn0 - nbuf;
           jg1 = jn1 + nbuf;
        }
        if (ndomZ==1) {
           kn0 = kg0 = 0;
           kn1 = kg1 = NZ;
        } else {
           kn0 = round(((NZ*1.0)/ndomZ)*(idZ));
           kn1 = round(((NZ*1.0)/ndomZ)*(idZ+1));
           kg0 = kn0 - nbuf;
           kg1 = kn1 + nbuf;
        }

        if ( nbuf > in1-in0 ) FATAL("buffer size larger than native size in X !");
        if ( nbuf > jn1-jn0 ) FATAL("buffer size larger than native size in Y !");
        if ( nbuf > kn1-kn0 ) FATAL("buffer size larger than native size in Z !");

    }

    int Distributor3D::pack_array_all(double *f, double *buf)
    {
        int i, j, k;
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        if (splitDim==0)
            memcpy(buf, f, sizeof(double)*NX*NY*NZ);

        if (splitDim==1)
        {
            for(i=0;i<NX;i++) for(j=0;j<NY;j++) for(k=0;k<NZ;k++)
            {
                buf[j*NZ*NX+k*NX+i] = f[i*NY*NZ+j*NZ+k];
            }
        }

        if (splitDim==2)
        {
            for(i=0;i<NX;i++) for(j=0;j<NY;j++) for(k=0;k<NZ;k++)
            {
                buf[k*NX*NY+i*NY+j] = f[i*NY*NZ+j*NZ+k];
            }
        }

        return 0;
    }

    int Distributor3D::pack_array_native(double *f, double *buf)
    {
        int i, j, k;
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        if (splitDim==0)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                buf[i*NY*NZ+j*NZ+k] = f[i*NY*NZ+j*NZ+k];
            }
        }

        if (splitDim==1)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                buf[j*NZ*NX+k*NX+i] = f[i*NY*NZ+j*NZ+k];
            }
        }

        if (splitDim==2)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                buf[k*NX*NY+i*NY+j] = f[i*NY*NZ+j*NZ+k];
            }
        }

        return 0;
    }

    int Distributor3D::unpack_array_all(double *f, double *buf)
    {
        int i, j, k;
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        if (splitDim==0)
            memcpy(f, buf, sizeof(double)*NX*NY*NZ);

        if (splitDim==1)
        {
            for(i=0;i<NX;i++) for(j=0;j<NY;j++) for(k=0;k<NZ;k++)
            {
                f[i*NY*NZ+j*NZ+k] = buf[j*NZ*NX+k*NX+i];
            }
        }

        if (splitDim==2)
        {
            for(i=0;i<NX;i++) for(j=0;j<NY;j++) for(k=0;k<NZ;k++)
            {
                f[i*NY*NZ+j*NZ+k] = buf[k*NX*NY+i*NY+j];
            }
        }

        return 0;
    }

    int Distributor3D::unpack_array_native(double *f, double *buf)
    {
        int i, j, k;
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        if (splitDim==0)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                f[i*NY*NZ+j*NZ+k] = buf[i*NY*NZ+j*NZ+k];
            }
        }

        if (splitDim==1)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                f[i*NY*NZ+j*NZ+k] = buf[j*NZ*NX+k*NX+i];
            }
        }

        if (splitDim==2)
        {
            for(i=in0;i<in1;i++) for(j=jn0;j<jn1;j++) for(k=kn0;k<kn1;k++)
            {
                f[i*NY*NZ+j*NZ+k] = buf[k*NX*NY+i*NY+j];
            }
        }

        return 0;
    }

    int Distributor3D::distribute_array(double *f, double *buf)
    { /* Master node send data blocks to individual slave nodes */
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        if (domainID==MASTER) pack_array_all(f, buf);

        int chunksize, offset, source, dest;
        MPI_Status status;
        int tag1, tag2;

        chunksize = NX*NY*NZ/nDomains;
        tag1 = 1; tag2 = 2;

        /* call MPI to send the buffer from Master to Slaves, Master keeps 1st part */
        if (domainID==MASTER) {
            offset = chunksize;
            for (dest=1; dest<nDomains; dest++) {
                MPI_Send(&offset, 1, MPI_INT, dest, tag1, MPI_COMM_WORLD);
                MPI_Send(&buf[offset], chunksize, MPI_DOUBLE, dest, tag2, MPI_COMM_WORLD);
                printf("Sent %d elements to task %d offset= %d\n",chunksize,dest,offset);
                offset = offset + chunksize;
            }
        }

        /* Receive my portion of array from the master task */
        if (domainID != MASTER) {
            source = MASTER;
            MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
            MPI_Recv(&buf[offset], chunksize, MPI_DOUBLE, source, tag2, 
                     MPI_COMM_WORLD, &status);
        }

        if (domainID != MASTER) unpack_array_native(f, buf);
    }

    int Distributor3D::collect_array(double *f, double *buf)
    {/* Individual slave nodes send data blocks to master node */
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        pack_array_native(f, buf);

        int chunksize, offset, source, dest;
        MPI_Status status;
        int tag1, tag2;

        chunksize = NX*NY*NZ/nDomains;
        tag1 = 1; tag2 = 2;

        /* call MPI to send the buffer from Slaves to Master */
        if (domainID != MASTER) {
            offset = chunksize*domainID;
            dest = MASTER;
            MPI_Send(&offset, 1, MPI_INT, dest, tag1, MPI_COMM_WORLD);
            MPI_Send(&buf[offset], chunksize, MPI_DOUBLE, MASTER, tag2, MPI_COMM_WORLD);
        }

        /* Receive part of array from each Slave */
        if (domainID==MASTER) {
            for (source=1; source<nDomains; source++) {
                MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
                MPI_Recv(&buf[offset], chunksize, MPI_DOUBLE, source, tag2,
                         MPI_COMM_WORLD, &status);
            }
        }

        if (domainID==MASTER) unpack_array_all(f, buf);
    }

    int Distributor3D::update_neighbor(double *f, double *buf)
    {/* Each node communicate to its neighboring domains to exchange data in buffer layer */
        if ((splitDim!=0)&&(splitDim!=1)&&(splitDim!=2)) FATAL("unknow splitDim "<<splitDim<<" !");

        pack_array_native(f, buf);

        int chunksize, offset, source, dest, skinsize;
        MPI_Status status;
        int tag1, tag2;

        chunksize = NX*NY*NZ/nDomains;  
        tag1 = 1; tag2 = 2;

        if      (splitDim == 0) skinsize = nbuf*NY*NZ;
        else if (splitDim == 1) skinsize = nbuf*NX*NZ;
        else if (splitDim == 2) skinsize = nbuf*NX*NY;
        else                    skinsize = 0;

        /* call MPI to send the buffer between neighbors */
        {
            /* send to left neighbor */
            offset = chunksize*domainID;
            dest = (domainID+nDomains-1) % nDomains;
            MPI_Send(&offset, 1, MPI_INT, dest, tag1, MPI_COMM_WORLD);
            MPI_Send(&buf[offset], skinsize, MPI_DOUBLE, dest, tag2, MPI_COMM_WORLD);

            /* receive from right neighbor */
            source = (domainID+nDomains+1) % nDomains;
            MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
            INFO_Printf("domain[%d]: receive from %d  offset = %d skinsize = %d\n",domainID, source, offset, skinsize);
            MPI_Recv(&buf[offset], skinsize, MPI_DOUBLE, source, tag2, MPI_COMM_WORLD, &status);
        }

        {
            /* send to right neighbor */
            offset = chunksize*(domainID+1) - skinsize;
            dest = (domainID+nDomains+1) % nDomains;
            MPI_Send(&offset, 1, MPI_INT, dest, tag1, MPI_COMM_WORLD);
            MPI_Send(&buf[offset], skinsize, MPI_DOUBLE, dest, tag2, MPI_COMM_WORLD);

            /* receive from left neighbor */
            source = (domainID+nDomains-1) % nDomains;
            MPI_Recv(&offset, 1, MPI_INT, source, tag1, MPI_COMM_WORLD, &status);
            MPI_Recv(&buf[offset], skinsize, MPI_DOUBLE, source, tag2, MPI_COMM_WORLD, &status);
        }

        unpack_array_all(f, buf);
    }



#endif //_PARALLEL
