/**********  segtre_mig.c **********************************
*
*	This subroutine uses a Monte Carlo algorithm described in
*	Hudson,R. 1983. Theor.Pop.Biol. 23: 183-201, to produce
*	a history of a random sample of gametes under a neutral
*	Wright-Fisher model with recombination and geographic structure. 
*	Input parameters
*	are the sample size (nsam), the number of sites between
*	which recombination can occur (nsites), and the recombination
*	rate between the ends of the gametes (r). The function returns
*	nsegs, the number of segments the gametes were broken into
*	in tracing back the history of the gametes.  The histories of
*	these segments are passed back to the calling function in the
*	array of structures seglst[]. An element of this array,  seglst[i],
* 	consists of three parts: (1) beg, the starting point of
*	of segment i, (2) ptree, which points to the first node of the
*	tree representing the history of the segment, (3) next, which
*	is the index number of the next segment.
*	     A tree is a contiguous set of 2*nsam nodes. The first nsam
*	nodes are the tips of the tree, the sampled gametes.  The other
*	nodes are the nodes ancestral to the sampled gametes. Each node
*	consists of an "abv" which is the number of the node ancestral to
*	that node, an "ndes", which is not used or assigned to in this routine,
*	and a "time", which is the time (in units of 4N generations) of the
*	node. For the tips, time equals zero.
*	Returns a pointer to an array of segments, seglst.

**************************************************************************/

#include  "my.h"
#define SEGINC 80 

extern int flag;

int nchrom, begs, nsegs;
long nlinks ;
static int *nnodes = NULL ;  
double t;

static unsigned seglimit = SEGINC ;
static unsigned maxchr ;

struct seg{
	int beg;
	int end;
	int desc;
	};

struct chromo{
	int nseg;
	int pop;
	struct seg  *pseg;
	} ;

static struct chromo *chrom = NULL ;

struct node{
	int abv;
	int ndes;
	float time;
	} *ptree1, *ptree2;

struct segl {
	int beg;
	struct node *ptree;
	int next;
	}  ;
static struct segl *seglst = NULL ;

	struct segl *
segtre_mig(npop,nsam, inconfig, nsites, r, mig_rate, pnsegs, alphag, t0)
	int npop, inconfig[], nsites, nsam, *pnsegs;
	double r, mig_rate, alphag, t0 ;

