/* -*- mode: C -*-  */
/* vim:set ts=8 sw=2 sts=2 et: */
/* 
   IGraph R library.
   Copyright (C) 2003-2013  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_vector.h"
#include "igraph_matrix.h"
#include "igraph_random.h"
#include "igraph_constructors.h"
#include "igraph_games.h"

#include <float.h>		/* for DBL_EPSILON */
#include <math.h> 		/* for sqrt */

/**
 * \function igraph_sbm_game
 * Sample from a stochastic block model
 *
 * This function samples graphs from a stochastic block
 * model by (doing the equivalent of) Bernoulli
 * trials for each potential edge with the probabilities
 * given by the Bernoulli rate matrix, \p pref_matrix.
 * See Faust, K., &amp; Wasserman, S. (1992a). Blockmodels:
 * Interpretation and evaluation. Social Networks, 14, 5-–61.
 *
 * </para><para>
 * The order of the vertex ids in the generated graph corresponds to
 * the \p block_sizes argument.
 *
 * \param graph The output graph.
 * \param n Number of vertices.
 * \param pref_matrix The matrix giving the Bernoulli rates.
 *     This is a KxK matrix, where K is the number of groups.
 *     The probability of creating an edge between vertices from
 *     groups i and j is given by element (i,j).
 * \param block_sizes An integer vector giving the number of
 *     vertices in each group.
 * \param directed Boolean, whether to create a directed graph. If
 *     this argument is false, then \p pref_matrix must be symmetric.
 * \param loops Boolean, whether to create self-loops.
 * \return Error code.
 *
 * Time complexity: O(|V|+|E|+K^2), where |V| is the number of
 * vertices, |E| is the number of edges, and K is the number of
 * groups.
 *
 * \sa \ref igraph_erdos_renyi_game() for a simple Bernoulli graph.
 *
 */

int igraph_sbm_game(igraph_t *graph, igraph_integer_t n, 
		    const igraph_matrix_t *pref_matrix,
		    const igraph_vector_int_t *block_sizes,
		    igraph_bool_t directed, igraph_bool_t loops) {

  int no_blocks=igraph_matrix_nrow(pref_matrix);
  int from, to, fromoff=0;
  igraph_real_t minp, maxp;
  igraph_vector_t edges;
  
  /* ------------------------------------------------------------ */
  /* Check arguments                                              */
  /* ------------------------------------------------------------ */

  if (igraph_matrix_ncol(pref_matrix) != no_blocks) {
    IGRAPH_ERROR("Preference matrix is not square", 
		 IGRAPH_NONSQUARE);
  }

  igraph_matrix_minmax(pref_matrix, &minp, &maxp);
  if (minp < 0 || maxp > 1) { 
    IGRAPH_ERROR("Connection probabilities must in [0,1]", IGRAPH_EINVAL);
  }

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

  if (!directed && !igraph_matrix_is_symmetric(pref_matrix)) {
    IGRAPH_ERROR("Preference matrix must be symmetric for undirected graphs",
		 IGRAPH_EINVAL);
  }

  if (igraph_vector_int_size(block_sizes) != no_blocks) {
    IGRAPH_ERROR("Invalid block size vector length", IGRAPH_EINVAL);
  }

  if (igraph_vector_int_min(block_sizes) < 0) {
    IGRAPH_ERROR("Block size must be non-negative", IGRAPH_EINVAL);
  }

  if (igraph_vector_int_sum(block_sizes) != n) {
    IGRAPH_ERROR("Block sizes must sum up to number of vertices", 
		 IGRAPH_EINVAL);
  }

  IGRAPH_VECTOR_INIT_FINALLY(&edges, 0);  
  
  RNG_BEGIN();

  for (from = 0; from < no_blocks; from++) {
    double fromsize = VECTOR(*block_sizes)[from];
    int start = directed ? 0 : from;
    int i, tooff=0;
    for (i=0; i<start; i++) {
      tooff += VECTOR(*block_sizes)[i];
    }
    for (to = start; to < no_blocks; to++) {
      double tosize = VECTOR(*block_sizes)[to];
      igraph_real_t prob=MATRIX(*pref_matrix, from, to);
      double maxedges, last=RNG_GEOM(prob);
      if (directed && loops) {
	maxedges = fromsize * tosize;
	while (last < maxedges) {
	  int vto=floor(last/fromsize);
	  int vfrom=last - (igraph_real_t)vto * fromsize;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}
      } else if (directed && !loops && from!=to) {
	maxedges = fromsize * tosize;
	while (last < maxedges) {
	  int vto=floor(last/fromsize);
	  int vfrom=last - (igraph_real_t)vto * fromsize;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}	
      } else if (directed && !loops && from==to) {
	maxedges = fromsize * (fromsize-1);
	while (last < maxedges) {
	  int vto=floor(last/fromsize);
	  int vfrom=last - (igraph_real_t)vto * fromsize;
	  if (vfrom == vto) { vto=fromsize-1; }
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}
      } else if (!directed && loops && from!=to) {
	maxedges = fromsize * tosize;	
	while (last < maxedges) {
	  int vto=floor(last/fromsize);
	  int vfrom=last - (igraph_real_t)vto * fromsize;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}
      } else if (!directed && loops && from==to) {
	maxedges = fromsize * (fromsize+1) / 2.0;
	while (last < maxedges) {
	  long int vto=floor((sqrt(8*last+1)-1)/2);
	  long int vfrom=last-(((igraph_real_t)vto)*(vto+1))/2;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}	
      } else if (!directed && !loops && from!=to) {
	maxedges = fromsize * tosize;
	while (last < maxedges) {
	  int vto=floor(last/fromsize);
	  int vfrom=last - (igraph_real_t)vto * fromsize;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}
      } else /*!directed && !loops && from==to */ {
	maxedges = fromsize * (fromsize-1) / 2.0;
	while (last < maxedges) {
	  int vto=floor((sqrt(8*last+1)+1)/2);
	  int vfrom=last-(((igraph_real_t)vto)*(vto-1))/2;
	  igraph_vector_push_back(&edges, fromoff + vfrom);
	  igraph_vector_push_back(&edges, tooff + vto);
	  last += RNG_GEOM(prob);
	  last += 1;
	}
      }
      
      tooff += tosize;
    }
    fromoff += fromsize;
  }

  RNG_END();

  igraph_create(graph, &edges, n, directed);
  igraph_vector_destroy(&edges);
  IGRAPH_FINALLY_CLEAN(1);

  return 0;
}

