#include "MeanR0CI.h"

/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
//function declarations:
void ReadRiDist(int Npre, int Ntot, char *DataFile, double **RiDist, double **RiCumDist);
double CalcMeanR(int Npre, int Ntot, double **RiDist);
double SampleMeanR(int Npre, int Ntot, double **RiCumDist);
void sort(int n, double *arr);


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

long seed, rcount;
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
int main(int argc, char *argv[])
{
	char DataFile[255], outFile[255];
	FILE *out;
	int Npre, Ntot;
	int SampleSize;
	double **RiDist, **RiCumDist;
	double meanR, CImin, CImax, *sampledMeanR;
	int i;

	printf("MeanR0CI: main started.\n");

	if(argc<6) 
		ERR_CRITICAL("Syntax:\n<No of IPs pre control> <No of IPs> <Data File> <SampleSize> <seed>\n");
	sscanf_s(argv[1], "%d", &Npre);
	sscanf_s(argv[2], "%d", &Ntot);
	sscanf_s(argv[3], "%s", DataFile, sizeof(DataFile));
	sscanf_s(argv[4], "%d", &SampleSize);
	sscanf_s(argv[5], "%d", &seed);

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

	if( (RiDist=dmatrix(0, Npre, 0, Ntot+1))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for RiDist.\n");
	if( (RiCumDist=dmatrix(0, Npre, 0, Ntot+1))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for RiCumDist.\n");
	if( (sampledMeanR=(double *) calloc(SampleSize, sizeof(double)))==0 )
		ERR_CRITICAL("main: Unable to allocate memory for sampledMeanR.\n");

	ReadRiDist(Npre, Ntot, DataFile, RiDist, RiCumDist);

	meanR=CalcMeanR(Npre, Ntot, RiDist);
	for(i=0;i<SampleSize;i++) {
		sampledMeanR[i]=SampleMeanR(Npre, Ntot, RiCumDist);
	}
	sort(SampleSize, sampledMeanR);

	/*if(!(out=fopen("sampledMeanR.txt", "w")))
		ERR_CRITICAL("main: unable to open file for output.\n");
    //printf("List of sampled mean R:\n");
	for(i=0;i<SampleSize;i++) fprintf(out, "%lg\n", sampledMeanR[i]);
	fclose(out);
	*/

	CImin=sampledMeanR[(int) (SampleSize*(1.0-CI_WIDTH)/2.0+0.5)];
	CImax=sampledMeanR[(int) (SampleSize*(1.0+CI_WIDTH)/2.0+0.5)];

	sprintf_s(outFile, sizeof(outFile), "%s.meanCI.txt", DataFile);
	if(!(fopen_s(&out, outFile, "w")))
		ERR_CRITICAL("main: Unable to open file for output.\n");
	fprintf(out, "%lg (%lg - %lg)\n", meanR, CImin, CImax);
	fclose(out);
	printf("mean (95 percent CI):\n");
	printf("%lg (%lg - %lg)\n", meanR, CImin, CImax);

	free_dmatrix(RiDist, 0, Npre, 0, Ntot+1);
	free_dmatrix(RiCumDist, 0, Npre, 0, Ntot+1);
	free(sampledMeanR);

	return 0;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
void ReadRiDist(int Npre, int Ntot, char *DataFile, double **RiDist, double **RiCumDist)
{
	FILE *in;
	int i,j;
	double temp[5];

	if( (fopen_s(&in, DataFile, "r"))==0 )
		ERR_CRITICAL("ReadRiDist: Unable to open file for input.\n");

	for(i=0;i<Npre;i++) {
		for(j=0;j<5;j++) {
			if(fscanf_s(in, "%lg", &temp[j])<0)
				ERR_CRITICAL("ReadRiDist: input file does not contain enough data.\n");
		}
		for(j=0;j<Ntot+1;j++) {
			if(fscanf_s(in, "%lg", &RiDist[i][j])<0)
				ERR_CRITICAL("ReadRiDist: input file does not contain enough data.\n");
		}
	}

	fclose(in);

	//calculate cumulative distribution:
	for(i=0;i<Npre;i++) {
		RiCumDist[i][0]=RiDist[i][0];
		for(j=1;j<Ntot+1;j++) RiCumDist[i][j]=RiCumDist[i][j-1]+RiDist[i][j];
		if(RiCumDist[i][Ntot]!=1.0)
			fprintf(stderr, "ReadRiDist: RiDist[i=%d] sums up to %lg!=1\n", 
			i, RiCumDist[i][Ntot]);
	}
	

	printf("ReadRiDist: successfully read Ri distributions.\n");
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
double CalcMeanR(int Npre, int Ntot, double **RiDist)
{
	int i,j;
	double meanR=0.0;

	for(i=0;i<Npre;i++) {
		for(j=1;j<=Ntot;j++) meanR+=j*RiDist[i][j];
	}
	return meanR/Npre;
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
double SampleMeanR(int Npre, int Ntot, double **RiCumDist)
{
	int i, Ri;
	double p, mean=0.0;

	for(i=0;i<Npre;i++) {
		p=ran3(&seed);
		for(Ri=0;Ri<=Ntot;Ri++)
			if(p<RiCumDist[i][Ri]) break;
		mean+=(double) Ri;
	}
	return (mean/Npre);
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/
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);
}
/*-------------------------------------------------------------------------------------------*/
/*-------------------------------------------------------------------------------------------*/