/***********************************************************************/
/*               						       */
/*  Simulation of Galton-Watson tree		  		       */
/*  N. Berglund, december 2012 / february 2013                         */   
/*  http://www.univ-orleans.fr/mapmo/membres/berglund/simgw.html       */   
/*               						       */
/***********************************************************************/

#include <math.h>
#include <string.h>
#include <stdio.h>
#include <sys/types.h>

#define RAND_MAX 2147483647	/* 2^31 - 1 */

#define WINWIDTH 	560 
#define WINHEIGHT 	315 

#define PS_LEFT 	70
#define PS_RIGHT 	560
#define PS_BOTTOM 	330
#define PS_TOP 		680
#define PS_LIMITE 	50

#define NMAX 20000	/* max number of indiviuals */	
#define NGEN 16		/* number of generations */
#define NCHILD 4	/* max number of children */
#define NBINOMIAL 3	/* n of binomial law */
#define PBINOMIAL 0.45	/* p of binomial law */

#define XMIN -1.05
#define XMAX 1.05	/* x interval */
#define YMIN -0.05	
#define HMAX 2.0*WINHEIGHT/WINWIDTH
#define YMAX HMAX + 0.05	/* y interval */

#define R 0.01		/* radius of balls */

#define STEP 0.0001	/* integration step */
#define KSPRING 20.0	/* spring constant for attraction on graph */
#define KREPULSIVE 2.0	/* constant for repulsion force between siblings */
#define MAXWIDTH 0.1	/* max span of children in tree */

int gen;	/* current number of generations */
int nindiv;	/* number of individuals */

 /* Individuals are recorded as single list, one generation after the next */
 /* The information recorded is somewhat redundant, but this seems easier to handle */
 
double position[NMAX];
int igen[NGEN];		/* first individual of a generation */
int ngenindiv[NGEN];	/* number of individuals of a generation */
int generation[NMAX];	/* generation of individual */
int nchildren[NMAX];	/* number of children of individual */
int parent[NMAX];	/* index of parent */
int children[NMAX][NCHILD];	/* indices of children */

/***********************************************************************/
/* routines for manipulation of EPS file                               */ 
/***********************************************************************/

 typedef struct					/* port */
	{
		double x0, y0, x1, y1; 	/* dimensions    */
		double xmin, xmax, ymin, ymax;	/* local coordinates */
		int couleur;
		char texte[128];
	} t_port;


 int ps_compteur = 0, sauve_ps = 0, serie_ps = 0;
 FILE *ps_file;
 t_port graphport;


 void set_port_domain(port, xmin, xmax, ymin, ymax)	/* set port local coordinates */
 t_port *port;
 double xmin, xmax, ymin, ymax;
 {
	port->xmin = xmin;
	port->xmax = xmax;
	port->ymin = ymin;
	port->ymax = ymax;
 }

 int ps_test()		/* write "stroke - newpath" from time to time */
 {
	if (ps_compteur>PS_LIMITE)
	{
		ps_compteur = 0;
		fprintf(ps_file,"s\n");
		fprintf(ps_file,"n\n");
		return(1);
	}
	else return(0);
 }

 void close_ps()

 {
	fprintf(ps_file,"s\n");
	fprintf(ps_file,"showpage\n");
	fclose(ps_file);
 }
 
 void set_ps_bbox()	/* set bounding box of EPS file */

 {
	fprintf(ps_file, "%sBoundingBox: %i %i %i %i\n", "%%", PS_LEFT-2,
	PS_BOTTOM-2, PS_RIGHT+2, PS_TOP+2);
 }
 
 void init_eps()		/* initialise EPS file */
 {
	set_port_domain(&graphport, XMIN, XMAX, YMIN, YMAX);
	
	fprintf(ps_file, "%sPS-Adobe-2.0 EPSF-2.0", "%!");
	fprintf(ps_file, "\n");
	set_ps_bbox();

	fprintf(ps_file,"/l {lineto}	def\n");
	fprintf(ps_file,"/m {moveto}	def\n");
	fprintf(ps_file,"/n {newpath}	def\n");
	fprintf(ps_file,"/p {pnt}	def\n");
	fprintf(ps_file,"/s {stroke}	def\n");
	fprintf(ps_file,"/t {translate}	def\n");
	fprintf(ps_file,"/h {scale}	def\n");

	fprintf(ps_file,"newpath\n");
	fprintf(ps_file, "1 setlinewidth\n");
	fprintf(ps_file, "1 setlinejoin\n");
 }

  
 double ps_x_coord(x, port)	/* convert local coordinates to EPS coordinates */ 
 double x;
 t_port *port;

 {
	double coord;

	coord = PS_LEFT + (x-port->xmin)*(PS_RIGHT-PS_LEFT)/(port->xmax-port->xmin);
	return(coord);
 }

 double ps_y_coord(y, port)	/* convert local coordinates to EPS coordinates */ 
 double y;
 t_port *port;

 {
	double coord;

	coord = PS_BOTTOM + (y-port->ymin)*(PS_TOP-PS_BOTTOM)/(port->ymax-port->ymin);
	return(coord);
 }

 void ps_translate(x, y, port)		/* translate coordinate system */
 double x, y;
 t_port *port;
	
 {
	if(sauve_ps)
	{
		fprintf(ps_file, "%.4lg ", x);
		fprintf(ps_file, "%.4lg ", y);
		fprintf(ps_file, "t\n");
		ps_compteur += 1;
		ps_test();
	}
 }

 void ps_setgray(gray, port)		/* set gray value */ 
 double gray;
 t_port *port;

 {
	if(sauve_ps)
	{
		fprintf(ps_file, "s\nn\n");
		fprintf(ps_file, "%.3lg setgray\n", gray);
		ps_compteur = 0;
	}
 }

 void ps_ligne(x1, y1, x2, y2, port)	/* draw line */ 
 double x1, x2, y1, y2;
 t_port *port;

 {
	if(sauve_ps)
	{
		fprintf(ps_file, "%.4lg ", ps_x_coord(x1, port));
		fprintf(ps_file, "%.4lg ", ps_y_coord(y1, port));
		fprintf(ps_file, "m\n");
		fprintf(ps_file, "%.4lg ", ps_x_coord(x2, port));
		fprintf(ps_file, "%.4lg ", ps_y_coord(y2, port));
		fprintf(ps_file, "l\n");
		ps_compteur += 2;
		ps_test();
	}
 }
 
 void ps_dotfill(x, y, r, gray, port)		/* draw filled circle */ 
 double x, y, r, gray;
 t_port *port;

 {
	if(sauve_ps)
	{
		ps_translate(ps_x_coord(x,port), ps_y_coord(y,port));
		ps_setgray(gray, port);
		fprintf(ps_file, "0 0 %.5lg 0 360 arc \n", r);
		fprintf(ps_file, "closepath fill\n");
		ps_setgray(0.0, port);
		fprintf(ps_file, "0 0 %.5lg 0 360 arc \n", r);
		ps_translate(-ps_x_coord(x,port), -ps_y_coord(y,port));

	}
 }

