#include "MaximumLikelihoodEstimation.h"

/*-------------------------------------------------------------------------------------------*/
//Function declarations:
void ReadDates(int noDates, char *DatesFile, int *dates);
int ReadParam(char *ParamFile, param *GP);
double neglogLikelihood(int noDates, int *dates, int *pfail, workparam Pi);
	double PDF(double x, double kappa, double eta);
		double gammaPDF(double x, double alpha, double beta);
		double weibullPDF(double x, double kappa, double eta);
double MLE(char *InputFile , int noDates, int *dates, int ndim, int DoFit, param *GP);
	void ConstructLHTable(int num, int paramnum, double **SamplePoints);
	void FindMin(int i, int ndim, int noDates, int *dates, param *GP, workparam P[], double L[],  
				 workparam Pres[], double Lres[], int itres[]);
	int amoeba(int noDates, int *dates, workparam P[], double L[], double ftol, int ndim, 
				double (*funk)(int, int *, int*, workparam), int *nfunk);
		double amotry(int noDates, int *dates, workparam P[], double L[], workparam psum, 
			int ndim, double (*funk)(int, int *, int *, workparam), int worst, double fac);
	void OutputRes(char *OutputFile, int ndim, int DoFit, workparam Pres[], double Lres[], 
				   int itres[]);
int CalcCIs(char *InputFile, int noDates, int *dates, int ndim, param *GP, double optL);


/*-------------------------------------------------------------------------------------------*/
//Global variables:
double **matrix;
int DIST_SWITCH;	//0 for gamma distribution, 1 for weibull distribution, 
					//2 for negative binomial distribution.
					//currently, only 1=weibull works.

long seed, rcount;

