/*
To be included by a file that desires rendezvous rank exchange functionality.
Need to define HAVEWANT_t, HAVEWANT_alltoallv, and HAVEWANT2Int
*/

#ifdef have2want_cpp
#error "This implementation can only be included once"
/* The static function names could involve a macro name. */
#endif

#define have2want_cpp

/*

A rank owns a set of HAVEWANT_t keys and wants information associated with
a set of HAVEWANT_t keys owned by unknown ranks.  Owners do not know which
ranks want their information. Ranks that want info do not know which ranks
own that info.

The have_to_want function returns two new vectors of keys along with
associated count and displacement vectors of length nhost and nhost+1
respectively. Note that a send_to_want_displ[i+1] =
  send_to_want_cnt[i] + send_to_want_displ[i] .

send_to_want[send_to_want_displ[i] to send_to_want_displ[i+1]] contains
the keys from this rank for which rank i wants information.

recv_from_have[recv_from_have_displ[i] to recv_from_have_displ[i+1] contains
the keys from which rank i is sending information to this rank.

Note that on rank i, the order of keys in the rank j area of send_to_want
is the same order of keys on rank j in the ith area in recv_from_have.

The rendezvous_rank function is used to parallelize this computation
and minimize memory usage so that no single rank ever needs to know all keys.
*/

#ifndef HAVEWANT_t
#define HAVEWANT_t int
#endif

// round robin default rendezvous rank function
static int default_rendezvous(HAVEWANT_t key) {
  return key % nrnmpi_numprocs;
}

static int* cnt2displ(int* cnt) {
  int* displ = new int[nrnmpi_numprocs+1];
  displ[0] = 0;
  for (int i=0; i < nrnmpi_numprocs; ++i) {
    displ[i+1] = displ[i] + cnt[i];
  }
  return displ;
}

static int* srccnt2destcnt(int* srccnt) {
  int* destcnt = new int[nrnmpi_numprocs];
  nrnmpi_int_alltoall(srccnt, destcnt, 1);
  return destcnt;
}

static void rendezvous_rank_get(HAVEWANT_t* data, int size, 
  HAVEWANT_t* &sdata, int* &scnt, int* &sdispl,
  HAVEWANT_t* &rdata, int* &rcnt, int* &rdispl,
  int (*rendezvous_rank)(HAVEWANT_t)
){
  int nhost = nrnmpi_numprocs;
  int rank = nrnmpi_myid;

  // count what gets sent
  scnt = new int[nhost];
  for (int i=0; i < nhost; ++i) { scnt[i] = 0; }
  for (int i=0; i < size; ++i) {
    int r = (*rendezvous_rank)(data[i]);
    ++scnt[r];
  }

  sdispl = cnt2displ(scnt);
  rcnt = srccnt2destcnt(scnt);
  rdispl = cnt2displ(rcnt);
  sdata = new HAVEWANT_t[sdispl[nhost]];
  rdata = new HAVEWANT_t[rdispl[nhost]];
  // scatter data into sdata by recalculating scnt.
  for (int i=0; i < nhost; ++i) { scnt[i] = 0; }
  for (int i=0; i < size; ++i) {
    int r = (*rendezvous_rank)(data[i]);
    sdata[sdispl[r] + scnt[r]] = data[i];
    ++scnt[r];
  }      
  HAVEWANT_alltoallv(sdata, scnt, sdispl, rdata, rcnt, rdispl);
}

