#include "TotCI_Ri.h"

/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
//Function Declarations:

void ReadDates(int NoDates, char *DatesFile, int *Dates);
void ReadParams(int SampleGenTime, char *ParamFile, char *OptParamFile, double *kappa_opt, 
				double *eta_opt, double *kappa, double *eta);
void CalcRDistOpt(int NoDates, int *Dates, double kappa, double eta, double **matrix, 
			   double **RDist, double *R0list);
void CalcRDist(int index, int NoDates, int *Dates, double kappa, double eta, double **matrix, 
			   double **RDist, double **RSample);
	double weibullPDF(double x, double kappa, double eta);
void SampleRi(int index, int NoDates, double **RDist, double **RSampleTot);
int NDistOutput(char *OutputFile, int NSample, int NoDates, double **RSampleTot);
void CalcMeanR(int Npre, int NSample, double **RSampleTot, double *meanRSampleTot);
double CalcMeanROpt(int Npre, double *R0list);
void CalcMaxR(int NoDates, int NSample, double **RSampleTot, double *maxRSampleTot);
double CalcMaxROpt(int NoDates, double *R0list);
void sort(int n, double *arr);
void OutputMeanMax(char *OutputFile, double meanR, double maxR, double meanCImin, 
				double meanCImax, double maxCImin, double maxCImax);
void OutputData(int NoDates, char *OutputFile, int *Dates, double *R0list, double *CIminGT, 
				double *CImaxGT, double *CIminTot, double *CImaxTot);


/*-------------------------------------------------------------------------------------------*/
//Global variables:

