/* Computes the expected joint site-frequency spectrum for the isolation */
/* model in Wakeley and Hey (1997), using the equations in that paper.   */

#include <stdio.h>
#include <stdlib.h>
#include <math.h> 
#include <time.h> 

void MyError(char*);
void AllocateGlobals(int,int);
void FreeGlobals();
double term(int,int);
double myprod(int,int,int);
double gammln(double);
double factln(int);
double bico(int,int);
void MakeTables(int,int);
double mypolya(int,int,int,int);
void MakeCubes(int,int);
double P1(int,int,double);
double P2(int,int,double);
double a1(int);
void CalculateAll(int,int,double,double,double,double);
void CalculateFreqs(int,int,double,double,double,double);

double *cube1,*cube2,*carray,*array1,*array2,**nchoosek,**table1,**table2;

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);
}

void AllocateGlobals(n1,n2)
int n1,n2;
{
	carray = (double *) malloc( (size_t) (n1+n2+1)*(n1+n2+1)*sizeof(double) );
	if( !carray ) MyError("problem getting memory for carray in AllocateGlobals");
	nchoosek = (double **) malloc( (size_t) (n1+n2+1)*sizeof(double *) );
	if( !nchoosek ) MyError("problem getting memory for nchoosek in AllocateGlobals");
	array1 = (double *) malloc( (size_t) (n1+1)*(n1+1)*sizeof(double) );
	if( !array1 ) MyError("problem getting memory for array1 in AllocateGlobals");
	table1 = (double **) malloc( (size_t) (n1+1)*sizeof(double *) );
	if( !table1 ) MyError("problem getting memory for table1 in AllocateGlobals");
	array2 = (double *) malloc( (size_t) (n2+1)*(n2+1)*sizeof(double) );
	if( !array2 ) MyError("problem getting memory for array2 in AllocateGlobals");
	table2 = (double **) malloc( (size_t) (n2+1)*sizeof(double *) );
	if( !table2 ) MyError("problem getting memory for table2 in AllocateGlobals");
	cube1 = (double *) malloc( (size_t) (n1+1)*(n1+1)*(n1+1)*sizeof(double) );
	if( !cube1 ) MyError("problem getting memory for cube1 in AllocateGlobals");
	cube2 = (double *) malloc( (size_t) (n2+1)*(n2+1)*(n2+1)*sizeof(double) );
	if( !cube2 ) MyError("problem getting memory for cube2 in AllocateGlobals");
}

void FreeGlobals()
{
	free( (void *)nchoosek );
	free( (void *)carray );
	free( (void *)table1 );
	free( (void *)table2 );
	free( (void *)array1 );
	free( (void *)array2 );
	free( (void *)cube1 );
	free( (void *)cube2 );
}

double term(i,j)
int i,j;
{
	double di,dj;

	if(i < 2 || j < 2) MyError("problem in term");
	di = (double)i;
	dj = (double)j;
	return (di*(di-1.0))/(dj*(dj-1.0));
}

double myprod(n,nt,i) 
int n,nt,i;
{
	int j;
	double product;

	product = 1.0;
	for(j=nt;j<i;j++)
		product *= 1.0 - term(i,j);
	for(j=i+1;j<=n;j++)
		product *= 1.0 - term(i,j);

	return product;
}

double gammln(xx)
double xx;
{
	double x,y,tmp,ser;
	static double cof[6]={76.18009172947146,-86.50532032941677,
		24.01409824083091,-1.231739572450155,
		0.1208650973866179e-2,-0.5395239384953e-5};
	int j;

	y=x=xx;
	tmp=x+5.5;
	tmp -= (x+0.5)*log(tmp);
	ser=1.000000000190015;
	for (j=0;j<=5;j++) ser += cof[j]/++y;
	return -tmp+log(2.5066282746310005*ser/x);
}

double factln(n)
int n;
{
	static double a[101];

	if(n < 0) MyError("negative factorial in factln");
	if(n <= 1) return 0.0;
	if(n <= 100) return a[n] ? a[n] : (a[n]=gammln(n+1.0));
	else return gammln(n+1.0);
}

double bico(n,k)
int n,k;
{
	return floor(0.5+exp(factln(n)-factln(k)-factln(n-k)));
}