static void have_to_want(
  HAVEWANT_t* have, int have_size, HAVEWANT_t* want, int want_size,
  HAVEWANT_t* &send_to_want, int* &send_to_want_cnt, int* &send_to_want_displ,
  HAVEWANT_t* &recv_from_have, int* &recv_from_have_cnt, int* &recv_from_have_displ,
  int (*rendezvous_rank)(HAVEWANT_t)
){
  // 1) Send have and want to the rendezvous ranks.
  // 2) Rendezvous rank matches have and want.
  // 3) Rendezvous ranks tell the want ranks which ranks own the keys
  // 4) Ranks that want tell owner ranks where to send.

  int nhost = nrnmpi_numprocs;
  int rank = nrnmpi_myid;

  // 1) Send have and want to the rendezvous ranks.
  HAVEWANT_t *have_s_data, *have_r_data;
  int *have_s_cnt, *have_s_displ, *have_r_cnt, *have_r_displ;
  rendezvous_rank_get(have, have_size,
    have_s_data, have_s_cnt, have_s_displ,
    have_r_data, have_r_cnt, have_r_displ,
    rendezvous_rank
  );
  // assume it is an error if two ranks have the same key so create
  // hash table of key2rank. Will also need it for matching have and want
  HAVEWANT2Int havekey2rank = HAVEWANT2Int(have_r_displ[nhost]);
  for (int r=0; r < nhost; ++r) {
    for (int i=0; i < have_r_cnt[r]; ++i) {
      HAVEWANT_t key = have_r_data[have_r_displ[r] + i];
      int srcrank;
      if (havekey2rank.find(key, srcrank)) {
	char buf[200];
sprintf(buf, "key %lld owned by multiple ranks\n", (long long)key);
        hoc_execerror(buf, 0);
      }
      havekey2rank[key] = r;
    }
  }
  delete [] have_r_data;
  delete [] have_r_cnt;
  delete [] have_r_displ;

  HAVEWANT_t *want_s_data, *want_r_data;
  int *want_s_cnt, *want_s_displ, *want_r_cnt, *want_r_displ;
  rendezvous_rank_get(want, want_size,
    want_s_data, want_s_cnt, want_s_displ,
    want_r_data, want_r_cnt, want_r_displ,
    rendezvous_rank
  );

  // 2) Rendezvous rank matches have and want.
  //    we already have made the havekey2rank map.
  // Create an array parallel to want_r_data which contains the ranks that
  // have that data.
  int n = want_r_displ[nhost];
  int* want_r_ownerranks = new int[n];
  for (int r=0; r < nhost; ++r) {
    for (int i=0; i < want_r_cnt[r]; ++i) {
      int ix = want_r_displ[r] + i;
      HAVEWANT_t key = want_r_data[ix];
      int srcrank;
      if (!havekey2rank.find(key, srcrank)) {
	char buf[200];
sprintf(buf, "key = %lld is wanted but does not exist\n", (long long)key);
        hoc_execerror(buf, 0);
      }
      want_r_ownerranks[ix] = srcrank;
    }
  }
  delete [] want_r_data;

  // 3) Rendezvous ranks tell the want ranks which ranks own the keys
  // The ranks that want keys need to know the ranks that own those keys.
  // The want_s_ownerranks will be parallel to the want_s_data.
  // That is, each item defines the rank from which information associated
  // with that key is coming from
  int* want_s_ownerranks = new int[want_s_displ[nhost]];
  nrnmpi_int_alltoallv(want_r_ownerranks, want_r_cnt, want_r_displ,
    want_s_ownerranks, want_s_cnt, want_s_displ);
  delete [] want_r_ownerranks;
  delete [] want_r_cnt;
  delete [] want_r_displ;

  // 4) Ranks that want tell owner ranks where to send.
  // Finished with the rendezvous ranks. The ranks that want keys know the
  // owner ranks for those keys. The next step is for the want ranks to
  // tell the owner ranks where to send.
  // The parallel want_s_ownerranks and want_s_data are now uselessly ordered
  // by rendezvous rank. Reorganize so that want ranks can tell owner ranks
  // what they want.
  n = want_s_displ[nhost];
  delete [] want_s_displ;
  for (int i=0; i < nhost; ++i) { want_s_cnt[i] = 0; }
  HAVEWANT_t* old_want_s_data = want_s_data;
  want_s_data = new HAVEWANT_t[n];
  // compute the counts
  for (int i=0; i < n; ++i) {
    int r = want_s_ownerranks[i];
    ++want_s_cnt[r];
  }
  want_s_displ = cnt2displ(want_s_cnt);
  for (int i=0; i < nhost; ++i) { want_s_cnt[i] = 0; } // recount while filling
  for (int i=0; i < n; ++i) {
    int r = want_s_ownerranks[i];
    HAVEWANT_t key = old_want_s_data[i];
    want_s_data[want_s_displ[r] + want_s_cnt[r]] = key;
    ++want_s_cnt[r];
  }
  delete [] old_want_s_data;
  want_r_cnt = srccnt2destcnt(want_s_cnt);
  want_r_displ = cnt2displ(want_r_cnt);
  want_r_data = new HAVEWANT_t[want_r_displ[nhost]];
  HAVEWANT_alltoallv(want_s_data, want_s_cnt, want_s_displ,
    want_r_data, want_r_cnt, want_r_displ);
  // now the want_r_data on the have_ranks are grouped according to the ranks
  // that want those keys.

  send_to_want = want_r_data;
  send_to_want_cnt = want_r_cnt;
  send_to_want_displ = want_r_displ;
  recv_from_have = want_s_data;
  recv_from_have_cnt = want_s_cnt;
  recv_from_have_displ = want_s_displ;
}
