/*
 * Copyright (c) 2008 The Trustees of Indiana University and Indiana
 *                    University Research and Technology
 *                    Corporation.  All rights reserved.
 *
 * Author(s): Torsten Hoefler <htor@cs.indiana.edu>
 *
 */
#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>

int MPI_Virtual_graph_create(MPI_Comm comm_old, int n, int *nodes,
                             int *degrees, int *tgtnodes, int reorder,
                             MPI_Comm *newcomm) {
  int rank, p;
  MPI_Comm_rank(comm_old, &rank);
  MPI_Comm_size(comm_old, &p);
  int *sindex = (int*)calloc(1,p*sizeof(int));
  int *rindex = (int*)calloc(1,p*sizeof(int));
  int *rindexscan = (int*)calloc(1,p*sizeof(int));

  for(int i = 0; i<n; i++) {
    sindex[nodes[i]]++;
  }
  MPI_Allreduce(sindex, rindex, p, MPI_INT, MPI_SUM, comm_old);
  MPI_Exscan(sindex, rindexscan, p, MPI_INT, MPI_SUM, comm_old);
  /*for(int i = 0; i<p; i++) {
    printf("[%i] rindexscan[%i] = %i\n", rank, i, rindexscan[i]);
  }*/


  int sum=0;
  for(int i = 0; i<p; i++) {
    sum += rindex[i];
  }
  int *edges= (int*)calloc(1,sum*sizeof(int));
  int *redges= (int*)calloc(1,sum*sizeof(int));

  int tgtindex=-1;
  for(int i = 0; i<n; i++) {
    for(int j = 0; j<degrees[i]; j++) {
      tgtindex ++;
      // find index where the current target host (nodes[i]) starts
      int ind=0;
      for(int prev=0; prev<nodes[i];prev++) {
        ind+=rindex[prev];
      }
      ind+=rindexscan[nodes[i]]+j;

      //printf("rank: %i, host: %i, rindexscan: %i, ind: %i (%i): %i\n", rank, nodes[i], rindexscan[nodes[i]], ind, tgtindex, tgtnodes[tgtindex]);

      // fill edges array
      edges[ind] = tgtnodes[tgtindex];
    }
  }

  MPI_Allreduce(edges, redges, sum, MPI_INT, MPI_SUM, comm_old);

  // transform rindex into index array
  for(int i=1; i<p;i++) {
    rindex[i]+=rindex[i-1];
  }

  /*if(!rank) printf("indices:\n");
  if(!rank) for(int i=0; i<p; i++) printf("%i\n", rindex[i]);
  if(!rank) printf("edges:\n");
  if(!rank) for(int i=0; i<sum; i++) printf("%i\n", redges[i]);*/

  
  MPI_Graph_create(comm_old, p, rindex, redges, reorder, newcomm);

}

int MPI_Virtual_neighbors_count(MPI_Comm comm, int *inneighbors, int *outneighbors) {

  int rank, p;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &p);

  MPI_Graph_neighbors_count(comm, rank, outneighbors);

  // this is local
  // get all neighbors who point at me ... should have been cached
  // during creation - it's only a POC
  *inneighbors=0;
  for(int i=0; i<p; i++) {
    int nneighbors;
    MPI_Graph_neighbors_count(comm, i, &nneighbors);
    int *neighbors = (int*)malloc(nneighbors*sizeof(int));
    MPI_Graph_neighbors(comm, i, nneighbors, neighbors);
    for(int j=0; j<nneighbors; j++) {
      if(neighbors[j] == rank) (*inneighbors)++;
    }
  }
}

int MPI_Virtual_neighbors(MPI_Comm comm, int inmaxneighbors, int *inneighbors, int outmaxneighbors, int *outneighbors) {

  int rank, p;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &p);

  MPI_Graph_neighbors(comm, rank, outmaxneighbors, outneighbors);
  
  // this is local
  // get all neighbors who point at me ... should have been cached
  // during creation - it's only a POC
  int inneighbor = 0;
  for(int i=0; i<p; i++) {
    int nneighbors;
    MPI_Graph_neighbors_count(comm, i, &nneighbors);
    int *neighbors = (int*)malloc(nneighbors*sizeof(int));
    MPI_Graph_neighbors(comm, i, nneighbors, neighbors);
    for(int j=0; j<nneighbors; j++) {
      if(neighbors[j] == rank) {
        inneighbors[inneighbor] = i;
        if(inmaxneighbors==++inneighbor) return MPI_SUCCESS;
      }
    }
  }
}

int main(int *argc, char **argv) {
  
  int rank, p;

  MPI_Init(NULL, NULL);

  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &p);

  int *nodes = (int*)malloc(2*sizeof(int)); // everybody contributes two nodes
  int *degrees = (int*)malloc(2*sizeof(int)); // everybody contributes two edges
  int *tgtnodes = (int*)malloc(2*sizeof(int));

  nodes[0] = rank;
  degrees[0] = 1;
  tgtnodes[0] = (rank+1)%p; // a ring
  nodes[1] = (rank+1)%p;
  degrees[1] = 1;
  tgtnodes[1] = rank; // a double ring
  MPI_Comm newcomm;


  MPI_Virtual_graph_create(MPI_COMM_WORLD, 2, nodes, degrees, tgtnodes, 
                           0, &newcomm);

  int inneighbors, outneighbors;
  MPI_Virtual_neighbors_count(newcomm, &inneighbors, &outneighbors);
  printf("rank: %i has %i inneighbors and %i outneighbors\n", rank, inneighbors, outneighbors);

  int *inneigh=(int*)malloc(inneighbors*sizeof(int));
  int *outneigh=(int*)malloc(outneighbors*sizeof(int));
  MPI_Virtual_neighbors(newcomm, inneighbors, inneigh, outneighbors, outneigh);
  if(!rank) {
    printf("rank %i has outneighbors:", rank);
    for(int i=0; i<outneighbors; i++) printf(" %i", outneigh[i]);
    printf(" and inneighbors:");
    for(int i=0; i<inneighbors; i++) printf(" %i", inneigh[i]);
    printf("\n");
  }
  

  MPI_Finalize();

}