#include <cassert>
#include <cstdlib>
#include <mpi.h>
#include <printf.hpp>
#include <hpc/matvec/densevector.hpp>
#include <hpc/matvec/iterators.hpp>
#include <hpc/matvec/gematrix.hpp>

int main(int argc, char** argv) {
   MPI_Init(&argc, &argv);

   int nof_processes; MPI_Comm_size(MPI_COMM_WORLD, &nof_processes);
   int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   assert(nof_processes == 2);

   using namespace hpc::matvec;
   using namespace hpc::aux;

   std::size_t nof_rows = 3;
   std::size_t nof_cols = 7;

   if (rank == 0) {
      GeMatrix<double> A(nof_rows, nof_cols, Order::RowMajor);
      for (auto [i, j, Aij]: A) {
	 Aij = i * 100 + j;
      }
      auto row = A.row(2, 0);
      auto col = A.col(0, 0);

      MPI_Send(/* row to rank 1 */);
      MPI_Send(/* col to rank 1 */);

      /* receive it back for verification */
      DenseVector<double> vec1(nof_cols), vec2(nof_rows);
      MPI_Status status;
      MPI_Recv(/* vec1 from rank 1 */);
      MPI_Recv(/* vec2 from rank 1 */);

      /* verify it */
      for (auto [i, xi]: vec1) {
	 if (vec1(i) != row(i)) {
	    fmt::printf("verification failed for row(%d): %lg vs %lg\n",
	       i, vec1(i), row(i));
	 }
      }
      for (auto [i, xi]: vec2) {
	 if (vec2(i) != col(i)) {
	    fmt::printf("verification failed for col(%d): %lg vs %lg\n",
	       i, vec2(i), col(i));
	 }
      }
   } else {
      DenseVector<double> vec1(nof_cols), vec2(nof_rows);
      MPI_Status status;
      MPI_Recv(/* vec1 from rank 0 */);
      MPI_Recv(/* vec2 from rank 0 */);

      /* send it back for verification */
      MPI_Send(/* vec1 to rank 0 */);
      MPI_Send(/* vec2 to rank 0 */);
   }
   MPI_Finalize();
}