/* -*- mode: C++ -*-  */
/* 
   IGraph library.
   Copyright (C) 2010-2012  Gabor Csardi <csardi.gabor@gmail.com>
   334 Harvard street, Cambridge, MA 02139 USA
   
   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.
   
   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.
   
   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 
   02110-1301 USA

*/

#include "igraph_interface.h"
#include "igraph_community.h"
#include "igraph_memory.h"
#include "igraph_constructors.h"
#include "igraph_attributes.h"
#include "igraph_foreign.h"
#include "igraph_hrg.h"
#include "igraph_random.h"

#include "hrg_dendro.h"
#include "hrg_graph.h"
#include "hrg_graph_simp.h"

using namespace fitHRG;

/**
 * \section hrg_intro Introduction
 * 
 * <para>A hierarchical random graph is an ensemble of undirected
 * graphs with \c n vertices. It is defined via a binary tree with \c
 * n leaf and \c n-1 internal vertices, where the
 * internal vertices are labeled with probabilities.
 * The probability that two vertices are connected in the random graph
 * is given by the probability label at their closest common
 * ancestor.
 * </para>
 * 
 * <para>Please read the following two articles for more about
 * hierarchical random graphs: A. Clauset, C. Moore, and M.E.J. Newman. 
 * Hierarchical structure and the prediction of missing links in networks.
 * Nature 453, 98 - 101 (2008); and A. Clauset, C. Moore, and M.E.J. Newman. 
 * Structural Inference of Hierarchies in Networks. In E. M. Airoldi
 * et al. (Eds.): ICML 2006 Ws, Lecture Notes in Computer Science
 * 4503, 1-13. Springer-Verlag, Berlin Heidelberg (2007).
 * </para>
 * 
 * <para>
 * igraph contains functions for fitting HRG models to a given network
 * (\ref igraph_hrg_fit), for generating networks from a given HRG
 * ensemble (\ref igraph_hrg_game, \ref igraph_hrg_sample), converting
 * an igraph graph to a HRG and back (\ref igraph_hrg_create, \ref
 * igraph_hrg_dendrogram), for calculating a consensus tree from a
 * set of sampled HRGs (\ref igraph_hrg_consensus) and for predicting
 * missing edges in a network based on its HRG models (\ref
 * igraph_hrg_predict). 
 * </para>
 * 
 * <para>The igraph HRG implementation is heavily based on the code
 * published by Aaron Clauset, at his website,
 * http://tuvalu.santafe.edu/~aaronc/hierarchy/
 * </para>
 */

namespace fitHRG {
  struct pblock { double L; int i; int j; };
}

int markovChainMonteCarlo(dendro *d, unsigned int period, 
			  igraph_hrg_t *hrg) {
  
  igraph_real_t bestL=d->getLikelihood();
  double  dL;
  bool    flag_taken;
  
  // Because moves in the dendrogram space are chosen (Monte
  // Carlo) so that we sample dendrograms with probability
  // proportional to their likelihood, a likelihood-proportional
  // sampling of the dendrogram models would be equivalent to a
  // uniform sampling of the walk itself. We would still have to
  // decide how often to sample the walk (at most once every n
  // steps is recommended) but for simplicity, the code here
  // simply runs the MCMC itself. To actually compute something 
  // over the set of sampled dendrogram models (in a Bayesian
  // model averaging sense), you'll need to code that yourself.
  
  // do 'period' MCMC moves before doing anything else
  for (unsigned int i=0; i<period; i++) {
    
    // make a MCMC move
    IGRAPH_CHECK(! d->monteCarloMove(dL, flag_taken, 1.0));
    
    // get likelihood of this D given G
    igraph_real_t cl= d->getLikelihood(); 
    if (cl > bestL) {
      // store the current best likelihood
      bestL = cl;
      // record the HRG structure
      d->recordDendrogramStructure(hrg);
    }
  }
  // corrects floating-point errors O(n)
  d->refreshLikelihood();	
  
  return 0;
}

int markovChainMonteCarlo2(dendro *d, int num_samples) {
  bool flag_taken;
  double dL, ptest = 1.0/(50.0*(double)(d->g->numNodes()));
  int sample_num=0, t=1, thresh = 200 * d->g->numNodes();
  
  // Since we're sampling uniformly at random over the equilibrium
  // walk, we just need to do a bunch of MCMC moves and let the
  // sampling happen on its own.
  while (sample_num < num_samples) {
    // Make a single MCMC move
    d->monteCarloMove(dL, flag_taken, 1.0);
		
    // We sample the dendrogram space once every n MCMC moves (on
    // average). Depending on the flags on the command line, we sample
    // different aspects of the dendrograph structure.
    if (t > thresh && RNG_UNIF01() < ptest) {
      sample_num++;
      d->sampleSplitLikelihoods(sample_num);
    }
		
    t++;

    // correct floating-point errors O(n)
    d->refreshLikelihood();	// TODO: less frequently    
  }
	
  return 0;
}