{
	int i, j, k, seg, dec, pop, c1, c2, ind, rchrom ;
	int migrant, source_pop, *config, flagint ;
	double  ran1(), trm, tcoal, ttemp ;
	double prec, nnm1, nnm0, mig, ran, coal_prob, prob, rdum ;
	char c;
	

/* Initialization */
	if( chrom == NULL ) {
	   maxchr = nsam + 20 ;
	   chrom = (struct chromo *)malloc( (unsigned)( maxchr*sizeof( struct chromo) )) ;
	  if( chrom == NULL ) perror( "malloc error. segtre");
	  }
	if( nnodes == NULL ){
		nnodes = (int*) malloc((unsigned)(seglimit*sizeof(int)))  ;
		if( nnodes == NULL ) perror("malloc error. segtre_mig");
		}
	if( seglst == NULL ) {
		seglst = (struct segl *)malloc((unsigned)(seglimit*sizeof(struct segl)) ) ;
		if( seglst == NULL ) perror("malloc error. segtre_mig.c 2");
		}

	config = (int *)malloc( (unsigned) ((npop+1)*sizeof(int) )) ;
	if( config == NULL ) perror("malloc error. segtre.");
	for(pop=0;pop<npop;pop++) config[pop] = inconfig[pop] ;
	for(pop=ind=0;pop<npop;pop++)
		for(j=0; j<inconfig[pop];j++,ind++) {
			
			chrom[ind].nseg = 1;
			if( !(chrom[ind].pseg = (struct seg*)malloc((unsigned)sizeof(struct seg)) ))
			  ERROR("calloc error. se1");

			(chrom[ind].pseg)->beg = 0;
			(chrom[ind].pseg)->end = nsites-1;
			(chrom[ind].pseg)->desc = ind ;
			chrom[ind].pop = pop ;
			}
	seglst[0].beg = 0;
	if( !(seglst[0].ptree = (struct node *)calloc((unsigned)(2*nsam),sizeof(struct node)) ))
		 perror("calloc error. se2");



	nnodes[0] = nsam - 1 ;
	nchrom=nsam;
	nlinks = ((long)(nsam))*(nsites-1) ;
	nsegs=1;
	t = 0.;
	r /= (nsites-1);
	flagint = 0 ;

/* Main loop */

	while( nchrom > 1 ) {
		prec = nlinks*r;

		mig = nchrom*mig_rate ;
		while( ( rdum = ran1() )  == 1.0 ) ;
		trm = -log(1.-rdum )/(prec+mig)  ;
		
		coal_prob = 0. ;
		for(pop=0; pop<npop ; pop++)
			 coal_prob += ((double)config[pop])*(config[pop]-1.) ;
			 
		while( ( rdum = ran1() )  == 1.0 ) ;
		if( alphag == 0 ) tcoal = -log( 1.-rdum ) /coal_prob ;
		else if( t >= t0 )  tcoal = -log( 1.-rdum ) /(coal_prob*exp(alphag*t0)) ;
		else tcoal = (1./alphag)*log( 1. + alphag*exp(-alphag*t)*( -1.0 *
	    	     log(1.-rdum) / coal_prob   ) ) ;
			 
			 
		ttemp = MIN( trm, tcoal) ;
		
		if( (t < t0) && ( (t+ttemp) > t0 ))  t = t0 ;
		else {
		t += ttemp ;	
		if( trm < tcoal) {
		   if( (ran = ran1()) < ( prec / (prec+mig) ) ){ 
		
		    /*recombination*/
		    	rchrom = re(nsam);
			    config[ chrom[rchrom].pop ] += 1 ;
			   }
		   else  {  /* migration event */
			  migrant = nchrom*ran1() ;
			  while( (source_pop=npop*ran1()) == chrom[migrant].pop ) ;
			  config[chrom[migrant].pop] -= 1;
			  config[source_pop] += 1;
			  chrom[migrant].pop = source_pop ;
			  }
			}
		else { 										 /* coalescent event */
			/* pick the two, c1, c2  */
			ran = ran1();
			prob = 0.0 ;
			for(pop=0;pop<npop;pop++) {
				prob += ((double)config[pop])*(config[pop]-1.)/coal_prob ;
				if( ran < prob ) break;
				}
			if( pop == npop ) pop = npop-1 ;
			pick2_chrom( pop, config, &c1,&c2);  /* c1 and c2 are chrom's to coalesce */
			dec = ca(nsam,nsites,c1,c2 );
			config[pop] -= dec ;
			}
		}
		}  
	*pnsegs = nsegs ;
	free(config); 
	return( seglst );
}

/******  recombination subroutine ***************************
   Picks a chromosome and splits it in two parts. If the x-over point
   is in a new spot, a new segment is added to seglst and a tree set up
   for it.   ****/


	int