void MakeTables(n1,n2)
int n1,n2;
{
	int x,y;

	for(x=0;x<(n1+n2+1);x++) {
		nchoosek[x] = carray + x*(n1+n2+1);
		for(y=0;y<=x;y++)
			nchoosek[x][y] = bico(x,y);
	}
	for(x=0;x<(n1+1);x++) { 
		table1[x] = array1 + (x*(n1+1));
		if(x>=2)
			for(y=2;y<(n1+1);y++) 
				table1[x][y] = myprod(n1,x,y); 
	}
	for(x=0;x<(n2+1);x++) { 
		table2[x] = array2 + (x*(n2+1));
		if(x>=2)
			for(y=2;y<(n2+1);y++) 
				table2[x][y] = myprod(n2,x,y); 
	}
}

double mypolya(k,i,nt,n) 
int k,i,nt,n;
{
	if(k==0 && i==0)
		return 1.0;
	else 
		if(k==nt && i==n)
			return 1.0;
		else
			if(nt<=n && k>0 && k<nt && i>=k && i<=(k+n-nt))
				return nchoosek[n-i-1][nt-k-1]*nchoosek[i-1][k-1]/nchoosek[n-1][nt-1];
			else 
				return 0.0;
}

void MakeCubes(n1,n2)
int n1,n2;
{
	int x,y,z,nt,i,k;

	for(x=0;x<=n1;x++) 
		for(y=0;y<=n1;y++)  
			for(z=0;z<=n1;z++)
				cube1[(n1+1)*(n1+1)*x+(n1+1)*y+z] = 0.0;
	for(nt=1;nt<=n1;nt++) 
		for(k=0;k<=nt;k++) 
			for(i=k;i<=(n1-nt+k);i++) 
				cube1[(n1+1)*(n1+1)*nt+(n1+1)*k+i] = mypolya(k,i,nt,n1); 
	for(x=0;x<=n2;x++) 
		for(y=0;y<=n2;y++)  
			for(z=0;z<=n2;z++)
				cube2[(n2+1)*(n2+1)*x+(n2+1)*y+z] = 0.0;
	for(nt=1;nt<=n2;nt++) 
		for(k=0;k<=nt;k++) 
			for(i=k;i<=(n2-nt+k);i++) 
				cube2[(n2+1)*(n2+1)*nt+(n2+1)*k+i] = mypolya(k,i,nt,n2); 
}

double P1(n1,nt1,t1) 
int n1,nt1;
double t1;
{
	int i;
	double sum;

	sum = 0.0; 
	if(nt1 == 1) {
		for(i=2;i<=n1;i++)
			sum += exp(-nchoosek[i][2]*t1)/table1[2][i];
		return (1.0 - sum);
	} else {
		for(i=nt1;i<=n1;i++)
			sum += nchoosek[i][2]*exp(-nchoosek[i][2]*t1)/table1[nt1][i];
			return (sum/nchoosek[nt1][2]);
	}
}

double P2(n2,nt2,t2) 
int n2,nt2;
double t2;
{
	int i;
	double sum;

	sum = 0.0; 
	if(nt2 == 1) {
		for(i=2;i<=n2;i++)
			sum += exp(-nchoosek[i][2]*t2)/table2[2][i];
		return (1.0 - sum);
	} else {
		for(i=nt2;i<=n2;i++)
			sum += nchoosek[i][2]*exp(-nchoosek[i][2]*t2)/table2[nt2][i];
			return (sum/nchoosek[nt2][2]);
	}
}

double a1(n) 
int n;
{
	int i;
	double sum;

	sum = 0.0; 
	if(n > 1)
		for(i=1;i<n;i++) 
			sum += 1.0/((double)i);
	return sum;
}

void CalculateAll(n1,n2,theta1,theta2,thetaA,tau) 
int n1,n2;
double theta1,theta2,thetaA,tau;
{
	int i,n1t,n2t;
	double t1,t2,sumleft1,sumleft2,sum1,sum2,sumf,sums,prob,fix,ESx1,ESx2,ESs,ESf;

	t1 = tau/theta1;
	t2 = tau/theta2;
	sumleft1 = 0.0; 
	for(i=2;i<=n1;i++) 
		sumleft1 += (1.0 - exp(-nchoosek[i][2]*t1))/(nchoosek[i][2]*table1[2][i]);
	sumleft2 = 0.0; 
	for(i=2;i<=n2;i++) 
		sumleft2 += (1.0 - exp(-nchoosek[i][2]*t2))/(nchoosek[i][2]*table2[2][i]);
	sum1 = 0.0;
	sum2 = 0.0;
	sums = 0.0;
	sumf = 0.0;
	for(n1t=1;n1t<=n1;n1t++) 
		for(n2t=1;n2t<=n2;n2t++) {
			prob = P1(n1,n1t,t1)*P2(n2,n2t,t2); 
			fix = (1.0/((double)n1t) + 1.0/((double)n2t))/bico(n1t+n2t,n1t); 
			sum1 += prob*( (a1(n1t+n2t) - fix - a1(n2t))*thetaA - a1(n1t)*theta1 ); 
			sum2 += prob*( (a1(n1t+n2t) - fix - a1(n1t))*thetaA - a1(n2t)*theta2 ); 
			sums += prob*(a1(n1t) + a1(n2t) + fix - a1(n1t+n2t)); 
			sumf += prob*fix; 
	}
	ESx1 = theta1*a1(n1) + sum1;
	ESx2 = theta2*a1(n2) + sum2;
	ESs = thetaA*sums;
	ESf = tau - theta1*sumleft1/2.0 - theta2*sumleft2/2.0 + thetaA*sumf;

	printf("%f\t%f\t%f\t%f\n\n",ESx1,ESx2,ESs,ESf); 
	fflush(stdout); 
}