int MCMCEquilibrium_Find(dendro *d, igraph_hrg_t *hrg) {

  // We want to run the MCMC until we've found equilibrium; we
  // use the heuristic of the average log-likelihood (which is
  // exactly the entropy) over X steps being very close to the
  // average log-likelihood (entropy) over the X steps that
  // preceded those. In other words, we look for an apparent
  // local convergence of the entropy measure of the MCMC.

  bool flag_taken;
  igraph_real_t dL, Likeli;
  igraph_real_t oldMeanL;
  igraph_real_t newMeanL=-1e-49;
  
  while (1) {
    oldMeanL = newMeanL;
    newMeanL = 0.0;
    for (int i=0; i<65536; i++) {
      IGRAPH_CHECK(! d->monteCarloMove(dL, flag_taken, 1.0));
      Likeli = d->getLikelihood();
      newMeanL += Likeli;
    }
    // corrects floating-point errors O(n)
    d->refreshLikelihood();
    if (fabs(newMeanL-oldMeanL)/65536.0 < 1.0) { break; }	
  }
  
  // Record the result
  if (hrg) { d->recordDendrogramStructure(hrg); }

  return 0;
}

int igraph_i_hrg_getgraph(const igraph_t *igraph,
			  dendro *d) {
  
  int no_of_nodes = igraph_vcount(igraph);
  int no_of_edges = igraph_ecount(igraph);
  int i;
  
  // Create graph
  d->g=new graph(no_of_nodes);  

  // Add edges
  for (i=0; i<no_of_edges; i++) {
    int from=IGRAPH_FROM(igraph, i);
    int to=IGRAPH_TO(igraph, i);
    if (from==to) { continue; }
    if (!d->g->doesLinkExist(from, to)) { d->g->addLink(from, to); }
    if (!d->g->doesLinkExist(to, from)) { d->g->addLink(to, from); }
  }
  
  d->buildDendrogram();

  return 0;
}

int igraph_i_hrg_getsimplegraph(const igraph_t *igraph, 
				dendro *d, simpleGraph **sg, 
				int num_bins) {

  int no_of_nodes = igraph_vcount(igraph);
  int no_of_edges = igraph_ecount(igraph);
  int i;
  
  // Create graphs
  d->g = new graph(no_of_nodes, true);
  d->g->setAdjacencyHistograms(num_bins);
  (*sg) = new simpleGraph(no_of_nodes);
  
  for (i=0; i<no_of_edges; i++) {
    int from=IGRAPH_FROM(igraph, i);
    int to=IGRAPH_TO(igraph, i);
    if (from==to) { continue; }
    if (!d->g->doesLinkExist(from, to)) { d->g->addLink(from, to); }
    if (!d->g->doesLinkExist(to, from)) { d->g->addLink(to, from); }
    if (!(*sg)->doesLinkExist(from, to)) { (*sg)->addLink(from, to); }
    if (!(*sg)->doesLinkExist(to, from)) { (*sg)->addLink(to, from); }
  }

  d->buildDendrogram();
  
  return 0;
}

/**
 * \function igraph_hrg_init
 * Allocate memory for a HRG.
 * 
 * This function must be called before passing an \ref igraph_hrg_t to
 * an igraph function.
 * \param hrg Pointer to the HRG data structure to initialize.
 * \param n The number of vertices in the graph that is modeled by
 *    this HRG. It can be zero, if this is not yet known.
 * \return Error code.
 * 
 * Time complexity: O(n), the number of vertices in the graph.
 */

int igraph_hrg_init(igraph_hrg_t *hrg, int n) {
  IGRAPH_VECTOR_INIT_FINALLY(&hrg->left,      n-1);
  IGRAPH_VECTOR_INIT_FINALLY(&hrg->right,     n-1);
  IGRAPH_VECTOR_INIT_FINALLY(&hrg->prob,      n-1);  
  IGRAPH_VECTOR_INIT_FINALLY(&hrg->edges,     n-1);  
  IGRAPH_VECTOR_INIT_FINALLY(&hrg->vertices,  n-1);  
  IGRAPH_FINALLY_CLEAN(5);
  return 0;
}

/** 
 * \function igraph_hrg_destroy
 * Deallocate memory for an HRG.
 * 
 * The HRG data structure can be reinitialized again with an \ref
 * igraph_hrg_destroy call.
 * \param hrg Pointer to the HRG data structure to deallocate.
 * 
 * Time complexity: operating system dependent.
 */