/**
 * \function igraph_hsbm_game
 * Hierarchical stochastic block model
 *
 * The function generates a random graph according to the hierarchical
 * stochastic block model.
 *
 * \param graph The generated graph is stored here.
 * \param n The number of vertices in the graph.
 * \param m The number of vertices per block. n/m must be integer.
 * \param rho The fraction of vertices per cluster,
 *        within a block. Must sum up to 1, and rho * m must be integer
 *        for all elements of rho.
 * \param C A square, symmetric numeric matrix, the Bernoulli rates for
 *        the clusters within a block. Its size must mach the size of the
 *        \code{rho} vector.
 * \param p The Bernoulli rate of connections between
 *        vertices in different blocks.
 * \return Error code.
 *
 * \sa \ref igraph_sbm_game() for the classic stochastic block model,
 * \ref igraph_hsbm_list_game() for a more general version.
 */

int igraph_hsbm_game(igraph_t *graph, igraph_integer_t n, 
		     igraph_integer_t m, const igraph_vector_t *rho,
		     const igraph_matrix_t *C, igraph_real_t p) {
  
  int b, i, k=igraph_vector_size(rho);
  igraph_vector_t csizes;
  igraph_real_t sq_dbl_epsilon=sqrt(DBL_EPSILON);
  int no_blocks=n / m;
  igraph_vector_t edges;
  int offset=0;

  if (n < 1) { 
    IGRAPH_ERROR("`n' must be positive for HSBM", IGRAPH_EINVAL); 
  }
  if (m < 1) {
    IGRAPH_ERROR("`m' must be positive for HSBM", IGRAPH_EINVAL);
  }
  if ((long) n  % (long) m) {
    IGRAPH_ERROR("`n' must be a multiple of `m' for HSBM", IGRAPH_EINVAL);
  }
  if (!igraph_vector_isininterval(rho, 0, 1)) {
    IGRAPH_ERROR("`rho' must be between zero and one for HSBM", 
		 IGRAPH_EINVAL);
  }
  if (igraph_matrix_min(C) < 0 || igraph_matrix_max(C) > 1) {
    IGRAPH_ERROR("`C' must be between zero and one for HSBM", IGRAPH_EINVAL);
  }
  if (fabs(igraph_vector_sum(rho) - 1.0) > sq_dbl_epsilon) {
    IGRAPH_ERROR("`rho' must sum up to 1 for HSBM", IGRAPH_EINVAL);
  }
  if (igraph_matrix_nrow(C) != k || igraph_matrix_ncol(C) != k) {
    IGRAPH_ERROR("`C' dimensions must match `rho' dimensions in HSBM",
		 IGRAPH_EINVAL);
  }
  if (!igraph_matrix_is_symmetric(C)) {
    IGRAPH_ERROR("`C' must be a symmetric matrix", IGRAPH_EINVAL);
  }
  if (p < 0 || p > 1) {
    IGRAPH_ERROR("`p' must be a probability for HSBM", IGRAPH_EINVAL);
  }
  for (i=0; i<k; i++) {
    igraph_real_t s=VECTOR(*rho)[i] * m;
    if (fabs(round(s)-s) > sq_dbl_epsilon) {
      IGRAPH_ERROR("`rho' * `m' is not integer in HSBM", IGRAPH_EINVAL);
    }
  }
  
  IGRAPH_VECTOR_INIT_FINALLY(&csizes, k);
  for (i=0; i<k; i++) {
    VECTOR(csizes)[i] = round(VECTOR(*rho)[i] * m);
  }
  
  IGRAPH_VECTOR_INIT_FINALLY(&edges, 0);

  RNG_BEGIN();

  /* Block models first */

  for (b=0; b<no_blocks; b++) {
    int from, to, fromoff=0;
    
    for (from = 0; from < k; from++) {
      int fromsize = VECTOR(csizes)[from];
      int i, tooff=0;
      for (i=0; i<from; i++) {
	tooff += VECTOR(csizes)[i];
      }
      for (to = from; to < k; to++) {
	int tosize = VECTOR(csizes)[to];
	igraph_real_t prob=MATRIX(*C, from, to);
	igraph_real_t maxedges;
	igraph_real_t last=RNG_GEOM(prob);
	if (from != to) {
	  maxedges = fromsize * tosize;	
	  while (last < maxedges) {
	    int vto=floor(last/fromsize);
	    int vfrom=last - (igraph_real_t)vto * fromsize;
	    igraph_vector_push_back(&edges, offset + fromoff + vfrom);
	    igraph_vector_push_back(&edges, offset + tooff + vto);
	    last += RNG_GEOM(prob);
	    last += 1;
	  }
	} else /* from==to */ {
	  maxedges = fromsize * (fromsize-1) / 2.0;
	  while (last < maxedges) {
	    int vto=floor((sqrt(8*last+1)+1)/2);
	    int vfrom=last-(((igraph_real_t)vto)*(vto-1))/2;
	    igraph_vector_push_back(&edges, offset + fromoff + vfrom);
	    igraph_vector_push_back(&edges, offset + tooff + vto);
	    last += RNG_GEOM(prob);
	    last += 1;
	  }
	}

	tooff += tosize;
      }
      fromoff += fromsize;
    }
    
    offset += m;
  }

  /* And now the rest, if not a special case */

  if (p == 1) {
    int fromoff=0, tooff=m;
    for (b=0; b<no_blocks; b++) {
      igraph_real_t fromsize = m;
      igraph_real_t tosize = n - tooff;
      int from, to;
      for (from=0; from<fromsize; from++) {
	for (to=0; to<tosize; to++) {
	  igraph_vector_push_back(&edges, fromoff + from);
	  igraph_vector_push_back(&edges, tooff + to);
	}
      }
      fromoff += m;
      tooff += m;
    }
  } else if (p > 0) {
    int fromoff=0, tooff=m;
    for (b=0; b<no_blocks; b++) {
      igraph_real_t fromsize = m;
      igraph_real_t tosize = n - tooff;
      igraph_real_t maxedges= fromsize * tosize;
      igraph_real_t last=RNG_GEOM(p);
      while (last < maxedges) {
	int vto = floor(last/fromsize);
	int vfrom = last - (igraph_real_t) vto * fromsize;
	igraph_vector_push_back(&edges, fromoff + vfrom);
	igraph_vector_push_back(&edges, tooff + vto);
	last += RNG_GEOM(p);
	last += 1;
      }
      
      fromoff += m;
      tooff += m;
    }
  }

  RNG_END();

  igraph_create(graph, &edges, n, /*directed=*/ 0);

  igraph_vector_destroy(&edges);
  igraph_vector_destroy(&csizes);
  IGRAPH_FINALLY_CLEAN(2);

  return 0;
}