re(nsam)
	int nsam;
{
	struct seg *pseg, *pseg2;
	int i, l, lsg, lsgm1, newsg, ic, jseg, k, is, in, spot;
	double ran1();


/* First generate a random x-over spot, then locate it as to chrom and seg. */

	spot = nlinks*ran1() + 1.;

    /* get chromosome # (ic)  */

	for( ic=0; ic<nchrom ; ic++) {
		lsg = chrom[ic].nseg ;
		lsgm1 = lsg - 1;
		pseg = chrom[ic].pseg;
		l = ( (pseg+lsgm1)->end ) - (pseg->beg);
		if( spot <= l ) break;
		spot -= l ;
		}
	is = pseg->beg + spot -1;

   /* get seg # (jseg)  */

	for( jseg=0; is >= (pseg+jseg)->end ; jseg++) ;
	if( is >= (pseg+jseg)->beg ) in=1;
	else in=0;
	newsg = lsg - jseg ;

   /* copy last part of chrom to nchrom  */

	nchrom++;
	if( nchrom >= maxchr ) {
	    maxchr += 20 ;
	    chrom = (struct chromo *)realloc( chrom, (unsigned)(maxchr*sizeof(struct chromo))) ;
	    if( chrom == NULL ) perror( "malloc error. segtre2");
	    }
	if( !( pseg2 = chrom[nchrom-1].pseg = (struct seg *)calloc((unsigned)newsg,sizeof(struct seg)) ) )
		ERROR(" alloc error. re1");
	chrom[nchrom-1].nseg = newsg;
	chrom[nchrom-1].pop = chrom[ic].pop ;
	pseg2->end = (pseg+jseg)->end ;
	if( in ) {
		pseg2->beg = is + 1 ;
		(pseg+jseg)->end = is;
		}
	else pseg2->beg = (pseg+jseg)->beg ;
	pseg2->desc = (pseg+jseg)->desc ;
	for( k=1; k < newsg; k++ ) {
		(pseg2+k)->beg = (pseg+jseg+k)->beg;
		(pseg2+k)->end = (pseg+jseg+k)->end;
		(pseg2+k)->desc = (pseg+jseg+k)->desc;
		}

	lsg = chrom[ic].nseg = lsg-newsg + in ;
	lsgm1 = lsg - 1 ;
	nlinks -= pseg2->beg - (pseg+lsgm1)->end ;
if( !(chrom[ic].pseg = 
     (struct seg *)realloc(chrom[ic].pseg,(unsigned)(lsg*sizeof(struct seg)) )) )
		perror( " realloc error. re2");
	if( in ) {
		begs = pseg2->beg;
		for( i=0,k=0; (k<nsegs-1)&&(begs > seglst[seglst[i].next].beg-1);
		   i=seglst[i].next, k++) ;
		if( begs != seglst[i].beg ) {
						/* new tree  */

	   	   if( nsegs >= seglimit ) {  
	   	   	  seglimit += SEGINC ;
	   	      nnodes = (int *)realloc( nnodes,(unsigned)(sizeof(int)*seglimit)) ; 
	   	      if( nnodes == NULL) perror("realloc error. 1. segtre_mig.c");
	   	      seglst =
	   	      	 (struct segl *)realloc( seglst,(unsigned)(sizeof(struct segl)*seglimit));
	   	      if(seglst == NULL ) perror("realloc error. 2. segtre_mig.c");
	   	      /*  printf("seglimit: %d\n",seglimit);  */
	   	      } 
	   	   seglst[nsegs].next = seglst[i].next;
	   	   seglst[i].next = nsegs;
	   	   seglst[nsegs].beg = begs ;
		   if( !(seglst[nsegs].ptree = (struct node *)calloc((unsigned)(2*nsam), sizeof(struct
			 node)) )) perror("calloc error. re3.");
		   nnodes[nsegs] = nnodes[i];
		   ptree1 = seglst[i].ptree;
		   ptree2 = seglst[nsegs].ptree;
		   nsegs++ ;
		   for( k=0; k<=nnodes[i]; k++) {
		      (ptree2+k)->abv = (ptree1+k)->abv ;
		      (ptree2+k)->time = (ptree1+k)->time;
		      }
		   }
	}
	return(ic) ;
}

/***** common ancestor subroutine **********************
   Pick two chromosomes and merge them. Update trees if necessary. **/

	int