void igraph_hrg_destroy(igraph_hrg_t *hrg) {
  igraph_vector_destroy(&hrg->left);
  igraph_vector_destroy(&hrg->right);
  igraph_vector_destroy(&hrg->prob);
  igraph_vector_destroy(&hrg->edges);
  igraph_vector_destroy(&hrg->vertices);
}

/**
 * \function igraph_hrg_size
 * Returns the size of the HRG, the number of leaf nodes.
 * 
 * \param hrg Pointer to the HRG.
 * \return The number of leaf nodes in the HRG.
 * 
 * Time complexity: O(1).
 */

int igraph_hrg_size(const igraph_hrg_t *hrg) {
  return igraph_vector_size(&hrg->left)+1;
}

/**
 * \function igraph_hrg_resize
 * Resize a HRG.
 * 
 * \param hrg Pointer to an initialized (see \ref igraph_hrg_init)
 *   HRG.
 * \param newsize The new size, i.e. the number of leaf nodes.
 * \return Error code.
 * 
 * Time complexity: O(n), n is the new size.
 */

int igraph_hrg_resize(igraph_hrg_t *hrg, int newsize) {
  int origsize=igraph_hrg_size(hrg);
  int ret=0;
  igraph_error_handler_t *oldhandler =
    igraph_set_error_handler(igraph_error_handler_ignore);

  ret  = igraph_vector_resize(&hrg->left, newsize-1);
  ret |= igraph_vector_resize(&hrg->right, newsize-1);
  ret |= igraph_vector_resize(&hrg->prob, newsize-1);
  ret |= igraph_vector_resize(&hrg->edges, newsize-1);
  ret |= igraph_vector_resize(&hrg->vertices, newsize-1);

  igraph_set_error_handler(oldhandler);

  if (ret) { 
    igraph_vector_resize(&hrg->left, origsize);
    igraph_vector_resize(&hrg->right, origsize);
    igraph_vector_resize(&hrg->prob, origsize);
    igraph_vector_resize(&hrg->edges, origsize);
    igraph_vector_resize(&hrg->vertices, origsize);
    IGRAPH_ERROR("Cannot resize HRG", ret);
  }  

  return 0;
}

/**
 * \function igraph_hrg_fit
 * Fit a hierarchical random graph model to a network
 * 
 * \param graph The igraph graph to fit the model to. Edge directions
 *   are ignored in directed graphs.
 * \param hrg Pointer to an initialized HRG, the result of the fitting
 *   is stored here. It can also be used to pass a HRG to the
 *   function, that can be used as the starting point of the Markov
 *   Chain Monte Carlo fitting, if the \c start argument is true.
 * \param start Logical, whether to start the fitting from the given
 *   HRG. 
 * \param steps Integer, the number of MCMC steps to take in the
 *   fitting procedure. If this is zero, then the fitting stop is a
 *   convergence criteria is fulfilled.
 * \return Error code.
 * 
 * Time complexity: TODO.
 */

int igraph_hrg_fit(const igraph_t *graph, 
		   igraph_hrg_t *hrg,
		   igraph_bool_t start,
		   int steps) {
  
  int no_of_nodes=igraph_vcount(graph);
  dendro *d;

  RNG_BEGIN();

  d = new dendro;  

  // If we want to start from HRG
  if (start) {
    d->clearDendrograph();
    if (igraph_hrg_size(hrg) != no_of_nodes) {
      delete d;
      IGRAPH_ERROR("Invalid HRG to start from", IGRAPH_EINVAL);
    }
    // Convert the igraph graph
    IGRAPH_CHECK(igraph_i_hrg_getgraph(graph, d));
    d->importDendrogramStructure(hrg);
  } else {
    // Convert the igraph graph
    IGRAPH_CHECK(igraph_i_hrg_getgraph(graph, d));
    IGRAPH_CHECK(igraph_hrg_resize(hrg, no_of_nodes));
  }

  // Run fixed number of steps, or until convergence
  if (steps > 0) {
    IGRAPH_CHECK(markovChainMonteCarlo(d, steps, hrg));
  } else {
    IGRAPH_CHECK(MCMCEquilibrium_Find(d, hrg));
  }
  
  delete d;

  RNG_END();

  return 0;

}

