#include <Partitioner.hpp>
#include <Messenger.hpp>

#include <parmetis.h>
#include <vector>

Partitioner::
Partitioner()
{
  std::cout << commRank() << " / " << commSize() << '\n' << std::flush;

  MPI_Barrier(MPI_COMM_WORLD);
  if (commSize() != 3) {
    Messenger::destroy();
    std::cerr << "Commsize=" << commSize() << ". Need 3 PE to test";
    std::exit(1);
  }

  idx_t wgtflag = 0;
  idx_t numflag = 0;
  idx_t ncon = 1;
  idx_t npart= 3;
  real_t wgts = 1./npart;
  std::vector<real_t> tpwgts{wgts, wgts, wgts};
  std::vector<real_t> ubvec{1.05};
  std::vector<idx_t> options{0,0,0};
  idx_t edgecut = 0;
  std::vector<idx_t> part(5);

  MPI_Comm mpi_comm_;
  MPI_Comm_dup(MPI_COMM_WORLD, &mpi_comm_);
  MPI_Comm* mpi_comm = & mpi_comm_;
  switch (commRank()) {
    case 0: {
      std::vector<idx_t> xadj   {0,2,5,8,11,13};
      std::vector<idx_t> adjncy {1,5,0,2,6,1,3,7,2,4,8,3,9};
      std::vector<idx_t> vtxdist{0,5,10,15};

      ParMETIS_V3_PartKway(&(vtxdist[0]), &(xadj[0]), &(adjncy[0]),
                           NULL, NULL, &wgtflag, &numflag,
                           &ncon, &npart, &(tpwgts[0]), &(ubvec[0]),
                           &(options[0]), &edgecut, &(part[0]), mpi_comm);
      break;
    }
    case 1: {
      std::vector<idx_t> xadj   {0,3,7,11,15,18};
      std::vector<idx_t> adjncy {0,6,10,1,5,7,11,2,6,8,12,3,7,9,13,4,8,14};
      std::vector<idx_t> vtxdist{0,5,10,15};

      ParMETIS_V3_PartKway(&(vtxdist[0]), &(xadj[0]), &(adjncy[0]),
                           NULL, NULL, &wgtflag, &numflag,
                           &ncon, &npart, &(tpwgts[0]), &(ubvec[0]),
                           &(options[0]), &edgecut, &(part[0]), mpi_comm);
      break;
    }
    case 2: {
      std::vector<idx_t> xadj   {0,2,5,8,11,13};
      std::vector<idx_t> adjncy {5,11,6,10,12,7,11,13,8,12,14,9,13};
      std::vector<idx_t> vtxdist{0,5,10,15};

      ParMETIS_V3_PartKway(&(vtxdist[0]), &(xadj[0]), &(adjncy[0]),
                           NULL, NULL, &wgtflag, &numflag,
                           &ncon, &npart, &(tpwgts[0]), &(ubvec[0]),
                           &(options[0]), &edgecut, &(part[0]), mpi_comm);
      break;
    }
    default: {
      std::cerr << "unexpected rank " << commRank() << "!\n";
      std::exit(0);
    }
  }

  std::cerr << commRank() << " FINISHED\n";
  MPI_Barrier(MPI_COMM_WORLD);

  Messenger::destroy();
  std::exit(0);
}