ca(nsam, nsites,c1,c2)
	int nsam,c1,c2;
	int  nsites;
{
	int yes1, yes2, seg1, seg2, seg ;
	int tseg, start, end, desc, k;
	struct seg *pseg;
	struct node *ptree;

	seg1=0;
	seg2=0;

	if( !(pseg = (struct seg *)calloc((unsigned)nsegs,sizeof(struct seg) ))) 
		perror("alloc error.ca1");

	tseg = -1 ;

	for( seg=0, k=0; k<nsegs; seg=seglst[seg].next, k++) {
		start = seglst[seg].beg;
		yes1 = isseg(start, c1, &seg1);
		yes2 = isseg(start, c2, &seg2);
		if( yes1 || yes2 ) {
			tseg++;
			(pseg+tseg)->beg=seglst[seg].beg;
			end = ( k< nsegs-1 ? seglst[seglst[seg].next].beg-1 : nsites-1 ) ;
			(pseg+tseg)->end = end ;

			if( yes1 && yes2 ) {
				nnodes[seg]++;
				if( nnodes[seg] >= (2*nsam-2) ) tseg--;
				else
					(pseg+tseg)->desc = nnodes[seg];
				ptree=seglst[seg].ptree;
				desc = (chrom[c1].pseg + seg1) ->desc;
				(ptree+desc)->abv = nnodes[seg];
				desc = (chrom[c2].pseg + seg2) -> desc;
				(ptree+desc)->abv = nnodes[seg];
				(ptree+nnodes[seg])->time = t;

				}
			else {
				(pseg+tseg)->desc = ( yes1 ?
				   (chrom[c1].pseg + seg1)->desc :
				  (chrom[c2].pseg + seg2)->desc);
				}
			}
		}
	nlinks -= links(c1);
	free(chrom[c1].pseg) ;
	if( tseg < 0 ) {
		free(pseg) ;
		chrom[c1].pseg = chrom[nchrom-1].pseg;
		chrom[c1].nseg = chrom[nchrom-1].nseg;
		chrom[c1].pop = chrom[nchrom-1].pop ;
		if( c2 == nchrom-1 ) c2 = c1;
		nchrom--;
		}
	else {
		if( !(pseg = (struct seg *)realloc(pseg,(unsigned)((tseg+1)*sizeof(struct seg)))))
			perror(" realloc error. ca1");
		chrom[c1].pseg = pseg;
		chrom[c1].nseg = tseg + 1 ;
		nlinks += links(c1);
		}
	nlinks -= links(c2);
	free(chrom[c2].pseg) ;
	chrom[c2].pseg = chrom[nchrom-1].pseg;
	chrom[c2].nseg = chrom[nchrom-1].nseg;
	chrom[c2].pop = chrom[nchrom-1].pop ;
	nchrom--;
	if(tseg<0) return( 2 );  /* decrease of nchrom is two */
	else return( 1 ) ;
}

/*** Isseg: Does chromosome c contain the segment on seglst which starts at
	    start? *psg is the segment of chrom[c] at which one is to begin 
	    looking.  **/

	int
isseg(start, c, psg)
	int start, c, *psg;
{
	int ns;
	struct seg *pseg;

	ns = chrom[c].nseg;
	pseg = chrom[c].pseg;

	for(  ; ( (pseg+(*psg))->beg <= start) &&
				   ((*psg) < ns ) ;  ++(*psg) )
		if( (pseg+(*psg))->end >= start ) return(1);
	return(0);
}
	


	int
pick2_chrom(pop,config,pc1,pc2)
	int pop, *pc1, *pc2, config[];
{
	int c1, c2, cs,cb,i, count;
	
	pick2(config[pop],&c1,&c2);
	cs = (c1>c2) ? c2 : c1;
	cb = (c1>c2) ? c1 : c2 ;
	i=count=0;
	for(;;){
		while( chrom[i].pop != pop ) i++;
		if( count == cs ) break;
		count++;
		i++;
		}
	*pc1 = i;
	i++;
	count++;
	for(;;){
		while( chrom[i].pop != pop ) i++;
		if( count == cb ) break;
		count++;
		i++;
		}
	*pc2 = i ;
}
	
	

/****  links(c): returns the number of links between beginning and end of chrom **/

	int
links(c)
	int c;
{
	int ns;

	ns = chrom[c].nseg - 1 ;

	return( (chrom[c].pseg + ns)->end - (chrom[c].pseg)->beg);
}