/**
 * \function igraph_hrg_sample
 * Sample from a hierarchical random graph model
 * 
 * Sample from a hierarchical random graph ensemble. The ensemble can
 * be given as a graph (\c input_graph), or as a HRG object (\c hrg).
 * If a graph is given, then first an MCMC optimization is performed
 * to find the optimal fitting model; then the MCMC is used to sample
 * the graph(s).
 * \param input_graph An igraph graph, or a null pointer. If not a
 *   null pointer, then a HRG is first fitted to the graph, possibly
 *   starting from the given HRG, if the \c start argument is true. If
 *   is is a null pointer, then the given HRG is used as a starting
 *   point, to  find the optimum of the Markov chain, before the
 *   sampling.
 * \param sample Pointer to an uninitialized graph, or a null
 *   pointer. If only one sample is requested, and it is not a null
 *   pointer, then the sample is stored here.
 * \param samples An initialized vector of pointers. If more than one
 *   samples are requested, then they are stored here. Note that to
 *   free this data structure, you need to call \ref igraph_destroy on
 *   each graph first, then \c free() on all pointers, and finally
 *   \ref igraph_vector_ptr_destroy.
 * \param no_samples The number of samples to generate.
 * \param hrg A HRG. It is modified during the sampling.
 * \param start Logical, whether to start the MCMC from the given
 *   HRG.
 * \return Error code.
 * 
 * Time complexity: TODO.
 */

int igraph_hrg_sample(const igraph_t *input_graph,
		      igraph_t *sample,
		      igraph_vector_ptr_t *samples,
		      int no_samples,
		      igraph_hrg_t *hrg,
		      igraph_bool_t start) {

  int i;
  dendro *d;

  if (no_samples < 0) {
    IGRAPH_ERROR("Number of samples must be non-negative", IGRAPH_EINVAL);
  }

  if (!sample && !samples) {
    IGRAPH_ERROR("Give at least one of `sample' and `samples'",
		 IGRAPH_EINVAL);
  }

  if (no_samples != 1 && sample) {
    IGRAPH_ERROR("Number of samples should be one if `sample' is given", 
		 IGRAPH_EINVAL);
  }

  if (no_samples > 1 && !samples) {
    IGRAPH_ERROR("`samples' must be non-null if number of samples "
		 "is larger than 1", IGRAPH_EINVAL);
  }

  if (!start && !input_graph) {
    IGRAPH_ERROR("Input graph must be given if initial HRG is not used", 
		 IGRAPH_EINVAL);
  }

  if (!start) {
    IGRAPH_CHECK(igraph_hrg_resize(hrg, igraph_vcount(input_graph)));
  }

  if (input_graph && igraph_hrg_size(hrg) != igraph_vcount(input_graph)) {
    IGRAPH_ERROR("Invalid HRG size, should match number of nodes", 
		 IGRAPH_EINVAL);
  }

  RNG_BEGIN();

  d = new dendro;

  // Need to find equilibrium first?
  if (start) {
    d->clearDendrograph();
    d->importDendrogramStructure(hrg);
  } else {
    IGRAPH_CHECK(MCMCEquilibrium_Find(d, hrg));
  }

  // TODO: free on error

  if (sample) {
    // A single graph
    d->makeRandomGraph();
    d->recordGraphStructure(sample);
    if (samples) {
      igraph_t *G=igraph_Calloc(1, igraph_t);
      if (!G) { IGRAPH_ERROR("Cannot sample HRG graphs", IGRAPH_ENOMEM); }
      d->recordGraphStructure(G);
      IGRAPH_CHECK(igraph_vector_ptr_resize(samples, 1));
      VECTOR(*samples)[0]=G;
    }
  } else {
    // Sample many     
    IGRAPH_CHECK(igraph_vector_ptr_resize(samples, no_samples));
    for (i=0; i<no_samples; i++) {
      igraph_t *G=igraph_Calloc(1, igraph_t);
      if (!G) { IGRAPH_ERROR("Cannot sample HRG graphs", IGRAPH_ENOMEM); }
      d->makeRandomGraph();
      d->recordGraphStructure(G);
      VECTOR(*samples)[i]=G;
    }
  }

  delete d;

  RNG_END();

  return 0;
}

/** 
 * \function igraph_hrg_game
 * Generate a hierarchical random graph
 * 
 * This function is a simple shortcut to \ref igraph_hrg_sample. 
 * It creates a single graph, from the given HRG.
 * \param graph Pointer to an uninitialized graph, the new graph is
 *   created here.
 * \param hrg The hierarchical random graph model to sample from. It
 *   is modified during the MCMC process.
 * \return Error code.
 * 
 * Time complexity: TODO.
 */

int igraph_hrg_game(igraph_t *graph,
		    const igraph_hrg_t *hrg) {
  return igraph_hrg_sample(/* input_graph= */ 0, /* sample= */ graph, 
			   /* samples= */ 0, /* no_samples=*/ 1,
			   /* hrg= */ (igraph_hrg_t*) hrg, 
			   /* start= */ 1);
}

/**
 * \function igraph_hrg_dendrogram
 * Create a dendrogram from a hierarchical random graph.
 * 
 * Creates the igraph graph equivalent of an \ref igraph_hrg_t data
 * structure. 
 * \param graph Pointer to an uninitialized graph, the result is
 *   stored here.
 * \param hrg The hierarchical random graph to convert.
 * \return Error code.
 * 
 * Time complexity: O(n), the number of vertices in the graph.
 */