/***********************************************************************/
/* specific EPS routines for Galton-Watson program                     */ 
/***********************************************************************/

 
double ycoord(int geni)		/* change y coordinate to get smaller height for later generations */ 
{
  double shrink = 1.04, stretch;
  
  stretch = 1.0/(1.0 - pow(shrink, -(double)NGEN));
  return(HMAX*(1.0+(pow(shrink,-(double)geni) - 1.0)*stretch));
  
//   return(HMAX-(double)geni/(double)gen*(double)HMAX);   /* alternative: linear interpolation */
//   return(HMAX*exp(-0.2*(double)geni/(double)NGEN));	/* alternative: hyperbolic interpolation */

}

void print_eps()	/* generate EPS and PDF files */ 
{
  int i, j, geni;
  double r, x, y, r1, x1, y1;
  
 
  ps_file = fopen("gw.eps", "w");
  sauve_ps = 1;

  init_eps();

  for (i=0; i<nindiv; i++)
  {
    geni = generation[i];
    y = ycoord(geni);
    x = position[i];
    
    for (j=0; j<nchildren[i]; j++)
    {
      y1 = ycoord(geni+1);
      x1 = position[children[i][j]];
      ps_ligne(x, y, x1, y1, &graphport);
    }

    ps_dotfill(x, y, 2.0, 0.0, &graphport);
  }


  close_ps();
  system("epstopdf gw.eps");
}

/***********************************************************************/
/* end of routines for manipulation of EPS file                        */ 
/***********************************************************************/


 
void init_rand(offset)		/* initialise random number generator to get different realisations */ 
int offset;
{
  int i;

  for (i=0; i<offset; i++) rand();		
}
 


double binomial(n, p) 
 /* returns random variable of binomial(n,p) distribution */
 int n;
 double p;
 {
	int i, binomial = 0, rnd;

	for (i=0; i<n; i++)
 		if ((double)rand()<p*RAND_MAX) binomial++;

// 	printf("binomial=%i\n", binomial);
	return(binomial);
 }



void init_gw()		/* initialise variables, generation 0 */
{
  nindiv = 1;
  gen = 0;
  
  position[0] = 0.0;
  igen[0] = 0;
  ngenindiv[0] = 1;
  generation[0] = 0;
}