/*  The next two items are from Numerical Recipes in C,   */

static int imaxarg1,imaxarg2;
#define IMAX(a,b) (imaxarg1=(a),imaxarg2=(b),(imaxarg1) > (imaxarg2) ? (imaxarg1) : (imaxarg2))

static int iminarg1,iminarg2;
#define IMIN(a,b) (iminarg1=(a),iminarg2=(b),(iminarg1) < (iminarg2) ? (iminarg1) : (iminarg2))

void CalculateFreqs(n1,n2,theta1,theta2,thetaA,tau) 
int n1,n2;
double theta1,theta2,thetaA,tau;
{
	int x,y,i1,i2,nt1,nt2,nt,k1,k2,k,k1min,k1max;
	double t1,t2;
	double pn1nt1n2nt2,pk1k2nt1nt2,pk1i1k2i2,*pn1nt1,*pn2nt2;
	double sumleft1,sumleft2,*sumx1,*sumx2,*sumarray,**sums,*mutfarray,**mutfreqs; 
	double Sx1,Sx2,Ss,Sf;

	mutfarray = (double *) malloc( (size_t) (n1+1)*(n2+1)*sizeof(double) );
	if( !mutfarray ) MyError("problem getting memory for mutfarray in CalculateFreqs");
	mutfreqs = (double **) malloc( (size_t) (n1+1)*sizeof(double *) );
	if( !mutfreqs ) MyError("problem getting memory for mutfreqs in CalculateFreqs");
	sumx1 = (double *) malloc( (size_t) (n1+1)*sizeof(double) );
	if( !sumx1 ) MyError("problem getting memory for sumarray in CalculateFreqs");
	sumx2 = (double *) malloc( (size_t) (n2+1)*sizeof(double) );
	if( !sumx2 ) MyError("problem getting memory for sumarray in CalculateFreqs");
	pn1nt1 = (double *) malloc( (size_t) (n1+1)*sizeof(double) );
	if( !pn1nt1 ) MyError("problem getting memory for sumarray in CalculateFreqs");
	pn2nt2 = (double *) malloc( (size_t) (n2+1)*sizeof(double) );
	if( !pn2nt2 ) MyError("problem getting memory for sumarray in CalculateFreqs");
	sumarray = (double *) malloc( (size_t) (n1+1)*(n2+1)*sizeof(double) );
	if( !sumarray ) MyError("problem getting memory for sumarray in CalculateFreqs");
	sums = (double **) malloc( (size_t) (n1+1)*sizeof(double *) );
	if( !sums ) MyError("problem getting memory for sums in CalculateFreqs");

	for(x=0;x<=n1;x++) {
		mutfreqs[x] = mutfarray + x*(n2+1); 
		sums[x] = sumarray + x*(n2+1);
		sumx1[x] = 0.0;
		for(y=0;y<=n2;y++) {
			mutfreqs[x][y] = 0.0;
			sums[x][y] = 0.0;
			if(x==0)
				sumx2[y] = 0.0;
		}
	}
	
	t1 = tau/theta1;
	t2 = tau/theta2;
	sumleft1 = 0.0; 
	for(i1=2;i1<=n1;i1++) 
		sumleft1 += (1.0 - exp(-nchoosek[i1][2]*t1))/(nchoosek[i1][2]*table1[2][i1]);
	sumleft2 = 0.0; 
	for(i2=2;i2<=n2;i2++) 
		sumleft2 += (1.0 - exp(-nchoosek[i2][2]*t2))/(nchoosek[i2][2]*table2[2][i2]);
	for(nt1=1;nt1<=n1;nt1++) { 
		pn1nt1[nt1] = P1(n1,nt1,t1);
		if(nt1>1)
			for(i1=1;i1<n1;i1++) 
				for(k1=1;k1<nt1;k1++)
					sumx1[i1] += pn1nt1[nt1]*cube1[(n1+1)*(n1+1)*nt1+(n1+1)*k1+i1]/((double)k1);
	}
	for(nt2=1;nt2<=n2;nt2++) {
		pn2nt2[nt2] = P2(n2,nt2,t2); 
		if(nt2>1)
			for(i2=1;i2<n2;i2++) 
				for(k2=1;k2<nt2;k2++)
					sumx2[i2] += pn2nt2[nt2]*cube2[(n2+1)*(n2+1)*nt2+(n2+1)*k2+i2]/((double)k2);
	}
	for(nt1=1;nt1<=n1;nt1++) 
		for(nt2=1;nt2<=n2;nt2++) {
			pn1nt1n2nt2 = pn1nt1[nt1]*pn2nt2[nt2]; 
			nt = nt1 + nt2;
			for(k=1;k<nt;k++) {
				k1min = IMAX(0,k-nt2);
				k1max = IMIN(k,nt1);
				for(k1=k1min;k1<=k1max;k1++) {
					k2 = k - k1;
					pk1k2nt1nt2 = nchoosek[nt1][k1]*nchoosek[nt2][k2]/nchoosek[nt][k];
					for(i1=k1;i1<=k1+n1-nt1;i1++) 
						for(i2=k2;i2<=k2+n2-nt2;i2++) {
						   	pk1i1k2i2 = cube1[(n1+1)*(n1+1)*nt1+(n1+1)*k1+i1]*cube2[(n2+1)*(n2+1)*nt2+(n2+1)*k2+i2];
							sums[i1][i2] += pn1nt1n2nt2*pk1k2nt1nt2*pk1i1k2i2/((double)k);
						}
				}
			}
		}

	for(i1=0;i1<=n1;i1++) 
		for(i2=0;i2<=n2;i2++) 
			mutfreqs[i1][i2] = thetaA*sums[i1][i2];
			
	for(i1=1;i1<n1;i1++) 
		mutfreqs[i1][0] += theta1/((double)i1) - theta1*sumx1[i1]; 
	for(i2=1;i2<n2;i2++) 
		mutfreqs[0][i2] += theta2/((double)i2) - theta2*sumx2[i2]; 
	mutfreqs[n1][0] += 0.5*(tau - theta1*sumleft1);
	mutfreqs[0][n2] += 0.5*(tau - theta2*sumleft2);
	
	putchar('\n');
	for(i1=0;i1<=n1;i1++) {
		for(i2=0;i2<=n2;i2++) 
			printf("%f\t",mutfreqs[i1][i2]);
		putchar('\n');
	}
	putchar('\n');
	fflush(stdout);
	
 	Sx1 = 0.0;
	for(x=1;x<n1;x++) {
		Sx1 += mutfreqs[x][0];
		Sx1 += mutfreqs[x][n2];
	}
	Sx2 = 0.0;
	for(x=1;x<n2;x++) {
		Sx2 += mutfreqs[0][x];
		Sx2 += mutfreqs[n1][x];
	}
	Ss = 0.0;
	for(x=1;x<n1;x++) 
		for(y=1;y<n2;y++) 
			Ss += mutfreqs[x][y];
	Sf = mutfreqs[n1][0] + mutfreqs[0][n2];
	printf("%f\t%f\t%f\t%f\n\n",Sx1,Sx2,Ss,Sf); 
	fflush(stdout);  


	free( (void *)pn1nt1 );
	free( (void *)pn2nt2 );
	free( (void *)sumx1 );
	free( (void *)sumx2 );
	free( (void *)sums );
	free( (void *)sumarray );
	free( (void *)mutfreqs );
	free( (void *)mutfarray );
} 

int main()
{ 
	int n1,n2; 
	double theta1,theta2,thetaA,tau; 
	time_t start,finish;

	n1 = 10;
	n2 = 6;
	theta1 = 1.0;
	theta2 = 1.0;
	thetaA = 1.0;
	tau = 0.5;

	AllocateGlobals(n1,n2);
	MakeTables(n1,n2); 
	MakeCubes(n1,n2); 

	time(&start);
	CalculateFreqs(n1,n2,theta1,theta2,thetaA,tau); 
	time(&finish);
	/* printf("\n%ld\n\n",(long int)(finish-start)); 
	fflush(stdout); */
	
	CalculateAll(n1,n2,theta1,theta2,thetaA,tau);

	FreeGlobals();
	
	return 0;
}