int igraph_hrg_dendrogram(igraph_t *graph,
			  const igraph_hrg_t *hrg) {
  
  int orig_nodes=igraph_hrg_size(hrg);
  int no_of_nodes=orig_nodes * 2 - 1;
  int no_of_edges=no_of_nodes-1;
  igraph_vector_t edges;
  int i, idx=0;
  igraph_vector_ptr_t vattrs;
  igraph_vector_t prob;
  igraph_attribute_record_t rec = { "probability", 
				    IGRAPH_ATTRIBUTE_NUMERIC,
				    &prob };
  
  // Probability labels, for leaf nodes they are IGRAPH_NAN
  IGRAPH_VECTOR_INIT_FINALLY(&prob, no_of_nodes);
  for (i=0; i<orig_nodes; i++) {
    VECTOR(prob)[i] = IGRAPH_NAN;
  }
  for (i=0; i<orig_nodes-1; i++) {
    VECTOR(prob)[orig_nodes+i] = VECTOR(hrg->prob)[i];
  }

  IGRAPH_VECTOR_INIT_FINALLY(&edges, no_of_edges * 2);
  IGRAPH_CHECK(igraph_vector_ptr_init(&vattrs, 1));
  IGRAPH_FINALLY(igraph_vector_ptr_destroy, &vattrs);
  VECTOR(vattrs)[0] = &rec;

  for (i=0; i<orig_nodes-1; i++) {
    int left=VECTOR(hrg->left)[i];
    int right=VECTOR(hrg->right)[i];

    VECTOR(edges)[idx++] = orig_nodes+i;
    VECTOR(edges)[idx++] = left < 0 ? orig_nodes-left-1 : left;
    VECTOR(edges)[idx++] = orig_nodes+i;
    VECTOR(edges)[idx++] = right < 0 ? orig_nodes-right-1 : right;
  }

  IGRAPH_CHECK(igraph_empty(graph, 0, IGRAPH_DIRECTED));
  IGRAPH_FINALLY(igraph_destroy, graph);
  IGRAPH_CHECK(igraph_add_vertices(graph, no_of_nodes, &vattrs));
  IGRAPH_CHECK(igraph_add_edges(graph, &edges, 0));

  igraph_vector_ptr_destroy(&vattrs);
  igraph_vector_destroy(&edges);
  igraph_vector_destroy(&prob);
  IGRAPH_FINALLY_CLEAN(4);	// + 1 for graph
  
  return 0;
}

/** 
 * \function igraph_hrg_consensus
 * Calculate a consensus tree for a HRG.
 * 
 * The calculation can be started from the given HRG (\c hrg), or (if
 * \c start is false), a HRG is first fitted to the given graph.
 * 
 * \param graph The input graph.
 * \param parents An initialized vector, the results are stored
 *   here. For each vertex, the id of its parent vertex is stored, or
 *   -1, if the vertex is the root vertex in the tree. The first n
 *   vertex ids (from 0) refer to the original vertices of the graph,
 *   the other ids refer to vertex groups.
 * \param weights Numeric vector, counts the number of times a given
 *   tree split occured in the generated network samples, for each
 *   internal vertices. The order is the same as in \c parents.
 * \param hrg A hierarchical random graph. It is used as a starting
 *   point for the sampling, if the \c start argument is true. It is
 *   modified along the MCMC.
 * \param start Logical, whether to use the supplied HRG (in \c hrg)
 *   as a starting point for the MCMC.
 * \param num_samples The number of samples to generate for creating
 *   the consensus tree.
 * \return Error code.
 * 
 * Time complexity: TODO.
 */

int igraph_hrg_consensus(const igraph_t *graph,
			 igraph_vector_t *parents,
			 igraph_vector_t *weights,
			 igraph_hrg_t *hrg,
			 igraph_bool_t start, 
			 int num_samples) {

  dendro *d;

  if (start && !hrg) {
    IGRAPH_ERROR("`hrg' must be given is `start' is true", IGRAPH_EINVAL);
  }

  RNG_BEGIN();
  
  d = new dendro;
  
  if (start) {
    d->clearDendrograph();
    IGRAPH_CHECK(igraph_i_hrg_getgraph(graph, d));
    d->importDendrogramStructure(hrg);
  } else {
    IGRAPH_CHECK(igraph_i_hrg_getgraph(graph, d));
    if (hrg) { igraph_hrg_resize(hrg, igraph_vcount(graph)); }
    IGRAPH_CHECK(MCMCEquilibrium_Find(d, hrg));
  }

  IGRAPH_CHECK(markovChainMonteCarlo2(d, num_samples));

  d->recordConsensusTree(parents, weights);

  delete d;
  
  RNG_END();

  return 0;
}

