#ifndef HPC_MPI_COPY_H
#define HPC_MPI_COPY_H 1

#include <cassert>
#include <mpi.h>
#include <hpc/matvec/gematrix.h>
#include <hpc/matvec/isgematrix.h>
#include <hpc/matvec/copy.h>
#include <hpc/mpi/gematrix.h>
#include <hpc/mpi/matrix.h>

namespace hpc { namespace mpi {

// scatter
template <typename T, typename Index>
void
copy(const hpc::matvec::GeMatrix<T, Index> &A, GeMatrix<T, Index> &B)
{
    if (B.grid.rank == 0) {
        MPI_Request requests[B.grid.nof_processes-1];
        int ri = 0;
        for (int i = 0; i < B.grid.nof_processes; ++i) {
            int coords[2];
            MPI_Cart_coords(B.grid.comm, i, 2, coords);
            auto A_ = A(B.getRowOffset(coords[0]), B.getColOffset(coords[1]),
                        B.getNumRows(coords[0]), B.getNumCols(coords[1]));
            if (i == 0) {
                hpc::matvec::copy(A_, B.buffer);
            } else {
                MPI_Isend((void*)&A_(0, 0), 1, get_type(A_),
                          i, 0, B.grid.comm, &requests[ri++]);
            }
        }
        for (auto& request: requests) {
            MPI_Status status;
            MPI_Wait(&request, &status);
        }
    } else {
        MPI_Status status;
        T &B00 = B.buffer(0, 0);
        MPI_Recv((void*)&B00, 1, get_type(B.buffer), 0, 0, B.grid.comm, &status);
    }
}

// gather
template <typename T, typename Index>
void
copy(const GeMatrix<T, Index> &A, hpc::matvec::GeMatrix<T, Index> &B)
{
    if (A.grid.rank == 0) {
        MPI_Request requests[A.grid.nof_processes-1];
        int ri = 0;
        for (int i = 0; i < A.grid.nof_processes; ++i) {
            int coords[2];
            MPI_Cart_coords(A.grid.comm, i, 2, coords);
            auto B_ = B(A.getRowOffset(coords[0]), A.getColOffset(coords[1]),
                        A.getNumRows(coords[0]), A.getNumCols(coords[1]));
            if (i == 0) {
                hpc::matvec::copy(A.buffer, B_);
            } else {
                MPI_Irecv(&B_(0, 0), 1, get_type(B_),
                          i, 0, A.grid.comm, &requests[ri++]);
            }
        }
        for (auto& request: requests) {
            MPI_Status status;
            MPI_Wait(&request, &status);
        }
    } else {
        const T &A00 = A.buffer(0, 0);
        MPI_Send((void *)&A00, 1, get_type(A.buffer), 0, 0, A.grid.comm);
    }
}

} } // namespaces mpi, hpc

#endif // HPC_MPI_COPY_H