/*  EXACTLIKE.C - this program computes the likelihood surface as in figure 3   */
/*  of Wakeley and Takahashi 2003 Mol. Biol. Evol. 20:208-213.  Unfortunately,  */
/*  you have to enter the data (sample size, total number of segregating sites, */
/*  and number of singletons) in main() below, then compile and run............ */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
/* #include "/home/john/Cprogs/Headers/deviates2.h" */
#include "Obaku:John:Cprograms:Headers:random.h" 
/* the line(s) above includes a file with the random number generator I like; */
/* put your own or a system supplied one in unirandom()/sunirandom() below    */

#define MAXDESC 141  /* The maximum number of descendents of a lineage (in one generation) */

struct node {   
	int id;
	int numdesc;
	double time;
	struct node *desc[MAXDESC];
};  /* the nodes in the genealogy */

struct ancestor {
	int id;
	int numdesc;
	int desc[MAXDESC];
};  /* lists results of balls-in-boxes of MakeTree(),PickAncestors() */

/* MYERROR - simple error alert */ 

void MyError(error_text)
char  error_text[];
{
	fprintf(stderr,"\nJoe Mamas run-time error...\n");
	fprintf(stderr,"\n\t%s\n",error_text);
	fprintf(stderr,"\n...now exiting to system...\n\n");
	exit(1);
}

/*	SUNIRANDOM - Seed your random number generator	*/

void sunirandom(seed)
unsigned int seed;
{
	/* put your own srandom here */
	srandom(seed);
}

/*	UNIRANDOM - Uniform [0,1) random deviates.	*/

#define IM 2147483648.0

double unirandom()
{
	/* put your own random number generator here */
	return (double)random()/IM; 
}

#undef IM

/*	EXPONENTIALDEV - mean 1, deviates. For mean m, multiple by m.	*/

double exponentialdev() 
{ 
	double dum;
 
	do 
		dum = unirandom();
	while( dum == 0.0 ); 
	return -log(dum); 
} 

void InitTree(numseqs,tree)
int numseqs;
struct node *tree;
{
	int x,y;

	for(x=0;x<(2*numseqs);x++) {
		(tree + x)->id = x;
		(tree + x)->numdesc = 0;
		(tree + x)->time = 0.0;
		for(y=0;y<MAXDESC;y++)
			(tree + x)->desc[y] = NULL;
	}
}

/*  PICKANCESTORS - throws nseqs balls into Ne boxes and records the numbers of balls  */
/*  in each box in the ancestors array; increments the current time (in generations).  */ 

void PickAncestors(nseqs,numancp,twoN,time,ancestors)
int nseqs,*numancp;
double twoN,*time;
struct ancestor *ancestors;
{
	int y,z,pick;
	
	if(nseqs==2) {
		*numancp = 1;
		ancestors[0].numdesc = 2;
		ancestors[0].desc[0] = 0;
		ancestors[0].desc[1] = 1;
		*time += twoN*exponentialdev();
	} else {
		do {
			*numancp = 0;
			for(y=0;y<nseqs;y++) {
				ancestors[y].id = ((int)twoN) + 1;
				ancestors[y].numdesc = 0;
				for(z=0;z<MAXDESC;z++)
					ancestors[y].desc[z] = 0;
			}
			for(y=0;y<nseqs;y++) {
				pick = (int)(twoN*unirandom()); 
				for(z=0;z<*numancp;z++) {
					if(pick == ancestors[z].id) {
						if( ancestors[z].numdesc == MAXDESC ) MyError("too many descendents");
						ancestors[z].desc[ancestors[z].numdesc] = y;
						ancestors[z].numdesc += 1;
						z = *numancp+1;
					}
				}
				if(z==*numancp) {
					ancestors[*numancp].id = pick;
					ancestors[*numancp].desc[0] = y;
					ancestors[*numancp].numdesc += 1;
					*numancp += 1;
				}
			}
			*time += 1.0;
		} while(*numancp==nseqs);
	}
}

/*  MAKETREE - uses PickAncestors() successively and puts the result into the genealogy,  */
/*  one node for each ancestor that has more than one descendent; increment tree length.  */

void MakeTree(numseqs,rootp,twoN,Ttotp,tree)
int numseqs,*rootp;
double twoN,*Ttotp;
struct node *tree;
{
	int x,y,z,cnode,numanc,current,flag;
	double t;
	struct node **list;
	struct ancestor *ancestors;
	
	list = (struct node **)malloc( (size_t)numseqs*sizeof(struct node *) );
	if(list == NULL) MyError("problem with list memory allocation");
	ancestors = (struct ancestor *)malloc( (size_t)numseqs*sizeof(struct ancestor) );
	if(ancestors == NULL) MyError("problem with ancestors memory allocation");
	
	for(x=0;x<numseqs;x++) 
		list[x] = tree + x;
		
	t = 0.0;
	*Ttotp = 0.0;
	x = numseqs;
	cnode = numseqs;
	