int MCMCEquilibrium_Sample(dendro *d, int num_samples) {

  // Because moves in the dendrogram space are chosen (Monte
  // Carlo) so that we sample dendrograms with probability
  // proportional to their likelihood, a likelihood-proportional
  // sampling of the dendrogram models would be equivalent to a
  // uniform sampling of the walk itself. We would still have to
  // decide how often to sample the walk (at most once every n steps
  // is recommended) but for simplicity, the code here simply runs the
  // MCMC itself. To actually compute something over the set of
  // sampled dendrogram models (in a Bayesian model averaging sense),
  // you'll need to code that yourself.
  
  double dL;
  bool flag_taken;
  int sample_num=0;
  int t=1, thresh=100 * d->g->numNodes();
  double ptest=1.0/10.0/d->g->numNodes();

  while (sample_num < num_samples) {
    d->monteCarloMove(dL, flag_taken, 1.0);
    if (t > thresh && RNG_UNIF01() < ptest) {
      sample_num++;
      d->sampleAdjacencyLikelihoods();
    }
    d->refreshLikelihood();	// TODO: less frequently
    t++;
  }
  
  return 0;
}

int QsortPartition (pblock* array, int left, int right, int index) {
  pblock p_value, temp;
  p_value.L = array[index].L;
  p_value.i = array[index].i;
  p_value.j = array[index].j;
  
  // swap(array[p_value], array[right])
  temp.L = array[right].L;
  temp.i = array[right].i;
  temp.j = array[right].j;
  array[right].L = array[index].L;
  array[right].i = array[index].i;
  array[right].j = array[index].j;
  array[index].L = temp.L;
  array[index].i = temp.i;
  array[index].j = temp.j;
	
  int stored = left;
  for (int i=left; i<right; i++) {
    if (array[i].L <= p_value.L) {
      // swap(array[stored], array[i])
      temp.L = array[i].L;
      temp.i = array[i].i;
      temp.j = array[i].j;
      array[i].L = array[stored].L;
      array[i].i = array[stored].i;
      array[i].j = array[stored].j;
      array[stored].L = temp.L;
      array[stored].i = temp.i;
      array[stored].j = temp.j;
      stored++;
    }
  }
  // swap(array[right], array[stored])
  temp.L = array[stored].L;
  temp.i = array[stored].i;
  temp.j = array[stored].j;
  array[stored].L = array[right].L;
  array[stored].i = array[right].i;
  array[stored].j = array[right].j;
  array[right].L  = temp.L;
  array[right].i  = temp.i;
  array[right].j  = temp.j;
  
  return stored;
}

void QsortMain (pblock* array, int left, int right) {
  if (right > left) {
    int pivot = left;
    int part  = QsortPartition(array, left, right, pivot);
    QsortMain(array, left,   part-1);
    QsortMain(array, part+1, right  );
  }
  return;
}

int rankCandidatesByProbability(simpleGraph *sg, dendro *d, 
				pblock *br_list, int mk) {
  int mkk=0;
  int n=sg->getNumNodes();
  for (int i=0; i<n; i++) {
    for (int j=i+1; j<n; j++) {
      if (sg->getAdjacency(i, j) < 0.5) {
	double temp=d->g->getAdjacencyAverage(i, j);
	br_list[mkk].L = temp * (1.0 + RNG_UNIF01()/1000.0);
	br_list[mkk].i = i;
	br_list[mkk].j = j;
	mkk++;
      }
    }
  }
  
  // Sort the candidates by their average probability
  QsortMain(br_list, 0, mk-1);  

  return 0;
}

int recordPredictions(pblock *br_list, igraph_vector_t *edges, 
		      igraph_vector_t *prob, int mk) {
    
  IGRAPH_CHECK(igraph_vector_resize(edges, mk*2));
  IGRAPH_CHECK(igraph_vector_resize(prob, mk));
  
  for (int i=mk-1, idx=0, idx2=0; i>=0; i--) {
    VECTOR(*edges)[idx++] = br_list[i].i;
    VECTOR(*edges)[idx++] = br_list[i].j;
    VECTOR(*prob)[idx2++] = br_list[i].L;
  }

  return 0;
}