/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int main(int argc, char *argv[]) {
	int noDates;
	char InputFile[255], ParamFile[255];
	int *dates;
	param *GP;
	int DoFit, ndim;
	double neg2lnL;
	int i, fail;

	if(argc<7) 
		ERR_CRITICAL("Syntax:\n<DIST_SWITCH> <No of dates> <Dates File> <Param File> <DoFit> <seed>\n");
	sscanf(argv[1],"%d", &DIST_SWITCH); 
	sscanf(argv[2],"%d", &noDates); 
	sscanf(argv[3],"%s", InputFile); 
	sscanf(argv[4],"%s", ParamFile);
	sscanf(argv[5],"%d", &DoFit);
	sscanf(argv[6],"%d", &seed); 

	if(seed>0) seed=-seed;
	if(seed==0) {
		seed=(long)time(NULL);
		if(seed>0) seed=-seed;
	}

	if(noDates<=0) 
		ERR_CRITICAL("positive number of dates necessary.\n");

	if(!(dates=(int *) calloc(noDates+1, sizeof(int))))
		ERR_CRITICAL("Unable to allocate memory for list of dates.\n");
	if(!(GP=(param *) malloc(sizeof(param))))
		ERR_CRITICAL("Unable to allcoate memory for parameters GP.\n");
	if(!(matrix=dmatrix(0, noDates, 0, noDates)))
		ERR_CRITICAL("main: Unable to allocate memory for matrix.\n");


	ReadDates(noDates, InputFile, dates);
	ndim=ReadParam(ParamFile, GP);

	neg2lnL=MLE(InputFile , noDates, dates, ndim, DoFit, GP);
	fail=CalcCIs(InputFile, noDates, dates, ndim, GP, neg2lnL);
	for(i=0;i<DoFit && fail==1; i++) {
		printf("recalc # %d.\n", i);
		neg2lnL=MLE(InputFile, noDates, dates, ndim, 0, GP);
		fail=CalcCIs(InputFile, noDates, dates, ndim, GP, neg2lnL);
	}





	free(dates);
	free(GP);
	free_dmatrix(matrix, 0, noDates, 0, noDates);
	return 0;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void ReadDates(int noDates, char *DatesFile, int *dates)
{
	FILE *indat;
	char filename[255];
	int i;
	int eofind=1;

	sprintf(filename, "%s.txt", DatesFile);

	if(!(indat=fopen(filename, "r")))
		ERR_CRITICAL("ReadDates: Unable to open file for input.\n");


	for(i=0;i<noDates && eofind>=0;i++){
		eofind=fscanf(indat, "%d", &dates[i]);
	}

	if(eofind<0) 
		ERR_CRITICAL("ReadDates: input file does not contain enough dates.\n");

	fclose(indat);
	printf("Successfully read list of dates.\n");

}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int ReadParam(char *ParamFile, param *GP){
	FILE *indat;
	char filename[255];
	int j;
	int scan=0;

	sprintf(filename, "%s.txt", ParamFile);

	if(!(indat=fopen(filename, "r")))
		ERR_CRITICAL("ReadParam: Unable to open file for input.\n");

	for(j=0;j<NUM_PARAM;j++) {
		fscanf(indat, "%lg %lg %lg %lg %d", &GP->index.array[j].Value, &GP->index.array[j].Min,
			&GP->index.array[j].Max, &GP->index.array[j].Scale, &GP->index.array[j].Scan);
		if(GP->index.array[j].Scan==1 || GP->index.array[j].Scan==2) scan++;
	}

	fclose(indat);
	printf("Successfully read Params.\n");
	return scan;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
double neglogLikelihood(int noDates, int *dates, int *pfail, workparam Pi) 
{
	int i,j,jmin;
	int timediffij;
	double lnL, Lsum;

	*pfail=0;

	//Calculate probability matrix:
	for(i=0;i<noDates;i++) for(j=0;j<noDates;j++) {
		timediffij=dates[j]-dates[i];
		if(timediffij>0 /*&& timediffij<=30*/ && i!=j){ //i can infect j if i was earlier, 
													//but less than 30 days before, 
													//and i cannot infect itself.
			matrix[i][j] = PDF(timediffij, Pi.index.named.kappa, Pi.index.named.eta);
		}
		else { matrix[i][j]=0; }
	}

	jmin=1;
	while(dates[jmin]==dates[jmin-1]) jmin++;

	lnL=0.0;
	for(j=jmin;j<noDates;j++) {
		Lsum=0.0;
		for(i=0;i<noDates;i++) Lsum+=matrix[i][j];
		if(Lsum==0.0) *pfail=1;
		lnL+=log(Lsum);
	}

	return -2.0*lnL;
}
/*-------------------------------------------------------------------------------------------*/
double PDF(double x, double kappa, double eta)
{
	switch(DIST_SWITCH) {
		case 1: return weibullPDF(x, kappa, eta);
			break;
		default:
			ERR_CRITICAL("Invalid DIST_SWITCH: 0 for gamma, 1 for weibull.\n");
		case 0: return gammaPDF(x, kappa, eta);
			break;
	}
}
/*-------------------------------------------------------------------------------------------*/
double gammaPDF(double x, double alpha, double beta)
//doesn't work anymore: neet to implement the cumulative distribution. 
{
	//mean = alpha * beta;
	//variance = alpha * beta * beta;
	return pow(x, alpha-1)*exp(-x/beta)/(exp(gammln(alpha))*pow(beta, alpha));
}

/*-------------------------------------------------------------------------------------------*/
double weibullPDF(double x, double kappa, double eta)
{ 
	//mean = 1/eta*exp(gammaln(1+1/kappa));
	//variance = 1/(eta*eta)*(exp(gammaln(1+2/kappa))-pow(exp(gammaln(1+1/kappa)), 2));
	if(x<=0.5) return 1.0 - exp(-pow(eta*(x+0.5), kappa));
	else return exp(-pow(eta*(x-0.5), kappa)) - exp(-pow(eta*(x+0.5), kappa));
	//return kappa*pow(eta, kappa)*pow(x, kappa-1)*exp(-pow(x*eta, kappa));
}
/*-------------------------------------------------------------------------------------------*/
/*double negBinPDF(int x, int r, double p)
//doesn't work !
{
	//mean = r*(1-p)/p;
	//variance = r*(1-p)/(p*p);
	return noverk(x+r-1, r-1)*pow(p, r)*pow(1-p, x);
}
/*-------------------------------------------------------------------------------------------*/
/*int noverk(int n, int k)
{
	int prod, i;

	if(n<0 || k<0) 
		ERR_CRITICAL("noverk: n<0 or k<0.\n");
	if(k>n) 
		ERR_CRITICAL("noverk: k>n.\n");
	if(k==0 || k==n) return 1;
	if(k>n/2.0) k=n-k;

	prod=n;
	for(i=n-1;i>n-k;i--) prod*=i;
	for(i=2;i<=k;i++) prod/=i;

	return prod;	
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
double MLE(char *InputFile , int noDates, int *dates, int ndim, int DoFit, param *GP)
{
	workparam *P, *Pres;
	double bestL, *L, *Lres;
	int *itres;
	double **SamplePoints;
	int loop;
	int i,j,k,l;



	if(!(P=(workparam *) calloc(ndim+2, sizeof(workparam))))
		ERR_CRITICAL("MLE: Unable to allocate memory for workparam P.\n");
	if(!(L=(double *) calloc(ndim+2, sizeof(double))))
		ERR_CRITICAL("MLE: Unable to allocate memory for L.\n");
	if(!(Pres=(workparam *) calloc(DoFit+2, sizeof(workparam))))
		ERR_CRITICAL("MLE: Unable to allocate memory for workparam Pres.\n");
	if(!(Lres=(double *) calloc(DoFit+2, sizeof(double))))
		ERR_CRITICAL("MLE: Unable to allocate memory for Lres.\n");
	if(!(itres=(int *) calloc(DoFit+2, sizeof(int))))
		ERR_CRITICAL("MLE: Unable to allocate memory for itres.\n");

	if(DoFit>0) {
		if(!(SamplePoints=dmatrix(0,DoFit+1,0,ndim+1)))
			ERR_CRITICAL("MLE: Unable to allocate memory for SamplePoints.\n");
		ConstructLHTable(DoFit, ndim, SamplePoints);
	
		loop=DoFit;

		for(i=0;i<loop;i++)
		{
			if(DoFit) printf("Draw # %d.\n",i);
			for(k=0, j=0;((DoFit>0)&&(j<ndim));j++)
			{
				if(GP->index.array[j].Scan==1) {
					P[0].index.array[k] = GP->index.array[j].Min 
						+ SamplePoints[i][k]
							* (GP->index.array[j].Max-GP->index.array[j].Min);
					for(l=1;l<ndim+1;l++) P[l].index.array[k]=P[0].index.array[k];
					P[k+1].index.array[k]+= GP->index.array[j].Scale;
					k++;
				}
				else if(GP->index.array[j].Scan==2) {
					if(GP->index.array[j].Max<0.0) {
						P[0].index.array[k] = -exp(log(-GP->index.array[j].Min)
							+ (log(-GP->index.array[j].Max)-log(GP->index.array[i].Min))
								*SamplePoints[i][k]);
						for(l=1;l<ndim+1;l++) P[l].index.array[k]=P[0].index.array[k];
						P[k+1].index.array[k]+= GP->index.array[j].Scale;
					}
					else {
						P[0].index.array[k] = exp(log(GP->index.array[j].Min)
						+ (log(GP->index.array[j].Max)-log(GP->index.array[j].Min))
							*SamplePoints[i][k]);
						for(l=1;l<ndim+1;l++) P[l].index.array[k]=P[0].index.array[k];
						P[k+1].index.array[k]+= GP->index.array[j].Scale;
					}
					k++;
				}
			}

			FindMin(i, ndim, noDates, dates, GP, P, L, Pres, Lres, itres);
		}
		OutputRes(InputFile, ndim, DoFit, Pres, Lres, itres);

		//write the MLestimates into GP:
		for(k=0, j=0;j<ndim;k++) {
			if(GP->index.array[k].Scan>0) {
				GP->index.array[k].Value=Pres[0].index.array[j];
				j++;
			}
		}
	
		bestL=Lres[0];
		printf("ML estimation completed.\n");

		free_dmatrix(SamplePoints, 0, DoFit+1, 0, ndim+1);
		free(Pres);
		free(Lres);
		free(itres);
	}

	else if (DoFit==0) {
		for(k=0, j=0;j<ndim;j++)
		{
			if(GP->index.array[j].Scan>0) {
				P[0].index.array[k] = GP->index.array[j].Value; 
				for(l=1;l<ndim+1;l++) P[l].index.array[k]=P[0].index.array[k];
				P[k+1].index.array[k]+= GP->index.array[j].Scale;
				k++;
			}
		}
		FindMin(0, ndim, noDates, dates, GP, P, L, Pres, Lres, itres);
		//write the MLestimates into GP:
		for(k=0, j=0;j<ndim;k++) {
			if(GP->index.array[k].Scan>0) {
				GP->index.array[k].Value=Pres[0].index.array[j];
				j++;
			}
		}
		bestL=Lres[0];


	}

	free(P);
	free(L);
	return bestL;
}
/*-------------------------------------------------------------------------------------------*/
void ConstructLHTable(int num, int paramnum, double **SamplePoints)
{
	int i,j,k,l;
	double r, f,ns,s;
	static FILE *dat;
	
	ns=(double) num;
	s=1.0/ns;	
	for(i=0;i<paramnum;i++)
	{	for(j=0;j<num;j++)
		{
			r=((double) j)*s;
			SamplePoints[j][i]=r+ran3(&seed)*s;
		}
		for(l=0;l<1000*num;l++)		//mixing up the columns
		{	j= (int) (ran3(&seed)*ns);	//j and k should both be <num, but non-negative
			k= (int) (ran3(&seed)*ns);
			f=SamplePoints[j][i];
			SamplePoints[j][i]=SamplePoints[k][i];
			SamplePoints[k][i]=f; 
		}
	}
}
/*-------------------------------------------------------------------------------------------*/
void FindMin(int i, int ndim, int noDates, int *dates, param *GP, workparam P[], double L[],  
			 workparam Pres[], double Lres[], int itres[])
{
	int j,k;
	int nfunk, it=0;
	int best;
	int fail, failsum=0;

	for(j=0;j<ndim+1;j++) {
		L[j]=neglogLikelihood(noDates, dates, &fail, P[j]);
		failsum+=fail;
	}

	if(failsum<ndim+1) {
		amoeba(noDates, dates, P, L, TINY, ndim, &neglogLikelihood, &nfunk);
		it=nfunk;
		//Re-initialising the lower points to restart the amoeba:
		for(best=0, j=1;j<ndim+1;j++) if(L[j]<L[best]) best=j;
		if(best!=0) {
			for(j=0;j<ndim;j++) P[0].index.array[j]=P[best].index.array[j];
			for(k=1;k<ndim+1;k++) {
				for(j=0;j<ndim;j++) P[k].index.array[j]=P[0].index.array[j];
				P[k].index.array[k-1]+=GP->index.array[k-1].Scale;
			}
		}
		amoeba(noDates, dates, P, L, TINY, ndim, &neglogLikelihood, &nfunk);
		it+=nfunk;	
	}

	best=0;
	for(j=1;j<ndim+1;j++) {
		if(L[j]<L[best]) best=j;
	}

	if(i==0) { 
		Lres[0]=L[best];
		for(j=0;j<ndim;j++) Pres[0].index.array[j]=P[best].index.array[j];
		itres[0]=it;
	}
	//if better than previous minimum, keep this one:
	else if(L[best]<Lres[0]) {
		Lres[0]=L[best];
		for(j=0;j<ndim;j++) Pres[0].index.array[j]=P[best].index.array[j];
		itres[0]=it;
	}

	//writing result into L/P/it res:
	for(j=0;j<ndim;j++) { 
		Pres[i+1].index.array[j]=P[best].index.array[j];
	}
	Lres[i+1]=L[best];
	itres[i+1]=it;
}
/*-------------------------------------------------------------------------------------------*/
int amoeba(int noDates, int *dates, workparam P[], double L[], double ftol, int ndim,
			double (*funk)(int, int *, int *, workparam), int *nfunk) 
{
	double amotry(int noDates, int *dates, workparam P[], double L[], workparam psum, 
		int ndim, double (*funk)(int, int *, int *, workparam), int worst, double fac);
	int i, worst, best, sworst, j, mpts=ndim+1;
	double rtol, sum, swap, Lsave, Ltry;
	workparam psum;
	int fail;

	*nfunk=0;
	GET_PSUM
	for(;;) {
		best=0;
		//First we must determine which point is the lowest(worst), next lowest and 
		//highest(best), by looping over the points in the simplex.
		if(L[0]>L[1])	{ sworst=1; worst=0; }
		else			{ sworst=0; worst=1; }
		for(i=0;i<mpts;i++) {
			if(L[i]<=L[best]) best=i;
			if(L[i]>L[worst]) {
				sworst=worst;
				worst=i;
			} else if (L[i]>L[sworst] && i!=worst) sworst=i;
		}
		rtol=2.0*fabs(L[best]-L[worst])/(fabs(L[worst])+fabs(L[best])+TINY);
		//Compute the fractional range from worst to best and return if satisfactory.
		if(rtol<ftol) {	//If returning, put best point and value in slot 1.
			SWAP(L[1],L[best])
			for(i=0;i<ndim;i++) SWAP(P[0].index.array[i], P[best].index.array[i])
			break;
		}
		if(*nfunk >=NMAX) {
			fprintf(stderr, "amoeba: NMAX exceeded\n");
			return 1;
		}
		*nfunk+=2;
		//Begin a new iteration. First extrapolate by a factor -1 through the face of the 
		//simplex across from the high point, i.e., reflect the simplex from the hight point.
		Ltry=amotry(noDates, dates, P, L, psum, ndim, funk, worst, 2.0);
		if(Ltry <=L[best])
			//Gives a result better than the best point, 
			//so try an additional extrapolation by a factor 2.
			Ltry=amotry(noDates, dates, P, L, psum, ndim, funk, worst, 2.0);
		else if (Ltry>=L[sworst]) {
			//The reflecte point is worse than the second worst, 
			//so look for an intermediate better point, i.e., do a one-dimensional contraction.
			Lsave=L[worst];
			Ltry=amotry(noDates, dates, P, L, psum, ndim, funk, worst, 0.5);
			if(Ltry>=Lsave) { //Can't seem to get rid of that bad point. Better 
				for(i=0;i<mpts;i++) {	//contract around the best point. 
					if(i!=best){
						for(j=0;j<ndim;j++)
							P[i].index.array[j]
								= psum.index.array[j]
								= 0.5* (P[i].index.array[j] + P[best].index.array[j]);
						L[i]=(*funk)(noDates, dates, &fail, psum);
					}
				}
				*nfunk += ndim;		//Keep track of function evaluations.
				GET_PSUM	//Recompute psum.
			}
		} else --(*nfunk);	//Correct the evaluation count.
	}	//Go back for the test of doneness and the next iteration.

	return 0;
}
/*-------------------------------------------------------------------------------------------*/
double amotry(int noDates, int *dates, workparam P[], double L[], workparam psum, int ndim,
			  double (*funk)(int, int *, int *, workparam), int worst, double fac)
{
	int j;
	double fac1, fac2, Ltry;
	workparam Ptry;
	int fail;

	fac1=(1.0-fac)/ndim;
	fac2=fac1-fac;
	for(j=0;j<ndim;j++) 
		Ptry.index.array[j] = psum.index.array[j]*fac1 - P[worst].index.array[j]*fac2;
	Ltry=(*funk)(noDates, dates, &fail, Ptry);	//Evaluate the function at the trial point. 
	if(Ltry <L[worst]) {	//If it's better than the worst, then replace the worst.
		L[worst]=Ltry;
		for(j=0;j<ndim;j++) {
			psum.index.array[j] += Ptry.index.array[j] - P[worst].index.array[j];
			P[worst].index.array[j] = Ptry.index.array[j];
		}
	}
	return Ltry;
}
/*-------------------------------------------------------------------------------------------*/
void OutputRes(char *OutputFile, int ndim, int DoFit, workparam Pres[], double Lres[], 
			   int itres[])
{
	FILE *out;
	char filename[255];
	int i,j;

	sprintf(filename, "%s.paramest.txt", OutputFile); 

	if(!(out=fopen(filename, "w")))
		ERR_CRITICAL("OutputRes: Unable to open file for output.\n");

	fprintf(out, "i\tno_it\t-2lnL\tkappa\teta\n");

	//range for i as i=0 stores the MLEs.
	for(i=0;i<=DoFit;i++) {
		fprintf(out, "%d\t%d\t%lg", i, itres[i], Lres[i]);
		for(j=0;j<ndim;j++) fprintf(out, "\t%lg", Pres[i].index.array[j]);
		fprintf(out, "\n");
	}

	fclose(out);
}

/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int CalcCIs(char *InputFile, int noDates, int *dates, int ndim, param *GP, double optL)
{
	FILE *out;
	char filename[255];
	int i,j,k,l,m;
	double xmin, xmax, dx;
	double **x, deltalnL[NO_X];
	workparam P, Pbest;
	double L, Lbest;
	int fail;

	Lbest=optL;
	if(!(x=dmatrix(0, ndim, 0, NO_X)))
		ERR_CRITICAL("CalcCIs: Unable to allocate memory for x.\n");


	for(k=0, j=0;(j<ndim && k<NUM_PARAM);k++) {
		//printf("Calculate profile likelihood for parameter %d. ", j);

		for(l=0,m=0; m<ndim; l++) {
			if(GP->index.array[l].Scan>0) {
				P.index.array[m]=GP->index.array[l].Value;
				for(i=0;i<NO_X;i++) x[m][i]=P.index.array[m];
				m++;
			}
		}

		switch(GP->index.array[k].Scan) {
			case 1: j++;
				xmin=GP->index.array[k].Min;
				xmax=GP->index.array[k].Max;
				dx=(xmax-xmin)/(NO_X-1);
				for(i=0;i<NO_X;i++) x[j-1][i]=xmin+i*dx;
				break;
			case 2: j++;
				xmin=GP->index.array[k].Min;
				xmax=GP->index.array[k].Max;
				dx=pow(xmax/xmin, 1.0/(NO_X-1));
				for(i=0;i<NO_X;i++) x[j-1][i]=xmin*pow(dx, i);
				break;
			default: break;
		}

		for(i=0;i<NO_X;i++) {
			P.index.array[j-1]=x[j-1][i];
			//printf("kappa = %lg\t eta=%lg\n", P.index.named.kappa, P.index.named.eta);
			L=neglogLikelihood(noDates, dates, &fail, P);
			if(L<Lbest-SMALL) {
				Lbest=L;
				for(m=0;m<ndim;m++) Pbest.index.array[m]=P.index.array[m];
			}
			deltalnL[i]=L-optL;
		}

		sprintf(filename, "%s.profL%d.txt", InputFile, j-1);
		if(!(out=fopen(filename, "w")))
			ERR_CRITICAL("CalcCIs: Unable to open file for output.\n");

		fprintf(out, "deltalnL\tkappa\teta\t(first line optimum)\n");
		fprintf(out, "%lg\t%lg\t%lg\n", 
			optL, GP->index.named.kappa.Value, GP->index.named.eta.Value);

		for(i=0;i<NO_X;i++) {
			fprintf(out, "%lg", deltalnL[i]);
			for(m=0;m<ndim;m++) fprintf(out, "\t%lg", x[m][i]);
			fprintf(out, "\n");
		}

		fclose(out);
		//printf("Completed.\n");
	}

	fail=0;
	if(Lbest<optL) {
		for(j=0,k=0; j<ndim && k<NUM_PARAM; k++) {
			if(GP->index.array[k].Scan==1 || GP->index.array[k].Scan==2) {
				GP->index.array[k].Value=Pbest.index.array[j];
				j++;
			}
		}
		fail=1;
	}
	//printf("CalcCIs: minimum -2lnL = %lg at kappa = %lg, eta = %lg\n", 
	//	Lbest, GP->index.named.kappa.Value, GP->index.named.eta.Value);

	free_dmatrix(x, 0, ndim, 0, NO_X);
	return fail;
}