/**
 * \function igraph_hsbm_list_game
 * Hierarchical stochastic block model, more general version
 *
 * The function generates a random graph according to the hierarchical
 * stochastic block model.
 *
 * \param graph The generated graph is stored here.
 * \param n The number of vertices in the graph.
 * \param mlist An integer vector of block sizes.
 * \param rholist A list of rho vectors (\c igraph_vector_t objects), one
 *        for each block.
 * \param Clist A list of square matrices (\c igraph_matrix_t objects),
 *        one for each block, giving the Bernoulli rates of connections
 *        within the block.
 * \param p The Bernoulli rate of connections between
 *        vertices in different blocks.
 * \return Error code.
 *
 * \sa \ref igraph_sbm_game() for the classic stochastic block model,
 * \ref igraph_hsbm_game() for a simpler general version.
 */

int igraph_hsbm_list_game(igraph_t *graph, igraph_integer_t n,
			  const igraph_vector_int_t *mlist,
			  const igraph_vector_ptr_t *rholist,
			  const igraph_vector_ptr_t *Clist,
			  igraph_real_t p) {

  int i, no_blocks=igraph_vector_ptr_size(rholist);
  igraph_real_t sq_dbl_epsilon=sqrt(DBL_EPSILON);
  igraph_vector_t csizes, edges;
  int b, offset=0;

  if (n < 1) { 
    IGRAPH_ERROR("`n' must be positive for HSBM", IGRAPH_EINVAL); 
  }
  if (no_blocks == 0) {
    IGRAPH_ERROR("`rholist' empty for HSBM", IGRAPH_EINVAL);
  }
  if (igraph_vector_ptr_size(Clist) != no_blocks && 
      igraph_vector_int_size(mlist) != no_blocks) {
    IGRAPH_ERROR("`rholist' must have same length as `Clist' and `m' "
		 "for HSBM", IGRAPH_EINVAL);
  }
  if (p < 0 || p > 1) {
    IGRAPH_ERROR("`p' must be a probability for HSBM", IGRAPH_EINVAL);
  }
  /* Checks for m's */
  if (igraph_vector_int_sum(mlist) != n) {
    IGRAPH_ERROR("`m' must sum up to `n' for HSBM", IGRAPH_EINVAL);
  }
  if (igraph_vector_int_min(mlist) < 1) { 
    IGRAPH_ERROR("`m' must be positive for HSBM", IGRAPH_EINVAL);
  }
  /* Checks for the rhos */
  for (i=0; i<no_blocks; i++) {
    const igraph_vector_t *rho=VECTOR(*rholist)[i];
    if (!igraph_vector_isininterval(rho, 0, 1)) {
      IGRAPH_ERROR("`rho' must be between zero and one for HSBM", 
		   IGRAPH_EINVAL);
    }
    if (fabs(igraph_vector_sum(rho) - 1.0) > sq_dbl_epsilon) {
      IGRAPH_ERROR("`rho' must sum up to 1 for HSBM", IGRAPH_EINVAL);
    }
  }
  /* Checks for the Cs */
  for (i=0; i<no_blocks; i++) {
    const igraph_matrix_t *C=VECTOR(*Clist)[i];
    if (igraph_matrix_min(C) < 0 || igraph_matrix_max(C) > 1) {
      IGRAPH_ERROR("`C' must be between zero and one for HSBM", 
		   IGRAPH_EINVAL);
    }
    if (!igraph_matrix_is_symmetric(C)) {
      IGRAPH_ERROR("`C' must be a symmetric matrix", IGRAPH_EINVAL);
    }
  }
  /* Check that C and rho sizes match */
  for (i=0; i<no_blocks; i++) {
    const igraph_vector_t *rho=VECTOR(*rholist)[i];
    const igraph_matrix_t *C=VECTOR(*Clist)[i];
    int k=igraph_vector_size(rho);
    if (igraph_matrix_nrow(C) != k || igraph_matrix_ncol(C) != k) {
      IGRAPH_ERROR("`C' dimensions must match `rho' dimensions in HSBM",
		   IGRAPH_EINVAL);
    }
  }
  /* Check that rho * m is integer */
  for (i=0; i<no_blocks; i++) {
    const igraph_vector_t *rho=VECTOR(*rholist)[i];
    igraph_real_t m=VECTOR(*mlist)[i];
    int j, k=igraph_vector_size(rho);
    for (j=0; j<k; j++) {
      igraph_real_t s=VECTOR(*rho)[j] * m;
      if (fabs(round(s)-s) > sq_dbl_epsilon) {
	IGRAPH_ERROR("`rho' * `m' is not integer in HSBM", IGRAPH_EINVAL);
      }
    }
  }

  IGRAPH_VECTOR_INIT_FINALLY(&csizes, 0);
  IGRAPH_VECTOR_INIT_FINALLY(&edges, 0);
  
  RNG_BEGIN();

  /* Block models first */
  
  for (b=0; b<no_blocks; b++) {
    int from, to, fromoff=0;
    const igraph_vector_t *rho=VECTOR(*rholist)[b];
    const igraph_matrix_t *C=VECTOR(*Clist)[b];
    igraph_real_t m=VECTOR(*mlist)[b];
    int k=igraph_vector_size(rho);
    
    igraph_vector_resize(&csizes, k);
    for (i=0; i<k; i++) {
      VECTOR(csizes)[i] = round(VECTOR(*rho)[i] * m);
    }
    
    for (from=0; from < k; from++) {
      int fromsize=VECTOR(csizes)[from];
      int i, tooff=0;
      for (i=0; i<from; i++) {
	tooff += VECTOR(csizes)[i];
      }
      for (to = from; to < k; to++) {
	int tosize = VECTOR(csizes)[to];
	igraph_real_t prob=MATRIX(*C, from, to);
	igraph_real_t maxedges;
	igraph_real_t last=RNG_GEOM(prob);
	if (from != to) {
	  maxedges = fromsize * tosize;	
	  while (last < maxedges) {
	    int vto=floor(last/fromsize);
	    int vfrom=last - (igraph_real_t)vto * fromsize;
	    igraph_vector_push_back(&edges, offset + fromoff + vfrom);
	    igraph_vector_push_back(&edges, offset + tooff + vto);
	    last += RNG_GEOM(prob);
	    last += 1;
	  }
	} else /* from==to */ {
	  maxedges = fromsize * (fromsize-1) / 2.0;
	  while (last < maxedges) {
	    int vto=floor((sqrt(8*last+1)+1)/2);
	    int vfrom=last-(((igraph_real_t)vto)*(vto-1))/2;
	    igraph_vector_push_back(&edges, offset + fromoff + vfrom);
	    igraph_vector_push_back(&edges, offset + tooff + vto);
	    last += RNG_GEOM(prob);
	    last += 1;
	  }
	}

	tooff += tosize;
      }
      fromoff += fromsize;
    }
    
    offset += m;
  }

  /* And now the rest, if not a special case */

  if (p == 1) {    
    int fromoff=0, tooff=VECTOR(*mlist)[0];
    for (b=0; b<no_blocks; b++) {
      igraph_real_t fromsize = VECTOR(*mlist)[b];
      igraph_real_t tosize = n - tooff;
      int from, to;
      for (from=0; from<fromsize; from++) {
	for (to=0; to<tosize; to++) {
	  igraph_vector_push_back(&edges, fromoff + from);
	  igraph_vector_push_back(&edges, tooff + to);
	}
      }
      fromoff += fromsize;
      if (b+1 < no_blocks) { tooff += VECTOR(*mlist)[b+1]; }
    }    
  } else if (p > 0) {    
    int fromoff=0, tooff=VECTOR(*mlist)[0];
    for (b=0; b<no_blocks; b++) {
      igraph_real_t fromsize = VECTOR(*mlist)[b];
      igraph_real_t tosize = n - tooff;
      igraph_real_t maxedges= fromsize * tosize;
      igraph_real_t last=RNG_GEOM(p);
      while (last < maxedges) {
	int vto = floor(last/fromsize);
	int vfrom = last - (igraph_real_t) vto * fromsize;
	igraph_vector_push_back(&edges, fromoff + vfrom);
	igraph_vector_push_back(&edges, tooff + vto);
	last += RNG_GEOM(p);
	last += 1;
      }
      
      fromoff += fromsize;
      if (b+1 < no_blocks) { tooff += VECTOR(*mlist)[b+1]; }
    }
  }
  
  RNG_END();

  igraph_create(graph, &edges, n, /*directed=*/ 0);

  igraph_vector_destroy(&edges);
  igraph_vector_destroy(&csizes);
  IGRAPH_FINALLY_CLEAN(2);  
  
  return 0;
}