/**
 * \function igraph_hrg_predict
 * Predict missing edges in a graph, based on HRG models
 * 
 * Samples HRG models for a network, and estimated the probability
 * that an edge was falsely observed as non-existent in the network.
 * \param graph The input graph.
 * \param edges The list of missing edges is stored here, the first
 *   two elements are the first edge, the next two the second edge,
 *   etc. 
 * \param prob Vector of probabilies for the existence of missing
 *   edges, in the order corresponding to \c edges.
 * \param hrg A HRG, it is used as a starting point if \c start is
 *   true. It is also modified during the MCMC sampling.
 * \param start Logical, whether to start the MCMC from the given HRG.
 * \param num_samples The number of samples to generate.
 * \param num_bins Controls the resolution of the edge
 *   probabilities. Higher numbers result higher resolution. 
 * \return Error code.
 * 
 * Time complexity: TODO.
 */

int igraph_hrg_predict(const igraph_t *graph,
		       igraph_vector_t *edges,
		       igraph_vector_t *prob,
		       igraph_hrg_t *hrg,
		       igraph_bool_t start, 
		       int num_samples, 
		       int num_bins) {

  dendro *d;
  pblock *br_list;
  int mk;
  simpleGraph *sg;

  if (start && !hrg) {
    IGRAPH_ERROR("`hrg' must be given is `start' is true", IGRAPH_EINVAL);
  }

  RNG_BEGIN();

  d = new dendro;

  IGRAPH_CHECK(igraph_i_hrg_getsimplegraph(graph, d, &sg, num_bins));

  mk = sg->getNumNodes() * (sg->getNumNodes()-1) / 2 - sg->getNumLinks()/2;
  br_list = new pblock[mk];
  for (int i=0; i<mk; i++) { 
    br_list[i].L = 0.0; 
    br_list[i].i = -1; 
    br_list[i].j = -1; 
  }
  
  if (start) {
    d->clearDendrograph();
    // this has cleared the graph as well.... bug?
    IGRAPH_CHECK(igraph_i_hrg_getsimplegraph(graph, d, &sg, num_bins));
    d->importDendrogramStructure(hrg);
  } else {
    if (hrg) { igraph_hrg_resize(hrg, igraph_vcount(graph)); }
    IGRAPH_CHECK(MCMCEquilibrium_Find(d, hrg));
  }

  IGRAPH_CHECK(MCMCEquilibrium_Sample(d, num_samples));
  IGRAPH_CHECK(rankCandidatesByProbability(sg, d, br_list, mk));
  IGRAPH_CHECK(recordPredictions(br_list, edges, prob, mk));

  delete d;
  delete sg;
  delete [] br_list;

  RNG_END();

  return 0;
}

/**
 * \function igraph_hrg_create
 * Create a HRG from an igraph graph.
 * 
 * \param hrg Pointer to an initialized \ref igraph_hrg_t. The result
 *    is stored here.
 * \param graph The igraph graph to convert. It must be a directed
 *    binary tree, with n-1 internal and n leaf vertices. The root
 *    vertex must have in-degree zero.
 * \param prob The vector of probabilities, this is used to label the
 *    internal nodes of the hierarchical random graph. The values
 *    corresponding to the leaves are ignored.
 * \return Error code.
 * 
 * Time complexity: O(n), the number of vertices in the tree.
 */

