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


/* prototypes */
void find_constrained_maxes(int K, double *p, double *lambda,
							int N, int *z);
							
/************************************************************** 	
	FUNCTION: main
	
	PURPOSE: Take in data from the user and return the complete or 
			 partial sample MLE's of allele freq and lambda, depending
			 on whether the user knows M or not.  
		
	AUTHOR:  Eric C. Anderson  (eriq@u.washington.edu  till August 2001
								eric.anderson@stanfordalumni.org after that )
							
	DATE: 4 JANUARY 2000
	
**************************************************************/



int main(void)
{
	int i,j,N,M,z[1000];
	double *CS_MLE_p,CS_MLE_lambda;  /* complete sample MLE's  */
	double *PS_MLE_p,PS_MLE_lambda;  /* partial sample MLE's */
	int M_known;
	int NA;  /* number of alleles  */
	char junk;
	
	printf("\nEnter the total number of sampled hosts, M.");
	printf("\nIf you only know the number of infected hosts, enter zero\n->");
	scanf("%d",&M_known);
	
	
	if(!M_known) {
		printf("\n\nAll right then.  Enter the number of infected hosts\n->",NA);
		scanf("%d",&N);
	}
	else {
		M = M_known;
	}
	
	printf("\n\nEnter the number of alleles\n->");
	scanf("%d",&NA);
	
	/* allocate memory:  */
	CS_MLE_p = (double *)malloc(NA * sizeof(double));	
	PS_MLE_p = (double *)malloc(NA * sizeof(double));	
		
	printf("\n");	
	for(j=0;j<NA;j++)  {  /* collect the data from the user */
		if(!M_known)
			printf("\nEnter z_j (< %d) for allele %d  -->  ",N,j+1);
		else
			printf("\nEnter z_j (< %d) for allele %d  -->  ",M,j+1);
			
		scanf("%d",&z[j]);   /* initialize to zero to count the number of ticks infected
					 with the different alleles  */
					
		if(z[j] > N && !M_known) {
			printf("\n\nSorry!  That is > %d, the number of infected hosts.\nQuitting!",N);
			return(0);
		}
		else if(M_known && z[j] >= M) {
			printf("\n\nSorry!  That is > %d, the number of hosts sampled.\nQuitting!",M);
			return(0);
		}
	}
	
	if(M_known) {
		CS_MLE_lambda = 0.0;
		for(j=0;j<NA;j++)  {  /* compute the unnormalized p's  */
			CS_MLE_p[j] = log(  (double)(M - z[j]) / (double)M);
			CS_MLE_lambda += CS_MLE_p[j];
		}
		for(j=0;j<NA;j++)  {  /* then normalize them  */
			CS_MLE_p[j] /= CS_MLE_lambda;
		}
		CS_MLE_lambda = -CS_MLE_lambda;  /* then fix the sign on lambda.  */
	}
	else {  /* if M not known, find the partial sample mle's  */
		find_constrained_maxes(NA,PS_MLE_p,&PS_MLE_lambda,N,z);
	}
				
	/*  now print out all those goodies  */
	if(M_known)  {
		printf("\n\nHere are the results");
		printf("\n\n\t\t\tCounts\t\tComplete_Sample_MLE's");
		for(i=0;i<NA;i++)  {
			printf("\nAllele %d:\t\t%d\t\t\t%.4f",i+1,z[i],CS_MLE_p[i]);
		}
		printf("\n\nLambda\t\t\t-----\t\t%.4f",CS_MLE_lambda);
	}
	else {
		printf("\n\nHere are the results");
		printf("\n\n\t\t\tCounts\t\tPartial_Sample_MLE's");
		for(i=0;i<NA;i++)  {
			printf("\nAllele %d:\t\t%d\t\t\t%.4f",i+1,z[i],PS_MLE_p[i]);
		}
		printf("\n\nLambda\t\t\t-----\t\t%.4f",PS_MLE_lambda);
	}
	
	free(CS_MLE_p);
	free(PS_MLE_p);
	printf("\nPress Enter to terminate program\n");
	junk=getchar();
	junk=getchar();
	return(0);
}


/************************************************************* 	
	FUNCTION: find_constrained_maxes
	
	PURPOSE: This is a little function to iteratively find the 
			 partial sample MLE's.
	

	int K	-->  number of alleles
	double *p  -->  array to hold the MLE's of p
	double *lambda  -->  output parameter for the value of lambda
	int N  -->  number of infected hosts
	int *z -->  array of the counts of hosts infected by different types.
	
	AUTHOR:  Eric C. Anderson  (eriq@u.washington.edu  till August 2001
								eric.anderson@stanfordalumni.org after that )
							
	DATE: 20 NOVEMBER 2000
	REVISED: 5 JANUARY 2001
	
	Note that if any of the z_j's is equal to N, the partial sample mle is
	undefined.  So, if that is the case, this thing just exits.
	
**************************************************************/

#define PRECIS  .0001  /* the precision level one wants */

void find_constrained_maxes(int K, double *p, double *lambda,
							int N, int *z)
{
	int i;
	double new_lam,old_lam=0.0;
	double normo;
	double tempN;
	int sum_zs = 0;
	
	/* start by computing Bruce's estimates  */
	normo = 0.0;
	for(i=0;i<K;i++)  {
		if(z[i] >= N)  {
			printf("\n\nError!! z[%d] = %d which is >= N = %d in function",i,z[i],N);
			printf("\nfind_constrained_maxes().  ...Exiting to System...\n\n");
			exit(1);
		}
		
		p[i] = -log( (double)(N-z[i]) / ((double) N) );
		normo += p[i]; 
		
		sum_zs += z[i];
	}
	
	/*  if sum_zs is less than N, we have to abort!!  */
	if(sum_zs < N)  {
		printf("\n\nError!  The sum of the z_j's is less than N, the");
		printf("\nnumber of infected hosts.  This can't be.  Please review");
		printf("\nyour data and try again.   ....Exiting to System....\n\n");
		exit(1);
	}
	
	new_lam = normo;
	for(i=0;i<K;i++)  {
		p[i] /= new_lam;
	}
	
	/* now enter a loop that goes till we get the desired precision */
	while(fabs(old_lam - new_lam) > PRECIS)  {
		old_lam = new_lam;
		
		normo = 0.0;
		tempN = (double)N / (1.0 - exp(-old_lam) );
		for(i=0;i<K;i++)  {
			p[i] = -log( (tempN-(double)z[i]) / (double) tempN );
			normo += p[i]; 
		}
		new_lam = normo;
		for(i=0;i<K;i++)  {
			p[i] /= new_lam;
		}
	}
	*lambda = new_lam;
}