#include <iostream>
#include <mpi.h>

#define BUFFER_SIZE 1

void my_MPI_Alltoall(int *sendbuf, int *recvbuf, int count, MPI_Comm comm,
                     int rank) {
  for (int i = 0; i < count; ++i) {
    recvbuf[rank] = sendbuf[rank];
    if (i == rank) {
      // recv from other proc
      for (int j = 0; j < count; ++j) {
        if (j != rank) {
          MPI_Recv(recvbuf + j, 1, MPI_INT, j, j, MPI_COMM_WORLD,
                   MPI_STATUS_IGNORE);
        }
      }
    } else {
      // send to one proc
      MPI_Send(sendbuf + i, 1, MPI_INT, i, rank, MPI_COMM_WORLD);
    }
  }
}

int main(int argc, char *argv[]) {
  int rank, size;
  int *send_buffer;
  int *recv_buffer_test;
  int *recv_buffer_val;
  // MPI init
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  // end MPI init
  send_buffer = new int[BUFFER_SIZE * size];
  recv_buffer_test = new int[BUFFER_SIZE * size];
  recv_buffer_val = new int[BUFFER_SIZE * size];

  // init send_buffer
  for (int i = 0; i < size; ++i) {
    send_buffer[i] = rank + size * i;
  }
  // use my_MPI_Alltoall
  my_MPI_Alltoall(send_buffer, recv_buffer_test, size, MPI_COMM_WORLD, rank);

  // use MPI_Alltoall
  MPI_Alltoall(send_buffer, BUFFER_SIZE, MPI_INT, recv_buffer_val, BUFFER_SIZE,
               MPI_INT, MPI_COMM_WORLD);

  // validate data
  bool foundErr = false;
  for (int i = 0; i < size; ++i) {
    if (recv_buffer_test[i] != recv_buffer_val[i]) {
      std::cout << "proc" << rank << ": "
                << "test[" << i << "]:" << recv_buffer_test[i] << ", val[" << i
                << "]:" << recv_buffer_val[i] << std::endl;
      foundErr = true;
    }
  }
  if (!foundErr) {
    std::cout <<"proc" << rank << ": no err found" << std::endl;
  }
  delete[] send_buffer;
  delete[] recv_buffer_val;
  delete[] recv_buffer_test;
  MPI_Finalize();
}