int igraph_hrg_create(igraph_hrg_t *hrg,
		      const igraph_t *graph, 
		      const igraph_vector_t *prob) {

  int no_of_nodes=igraph_vcount(graph);
  int no_of_internal=(no_of_nodes-1)/2;
  igraph_vector_t deg, idx;
  int root=0;
  int d0=0, d1=0, d2=0;
  int ii=0, il=0;
  igraph_vector_t neis;
  igraph_vector_t path;

  // --------------------------------------------------------
  // CHECKS
  // --------------------------------------------------------

  // At least three vertices are required
  if (no_of_nodes < 3) {
    IGRAPH_ERROR("HRG tree must have at least three vertices",
		 IGRAPH_EINVAL);
  }

  // Prob vector was given
  if (!prob) {
    IGRAPH_ERROR("Probability vector must be given for HRG", 
		 IGRAPH_EINVAL);
  }

  // Length of prob vector
  if (igraph_vector_size(prob) != no_of_nodes) {
    IGRAPH_ERROR("HRG probability vector of wrong size", IGRAPH_EINVAL);
  }

  // Must be a directed graph
  if (!igraph_is_directed(graph)) {
    IGRAPH_ERROR("HRG graph must be directed", IGRAPH_EINVAL);
  }

  // Number of nodes must be odd
  if (no_of_nodes % 2 == 0) {
    IGRAPH_ERROR("Complete HRG graph must have odd number of vertices", 
		 IGRAPH_EINVAL);
  }

  IGRAPH_VECTOR_INIT_FINALLY(&deg, 0);

  // Every vertex, except for the root must have in-degree one.
  IGRAPH_CHECK(igraph_degree(graph, &deg, igraph_vss_all(), IGRAPH_IN,
			     IGRAPH_LOOPS));
  for (int i=0; i<no_of_nodes; i++) {
    int d=VECTOR(deg)[i];
    switch (d) {
    case 0: d0++; root=i; break;
    case 1: d1++; break;
    default:
      IGRAPH_ERROR("HRG nodes must have in-degree one, except for the "
		   "root vertex", IGRAPH_EINVAL);
    }
  }
  if (d1 != no_of_nodes-1 || d0 != 1) {
    IGRAPH_ERROR("HRG nodes must have in-degree one, except for the "
		 "root vertex", IGRAPH_EINVAL);
  }
  
  // Every internal vertex must have out-degree two,
  // leaves out-degree zero
  d0=d1=d2=0;
  IGRAPH_CHECK(igraph_degree(graph, &deg, igraph_vss_all(), IGRAPH_OUT, 
			     IGRAPH_LOOPS));
  for (int i=0; i<no_of_nodes; i++) {
    int d=VECTOR(deg)[i];
    switch (d) {
    case 0: d0++; break;
    case 2: d2++; break;
    default:
      IGRAPH_ERROR("HRG nodes must have out-degree 2 (internal nodes) or "
		   "degree 0 (leaves)", IGRAPH_EINVAL);
    }            
  }
  
  // Number of internal and external nodes is correct
  // This basically checks that the graph has one component
  if (d0 != d2+1) {
    IGRAPH_ERROR("HRG degrees are incorrect, maybe multiple components?",
		 IGRAPH_EINVAL);
  }
  
  // --------------------------------------------------------
  // Graph is good, do the conversion
  // --------------------------------------------------------

  // Create an index, that maps the root node as first, then
  // the internal nodes, then the leaf nodes
  IGRAPH_VECTOR_INIT_FINALLY(&idx, no_of_nodes);
  VECTOR(idx)[root] = - (ii++) - 1;
  for (int i=0; i<no_of_nodes; i++) {
    int d=VECTOR(deg)[i];
    if (i==root) { continue; }
    if (d==2) { VECTOR(idx)[i] = - (ii++) - 1; }
    if (d==0) { VECTOR(idx)[i] = (il++); }
  }

  igraph_hrg_resize(hrg, no_of_internal+1);
  IGRAPH_VECTOR_INIT_FINALLY(&neis, 0);
  for (int i=0; i<no_of_nodes; i++) {
    int ri=VECTOR(idx)[i];
    if (ri >= 0) { continue; }
    IGRAPH_CHECK(igraph_neighbors(graph, &neis, i, IGRAPH_OUT));
    VECTOR(hrg->left )[-ri-1] = VECTOR(idx)[ (int) VECTOR(neis)[0] ];
    VECTOR(hrg->right)[-ri-1] = VECTOR(idx)[ (int) VECTOR(neis)[1] ];
    VECTOR(hrg->prob )[-ri-1] = VECTOR(*prob)[i];
  }

  // Calculate the number of vertices and edges in each subtree
  igraph_vector_null(&hrg->edges);
  igraph_vector_null(&hrg->vertices);
  IGRAPH_VECTOR_INIT_FINALLY(&path, 0);
  IGRAPH_CHECK(igraph_vector_push_back(&path, VECTOR(idx)[root]));
  while (!igraph_vector_empty(&path)) {
    int ri=igraph_vector_tail(&path);
    int lc=VECTOR(hrg->left)[-ri-1];
    int rc=VECTOR(hrg->right)[-ri-1];
    if (lc < 0 && VECTOR(hrg->vertices)[-lc-1]==0) {
      // Go left
      IGRAPH_CHECK(igraph_vector_push_back(&path, lc));
    } else if (rc < 0 && VECTOR(hrg->vertices)[-rc-1]==0) {
      // Go right
      IGRAPH_CHECK(igraph_vector_push_back(&path, rc));
    } else {
      // Subtrees are done, update node and go up
      VECTOR(hrg->vertices)[-ri-1] += 
	lc < 0 ? VECTOR(hrg->vertices)[-lc-1] : 1;
      VECTOR(hrg->vertices)[-ri-1] += 
	rc < 0 ? VECTOR(hrg->vertices)[-rc-1] : 1;
      VECTOR(hrg->edges)[-ri-1] += lc < 0 ? VECTOR(hrg->edges)[-lc-1]+1 : 1;
      VECTOR(hrg->edges)[-ri-1] += rc < 0 ? VECTOR(hrg->edges)[-rc-1]+1 : 1;
      igraph_vector_pop_back(&path);
    }
  }

  igraph_vector_destroy(&path);
  igraph_vector_destroy(&neis);
  igraph_vector_destroy(&idx);
  igraph_vector_destroy(&deg);
  IGRAPH_FINALLY_CLEAN(4);

  return 0;
}