int add_generation(n, p)	/* add one generation */ 
int n;
double p;
{
  int i, j, nc, ngen=0, iparent;
  double x, left, right, width;
  
  gen++;
  igen[gen] = nindiv; 
  for (i=0; i<ngenindiv[gen-1]; i++)
  {
    iparent = igen[gen-1]+i;	
    nc = binomial(n, p);	
    ngen+= nc;
    nchildren[iparent] = nc;	
    
    /* compute positions */
    x = position[iparent];
    
    /* determine position of left neighbour of parent */
    if (iparent == igen[generation[iparent]]) left = -1.0;
    else left = position[iparent-1];
    
    /* determine position of left neighbour of parent */
    if (iparent == igen[generation[iparent]] + ngenindiv[generation[iparent]] - 1) right = 1.0;
    else right = position[iparent+1];
    
    width = right - x;
    if (x - left < width) width = x - left;
    width = 0.25*width;
    if (width > MAXWIDTH) width = MAXWIDTH;
        
    for (j=0; j<nc; j++)
    {
      generation[nindiv] = gen;
      parent[nindiv] = iparent;
      children[iparent][j] = nindiv;
      
      if (nc==1) position[nindiv] = x;
      else position[nindiv] = x - width + 2.0*width*((double)j/(double)(nc-1));
      
      printf("%i %.3f\n", nindiv, position[nindiv]);
      nindiv++;
      
    }
    
  }
  ngenindiv[gen] = ngen;
  return(ngen);
    
}

void print_generations()	/* print generation tree on stdout */ 
{
  int i, j;
  
  printf("Generation=%i\n", gen);
  printf("Individuals=%i\n", nindiv);
  
  printf("igen: ");
  for (i=0; i<gen+1; i++) printf("%i ", igen[i]);
  printf("\n");
  
  printf("ngenindiv: ");
  for (i=0; i<gen+1; i++) printf("%i ", ngenindiv[i]);
  printf("\n");

  printf("generation: ");
  for (i=0; i<nindiv; i++) printf("%i ", generation[i]);
  printf("\n");

  printf("nchildren: ");
  for (i=0; i<nindiv; i++) printf("%i ", nchildren[i]);
  printf("\n");

  printf("parent: ");
  for (i=0; i<nindiv; i++) printf("%i ", parent[i]);
  printf("\n");

  printf("children: ");
  for (i=0; i<nindiv; i++) 
  {
    for (j=0; j<nchildren[i]; j++) printf("%i ", children[i][j]);
    printf(" - ");
  }
  printf("\n\n");

}

double attractive_force(double x)	/* harmonic attractive force between parents and children */
{
  return(-KSPRING*x);
}

double repulsive_force(double x)	/* logarithmic repulsive force between siblings */ 
{
  if (x==0.0) x=0.0000001;
  if (x>1.0) x=1.0;
  if (x<-1.0) x=-1.0;
  if (x>0.0) return(-KREPULSIVE*log(x));
  else return(KREPULSIVE*log(-x));
}

void relax(int nsteps)		/* evolve the tree by attraction/repulsion between sites */
{
    int n, i, j, k;
    double totalforce = 0.0;
    
    for (n=0; n<nsteps; n++)
    {
      for (i=0; i<nindiv; i++)
      {
        totalforce = 0.0;
	
	/* get attracted by parents */
	if (i>0) totalforce += attractive_force(position[i] - position[parent[i]]);
		
	/* get attracted by children */
	for (j=0; j<nchildren[i]; j++) 
	  totalforce += attractive_force(position[i] - position[children[i][j]]);
	
	
	/* get pushed away from left sibling */
	if ((i>0)&&(generation[i-1]==generation[i])) 
	  totalforce += repulsive_force(position[i] - position[i-1]);
		
	/* get pushed away from right sibling */
	if (generation[i+1]==generation[i]) 
	  totalforce += repulsive_force(position[i] - position[i+1]);
		
	/* get pushed away from left wall */
	if (i == igen[generation[i]]) 
	  totalforce += repulsive_force(position[i] + 1.0);
	  		
	/* get pushed away from right wall */
	if (i == igen[generation[i]] + ngenindiv[generation[i]] - 1) 
	  totalforce += repulsive_force(position[i] - 1.0);
		
	position[i] += totalforce*STEP;
      }      
    }
}


void graph()	/* generate Galton-Watson tree */ 
{
  int i = 0, j, pop = 1;
  
  init_gw();
  print_generations();
  
  while ((i<NGEN)&&(pop>0))
  {
     pop = add_generation(NBINOMIAL, PBINOMIAL);
     print_generations();
     for (j=0; j<50; j++) relax(100);
     if (pop==0) printf("Population extinct\n");
     i++;
  }
  print_eps();
 }


int main(int argc, char** argv)
{
  
  init_rand(3000);
  graph();
  
  return 0;
}