	while(x>1) {
		PickAncestors(x,&numanc,twoN,&t,ancestors);
		if(numanc < x) { 
			for(y=0;y<numanc;y++) 
				if(ancestors[y].numdesc > 1) {
					tree[cnode].numdesc = ancestors[y].numdesc;
					tree[cnode].time = t;
					for(z=0;z<tree[cnode].numdesc;z++) {
						tree[cnode].desc[z] = list[ancestors[y].desc[z]];
						if(z>0)
							list[ancestors[y].desc[z]] = NULL;
						else 
							list[ancestors[y].desc[z]] = tree + cnode;
					}
					*Ttotp += t*((double)tree[cnode].numdesc);
					for(z=0;z<tree[cnode].numdesc;z++)
						*Ttotp -= (tree+cnode)->desc[z]->time;
					cnode += 1; 
				}
			flag = 0;
			for(y=0;y<numanc;y++) 
				if(list[y]==NULL) {
					if(flag==0) {
						current = y+1; 
						flag = 1;
					}
					for(z=current;z<x;z++) 
						if(list[z]!=NULL) {
							list[y] = list[z];
							list[z] = NULL;
							current = z+1;
							z = x;
						}
				}
			x = numanc;
		}
	}
	*rootp = cnode-1;
	
	free( (void *)list );
	free( (void *)ancestors );
}

/*  EXTERNALTIME - counts the combined length of all external branches in the genealogy  */

double ExternalTime(numseqs,root,tree)
int numseqs,root;
struct node *tree;
{
	int x,y;
	double Tsum;

	Tsum = 0.0;
	for(x=root;x>=numseqs;x--) 
		for(y=0;y<tree[x].numdesc;y++) 
			if( (tree+x)->desc[y]->id < numseqs )
				Tsum += (tree[x].time - (tree+x)->desc[y]->time); 
	return Tsum;
}

void WriteTime(start,finish,fptr)
long start,finish;
FILE *fptr;
{
	double seconds,hours,minutes;
	
	seconds = (double)(finish-start);
	modf(seconds/3600.0,&hours);
	seconds -= 3600.0*hours;
	modf(seconds/60.0,&minutes);
	seconds -= 60.0*minutes;
	fprintf(fptr,"\n\nTotal Running Time = %1.0f:%1.0f:%1.0f\n\n",hours,minutes,seconds);
	fflush(fptr);
}

int main() 
{
	int x,y,z,numseqs,numreps,root,nstats;
	int Stot,S1;
	double part1,part2,Ttot,T1,twoN,theta,u,count,*sums;
	struct node *tree;
	time_t first,last,start,finish;
	FILE *fptr;
	int popsizes[32] = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32};
	double thetas[14] = {1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0};
	
	/* fptr = fopen("like_result","w"); */
	fptr = stdout; 
	
	numreps = 1000;
	nstats = 10;

 /* ENTER YOUR DATA HERE */
	numseqs = 141;
	Stot = 50;
	S1 = 31;
	/* Stot = 28;
	S1 = 5; */

	time(&first); 

	fprintf(fptr,"\n{");

	for(x=1;x<26;x++) {

		fprintf(fptr,"{");

		for(y=0;y<14;y++) {
	
			twoN = (double)popsizes[x];
			u = thetas[y]/(2.0*twoN);
	
			time(&start); 
			srandom( (unsigned)start );
	
			tree = (struct node *)malloc( (size_t)2*numseqs*sizeof(struct node) );
			if(tree == NULL) MyError("problem getting memory for tree in main");
			sums = (double *)malloc( (size_t)nstats*sizeof(double) );
			if(sums == NULL) MyError("problem getting memory for sums in main");

			for(z=0;z<nstats;z++) 
				sums[z] = 0.0;
		
			count = 0.0;
			for(z=0;z<numreps;z++) {
				InitTree(numseqs,tree);
				MakeTree(numseqs,&root,twoN,&Ttot,tree);
				/* WriteTreeInfo(root,tree,fptr); */
				T1 = ExternalTime(numseqs,root,tree);
				part1 = ((double)(Stot-S1))*log(u*(Ttot-T1)) - u*(Ttot-T1);
				part2 = ((double)S1)*log(u*T1) - u*T1;
				sums[2] += part1 + part2;
				sums[0] += Ttot;
				sums[1] += T1;
				count+=1.0;
			} 

			/* for(z=0;z<2;z++) 
				fprintf(fptr,"%lf\n",u*sums[z]/count); */
				/* fprintf(fptr,"%lf\n",sums[z]/(count*2.0*twoN)); */
			if(y<13)
				fprintf(fptr,"%5.4lf,",sums[2]/count);
			else
				if(x<25)
					fprintf(fptr,"%5.4lf},\n",sums[2]/count);
				else
					fprintf(fptr,"%5.4lf}",sums[2]/count);
			fflush(fptr); 
			
			free( (void *)tree );
			free( (void *)sums );
			
			time(&finish);
			/* WriteTime(start,finish,fptr); */
	
		}
	}

	fprintf(fptr,"}\n");

	time(&last);
	WriteTime(first,last,fptr); 
	
	fclose(fptr);
	
	return 0;
}