long seed, rcount;
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int main(int argc, char *argv[])
{
	char DatesFile[255], ParamFile[255], OptParamFile[255], OutputFile[255];
	int NoDates, Npre, SampleGenTime, SampleR, NoEps;
	int *Dates;
	double *kappa, *eta, kappa_opt, eta_opt;
	double **matrix, **RDist;
	double *R0list, *CIminGT, *CImaxGT, *CIminTot, *CImaxTot;
	double **RSampleGT, **RSampleTot, *meanRSampleTot, *maxRSampleTot;
	double meanR, meanRCImin, meanRCImax, maxR, maxRCImin, maxRCImax;
	int i, j;

	printf("Hello from TotCI_Ri.\n");

	if(argc<10)
		ERR_CRITICAL("Syntax: <Npre> <NoDates> <DatesFile> <ParamFile> <OptParamFile> <OutputFile> <SampleGenTime> <SampleR> <seed>\n");
	sscanf_s(argv[1], "%d", &Npre);
	sscanf_s(argv[2], "%d", &NoDates);
	sscanf_s(argv[3], "%s", DatesFile, sizeof(DatesFile));
	sscanf_s(argv[4], "%s", ParamFile, sizeof(ParamFile));
	sscanf_s(argv[5], "%s", OptParamFile, sizeof(OptParamFile));
	sscanf_s(argv[6], "%s", OutputFile, sizeof(OutputFile));
	sscanf_s(argv[7], "%d", &SampleGenTime);
	sscanf_s(argv[8], "%d", &SampleR);
	sscanf_s(argv[9], "%d", &seed);

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

	if( (Dates=(int *) calloc(NoDates, sizeof(int)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for Dates.\n");
	if( (kappa=(double *) calloc(SampleGenTime, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for kappa.\n");
	if( (eta=(double *) calloc(SampleGenTime, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for eta.\n");
	if( (matrix=dmatrix(0, NoDates-1, 0, NoDates-1))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for matrix.\n");
	if( (RDist=dmatrix(0, NoDates-1, 0, NoDates))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for RDist.\n");
	if( (R0list=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for R0list.\n");
	if( (RSampleGT=dmatrix(0, NoDates-1, 0, SampleGenTime))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for RSampleGT.\n");
	if( (RSampleTot=dmatrix(0, NoDates, 0, SampleGenTime*SampleR))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for RSampleTot.\n");
	if( (meanRSampleTot=(double *) calloc(SampleGenTime*SampleR, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for meanRSampleTot.\n");
	if( (maxRSampleTot=(double *) calloc(SampleGenTime, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for maxRSampleTot.\n");
	if( (CIminGT=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for CIminGT.\n");
	if( (CImaxGT=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for CImaxGT.\n");
	if( (CIminTot=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for CIminTot.\n");
	if( (CImaxTot=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for CImaxTot.\n");



	ReadDates(NoDates, DatesFile, Dates);
	ReadParams(SampleGenTime, ParamFile, OptParamFile, &kappa_opt, &eta_opt, kappa, eta);

	CalcRDistOpt(NoDates, Dates, kappa_opt, eta_opt, matrix, RDist, R0list);
	meanR=CalcMeanROpt(Npre, R0list);
	maxR=CalcMaxROpt(NoDates, R0list);

	for(i=0;i<SampleGenTime;i++) {
		CalcRDist(i, NoDates, Dates, kappa[i], eta[i], matrix, RDist, RSampleGT);
		for(j=0;j<SampleR;j++) {
			SampleRi(i*SampleR+j, NoDates, RDist, RSampleTot); 
		}
	}
	NoEps=NDistOutput(OutputFile, SampleGenTime*SampleR, NoDates, RSampleTot);
	printf("%d of %d simulated epis are of size %d.\n", NoEps, SampleGenTime*SampleR, NoDates);

	CalcMeanR(Npre, SampleGenTime*SampleR, RSampleTot, meanRSampleTot);
	CalcMaxR(NoDates, SampleGenTime, RSampleGT, maxRSampleTot);

	for(i=0;i<NoDates;i++) {
		sort(SampleGenTime, RSampleGT[i]);
		CIminGT[i]=RSampleGT[i][(int) (SampleGenTime*(1.0-CI_WIDTH)/2.0+0.5)];
		CImaxGT[i]=RSampleGT[i][(int) (SampleGenTime*(1.0+CI_WIDTH)/2.0+0.5)];

		sort(SampleGenTime*SampleR, RSampleTot[i]);
		CIminTot[i]=RSampleTot[i][(int) (SampleGenTime*SampleR*(1.0-CI_WIDTH)/2.0+0.5)];
		CImaxTot[i]=RSampleTot[i][(int) (SampleGenTime*SampleR*(1.0+CI_WIDTH)/2.0+0.5)];

		sort(SampleGenTime*SampleR, meanRSampleTot);
		meanRCImin=meanRSampleTot[(int) (SampleGenTime*SampleR*(1.0-CI_WIDTH)/2.0+0.5)];
		meanRCImax=meanRSampleTot[(int) (SampleGenTime*SampleR*(1.0+CI_WIDTH)/2.0+0.5)];

		sort(SampleGenTime, maxRSampleTot);
		maxRCImin=maxRSampleTot[(int) (SampleGenTime*(1.0-CI_WIDTH)/2.0+0.5)];
		maxRCImax=maxRSampleTot[(int) (SampleGenTime*(1.0+CI_WIDTH)/2.0+0.5)];
	}

	//Output mean and max R + CIs:
	OutputMeanMax(OutputFile, meanR, maxR, meanRCImin, meanRCImax, maxRCImin, maxRCImax);
	OutputData(NoDates, OutputFile, Dates, R0list, CIminGT, CImaxGT, CIminTot, CImaxTot);


	free(Dates);
	free(kappa);
	free(eta);
	free_dmatrix(matrix, 0, NoDates-1, 0, NoDates-1);
	free_dmatrix(RDist, 0, NoDates-1, 0, NoDates);
	free(R0list);
	free_dmatrix(RSampleGT, 0, NoDates-1, 0, SampleGenTime);
	free_dmatrix(RSampleTot, 0, NoDates, 0, SampleGenTime*SampleR);
	free(meanRSampleTot);
	free(maxRSampleTot);
	free(CIminGT);
	free(CImaxGT);
	free(CIminTot);
	free(CImaxTot);

	printf("Good Bye.\n");
	return 0;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void ReadDates(int NoDates, char *DatesFile, int *Dates)
{
	FILE *indat;
	int i;
	int eofind=1;

	if( (fopen_s(&indat, DatesFile, "r"))!=0)
		ERR_CRITICAL("ReadDates: Unable to open file for input.\n");

	for(i=0;i<NoDates && eofind>=0;i++){
		eofind=fscanf_s(indat, "%d", &Dates[i]);
	}
	if(eofind<0) 
		ERR_CRITICAL("ReadDates: dates input file does not contain enough dates.\n");
	fclose(indat);

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

}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void ReadParams(int SampleGenTime, char *ParamFile, char *OptParamFile, double *kappa_opt, 
				double *eta_opt, double *kappa, double *eta)
{
	FILE *indat;
	int j, eofind=1;
	double temp;

	if( (fopen_s(&indat, ParamFile, "r"))!=0)
		ERR_CRITICAL("ReadParam: Unable to open ParamFile for input.\n");

	for(j=0;j<SampleGenTime && eofind>=0;j++) {
		eofind=fscanf_s(indat, "%lg %lg %lg", &temp, &kappa[j], &eta[j]);
	}
	if(eofind<0) 
		ERR_CRITICAL("ReadParams: not enough values of kappa and eta in ParamFile.\n");
	fclose(indat);

	if( (fopen_s(&indat, OptParamFile, "r"))!=0)
		ERR_CRITICAL("ReadParam: Unable to open OptParamFile for input.\n");
	fscanf_s(indat, "%lg %lg", kappa_opt, eta_opt);
	fclose(indat);


	printf("Successfully read Params.\n");
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void CalcRDistOpt(int NoDates, int *Dates, double kappa, double eta, double **matrix, 
			   double **RDist, double *R0list)
//Calculates the estimates of R for each date, 
//and returns the cumulative distributions for Ri in matrix.
{
	int i, j, r, timediffij;
	double *norm;

	if( (norm=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("CalcRDist: unable to allocate memory for norm.\n");

	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] = weibullPDF(timediffij, kappa, eta);
			}
			else { matrix[i][j]=0; }
		}
	}

	//Normalisation of matrix
	for(j=0;j<NoDates;j++) {
		norm[j]=0.0;
		for(i=0;i<NoDates;i++) norm[j]+=matrix[i][j];
	}
	for(j=0;j<NoDates;j++) {
		if(norm[j]>0) for(i=0;i<NoDates;i++) matrix[i][j]/=norm[j];
		else for(i=0;i<NoDates;i++) matrix[i][j]=0.0;	//not necessary if matrix was all 
														//non-negative. (as should have been)
	}
	//Adding up for R0:
	for(i=0;i<NoDates;i++) {
		R0list[i]=0.0;
		for(j=0;j<NoDates;j++) R0list[i]+=matrix[i][j];
	}
	

	//Calculating the convoluted Bernoulli distribution:
	for(i=0;i<NoDates;i++) {
		RDist[i][0]=1.0-matrix[i][0];
		RDist[i][1]=matrix[i][0];
		for(r=2;r<=NoDates;r++) RDist[i][r]=0.0;

		for(j=1;j<NoDates;j++) {
			for(r=j+1;r>0;r--) {
				RDist[i][r]=RDist[i][r]*(1.0-matrix[i][j])+RDist[i][r-1]*matrix[i][j];
			}
			RDist[i][0]=RDist[i][0]*(1.0-matrix[i][j]);
		}
	}
	//Convert RDist into the cumulative distribution:
	for(i=0;i<NoDates;i++) {
		for(j=1;j<=NoDates;j++) RDist[i][j]+=RDist[i][j-1];
	}

}
/*-------------------------------------------------------------------------------------------*/
void CalcRDist(int index, int NoDates, int *Dates, double kappa, double eta, double **matrix, 
			   double **RDist, double **RSampleGT)
//Calculates the estimates of R for each date, 
//and returns the cumulative distributions for Ri in matrix.
{
	int i, j, r, timediffij;
	double *norm;

	if( (norm=(double *) calloc(NoDates, sizeof(double)))==0 )
		ERR_CRITICAL("CalcRDist: unable to allocate memory for norm.\n");

	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] = weibullPDF(timediffij, kappa, eta);
			}
			else { matrix[i][j]=0; }
		}
	}

	//Normalisation of matrix
	for(j=0;j<NoDates;j++) {
		norm[j]=0.0;
		for(i=0;i<NoDates;i++) norm[j]+=matrix[i][j];
	}
	for(j=0;j<NoDates;j++) {
		if(norm[j]>0) for(i=0;i<NoDates;i++) matrix[i][j]/=norm[j];
		else for(i=0;i<NoDates;i++) matrix[i][j]=0;	//not necessary if matrix was all 
													//non-negative. (as should have been)
	}
	//Adding up for R0:
	for(i=0;i<NoDates;i++) {
		RSampleGT[i][index]=0;
		for(j=0;j<NoDates;j++) RSampleGT[i][index]+=matrix[i][j];
	}
	

	//Calculating the convoluted Bernoulli distribution:
	for(i=0;i<NoDates;i++) {
		RDist[i][0]=1.0-matrix[i][0];
		RDist[i][1]=matrix[i][0];
		for(r=2;r<=NoDates;r++) RDist[i][r]=0.0;

		for(j=1;j<NoDates;j++) {
			for(r=j+1;r>0;r--) {
				RDist[i][r]=RDist[i][r]*(1.0-matrix[i][j])+RDist[i][r-1]*matrix[i][j];
			}
			RDist[i][0]=RDist[i][0]*(1.0-matrix[i][j]);
		}
	}
	//Convert RDist into the cumulative distribution:
	for(i=0;i<NoDates;i++) {
		for(j=1;j<=NoDates;j++) RDist[i][j]+=RDist[i][j-1];
	}

	free(norm);
}
/*-------------------------------------------------------------------------------------------*/
double weibullPDF(double x, double kappa, double eta)
{ 
	double PDF;
	//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) PDF = 1.0 - exp(-pow(eta*(x+0.5), kappa));
	else PDF = exp(-pow(eta*(x-0.5), kappa)) - exp(-pow(eta*(x+0.5), kappa));
	return PDF;
	//return kappa*pow(eta, kappa)*pow(x, kappa-1)*exp(-pow(x*eta, kappa));
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void SampleRi(int index, int NoDates, double **RDist, double **RSampleTot) 
{
	double p;
	int i, R, N=0, j;

	printf("Sampling Epi no %d\r", index);

	for(j=0;j<1000 && N!=NoDates;j++) {	
	//for unconstrained epidemic size, comment the for-loop out (just keep the body)
		N=0;
		for(i=0;i<NoDates;i++) {
			p=ran3(&seed);
			for(R=0;R<=NoDates;R++) {
				if(p<RDist[i][R])	
					break;
			}
			RSampleTot[i][index]=(double) R;
			N+=R;
		}
	}
	RSampleTot[NoDates][index]=(double) N;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int NDistOutput(char *OutputFile, int NSample, int NoDates, double **RSampleTot)
{
	int *Dist;
	int i, N, NoCond=0;
	char filename[255];
	FILE *out;

	if( (Dist=(int *) calloc(10*NoDates+1, sizeof(int)))==0 )
		ERR_CRITICAL("NDistOutput: Unable to allocate memory for Dist.\n");
	for(i=0;i<10*NoDates+1;i++) Dist[i]=0;

	for(i=0;i<NSample;i++) {
		N=(int) RSampleTot[NoDates][i];
		if(N==NoDates) NoCond++;
		if(N<10*NoDates) Dist[N]++;
		else Dist[10*NoDates]++;
	}
	
	sprintf_s(filename, sizeof(filename), "NDist.%s", OutputFile);
	if( fopen_s(&out, filename, "w") !=0 )
		ERR_CRITICAL("NDistOutput: Unable to open file for output.\n");
	fprintf_s(out, "N\tNoEpis\n");
	for(i=0;i<10*NoDates+1;i++) fprintf_s(out, "%d\t%d\n", i, Dist[i]);
	fclose(out);

	free(Dist);
	return NoCond;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void CalcMeanR(int Npre, int NSample, double **RSampleTot, double *meanRSampleTot)
{
	int index, i;
	double meanR;

	for(index=0;index<NSample;index++) {
		meanR=0.0;
		for(i=0;i<Npre;i++) meanR+=RSampleTot[i][index];
		meanRSampleTot[index]=meanR/Npre;
	}
}
/*-------------------------------------------------------------------------------------------*/
double CalcMeanROpt(int Npre, double *R0list)
{
	int i;
	double meanR=0.0;

	for(i=0;i<Npre;i++) meanR+=R0list[i];
	return meanR/Npre;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void CalcMaxR(int NoDates, int NSample, double **RSampleTot, double *maxRSampleTot)
{
	int index, i;

	for(index=0;index<NSample;index++) {
		maxRSampleTot[index]=-1.0;
		for(i=0;i<NoDates;i++) 
			if(maxRSampleTot[index]<RSampleTot[i][index]) 
				maxRSampleTot[index]=RSampleTot[i][index];
	}
}
/*-------------------------------------------------------------------------------------------*/
double CalcMaxROpt(int NoDates, double *R0list)
{
	int i;
	double maxR=-1.0;

	for(i=0;i<NoDates;i++) if(maxR<R0list[i]) maxR=R0list[i];
	return maxR;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void sort(int n, double *arr)
// This is quicksort from the Numerical Recipes, but the array runs reasonably from 0 ot n-1, 
// unlike in the original where it runs from 1 to n.
{
	int i,ir=n,j,k,l=1,*istack;
	int jstack=-1;
	double a, temp;

	if( (istack=(int *) calloc(NSTACK, sizeof(int)))==0)
		ERR_CRITICAL("sort: Unable to allocate memory for istack.\n");

	for(;;) {
		if(ir-l < M) {
			for(j=l+1;j<=ir;j++) {
				a=arr[j-1];
				for(i=j-1;i>=l;i--) {
					if(arr[i-1] <= a) break;
					arr[i]=arr[i-1];
				}
				arr[i]=a;
			}
			if(jstack == -1) break;
			ir=istack[jstack--];
			l=istack[jstack--];
		} else {
			k=(l+ir) >> 1;
			SWAP(arr[k-1], arr[l])
			if(arr[l-1] > arr[ir-1]) {
				SWAP(arr[l-1], arr[ir-1])
			}
			if(arr[l] > arr[ir-1]) {
				SWAP(arr[l], arr[ir-1])
			}
			if(arr[l-1] > arr[l]) {
				SWAP(arr[l-1],arr[l])
			}
			i=l+1;
			j=ir;
			a=arr[l];
			for (;;) {
				do i++; while (arr[i-1]<a);
				do j--; while (arr[j-1]>a);
				if (j < i) break;
				SWAP(arr[i-1],arr[j-1]);
			}
			arr[l]=arr[j-1];
			arr[j-1]=a;
			jstack +=2;
			if(jstack > NSTACK-1) 
				ERR_CRITICAL("STACK too small in sort.\n");
			if(ir-i+1 >= j-l) {
				istack[jstack]=ir;
				istack[jstack-1]=i;
				ir=j-1;
			} else {
				istack[jstack]=j-1;
				istack[jstack-1]=l;
				l=i;
			}
		}
	}

	free(istack);
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void OutputMeanMax(char *OutputFile, double meanR, double maxR, double meanCImin, 
				double meanCImax, double maxCImin, double maxCImax)
{
	FILE *out;
	char filename[255];

	sprintf_s(filename, sizeof(filename), "MeanMaxR%s", OutputFile);
	if( (fopen_s(&out, filename, "w"))!=0 )
		ERR_CRITICAL("OutputMeanMax: Unable to open file for output.\n");

	fprintf(out, "meanR\tmeanR_CImin\tmeanR_CImax\n%lg\t%lg\t%lg\n", 
		meanR, meanCImin, meanCImax);
	fprintf(out, "maxR\tmaxR_CImin\tmaxR_CImax\n%lg\t%lg\t%lg\n", 
		maxR, maxCImin, maxCImax);
	fclose(out);

}
/*-------------------------------------------------------------------------------------------*/
void OutputData(int NoDates, char *OutputFile, int *Dates, double *R0list, double *CIminGT, 
				double *CImaxGT, double *CIminTot, double *CImaxTot)
{
	int i;
	FILE *out;

	if( (fopen_s(&out, OutputFile, "w"))!=0 )
		ERR_CRITICAL("OutputData: Unable to open file for output.\n");
	fprintf(out, "Date\tR\tCIminGT\tCImaxGT\tCIminTot\tCImaxTot\n");
	for(i=0;i<NoDates;i++) {
		fprintf(out, "%d\t%lg\t%lg\t%lg\t%lg\t%lg\n", 
			Dates[i], R0list[i], CIminGT[i], CImaxGT[i], CIminTot[i], CImaxTot[i]);
	}
	fclose(out);

	printf("OutputData: Successfully written the data.\n");
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/