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


/**********************/
/*    MCMC            */
/**********************/
int sampleZ (int K, double *pz_s, double *pxz_s)
{
	int i, Z_XPQ;
	double zran, p1, p2, psum, sump;

	if(K==1)							
		Z_XPQ = 0;
	else
	{
		sump = 0;
		for (i=0;i<K;i++) 
		{
			sump += pxz_s[i]*pz_s[i];
		}	
		psum=0;
		for (i=0;i<K;i++) 
		{
			psum += pxz_s[i]*pz_s[i]/sump;
		}
		p1 = 0;
		p2 = 0;
		zran = whrandom()*psum; 
		for (i=0;i<K;i++) 
		{							
			p2 += pxz_s[i]*pz_s[i]/sump;
			if((zran>p1) && (zran<=p2))
			{
				Z_XPQ = i;
				i = K;
			}
			p1 = p2;
		}
	}
	return (Z_XPQ);
}


/**********************************/
/*** 6. to calculate likelihood ***/
/**********************************/

void est_xz_fn(int N,double *a,double **ax,double alpha_o,int times,int trials, int n,int Nt,int tt,int K,int M,int aN,int *xM,int **fdata,double ***PXK)
{
	int time, i,j,an,xm;
	//double lr,prob;
	double mu,sigma2;
	double ***pxz_s,*pxz_s_k,*ax_an,*rp_x;
	double **pz_s,*rp_z,alpha_s,*az;
	double PXZ;//P(X|K)
	int **ZXPQ,**nk,*mk;
	double p1, p2, zran;
	//double lr,prob;


	mk = (int*) calloc(n,sizeof(int));
	nk = (int**) calloc(K,sizeof(int*));
	for (i=0;i<K;i++)
		nk[i]= (int*) calloc(3,sizeof(int));

	az = (double*) calloc(K,sizeof(double));

	pz_s = (double**) calloc(N,sizeof(double*));
	for (j=0;j<N;j++)
		pz_s[j] = (double*) calloc(K,sizeof(double));
	rp_z = (double*) calloc(K,sizeof(double));
	ax_an = (double*) calloc(3,sizeof(double));

	ZXPQ = (int**) calloc(aN,sizeof(int*));
	for(i=0;i<aN;i++)
		ZXPQ[i] = (int*) calloc(N,sizeof(int));

	pxz_s = (double***) calloc(aN,sizeof(double**));
	for (an=0;an<aN;an++)
	{
		pxz_s[an] = (double**) calloc(K,sizeof(double*));
		for (i=0;i<K;i++)
			pxz_s[an][i]= (double*) calloc(xM[an],sizeof(double));
	}
	pxz_s_k = (double*) calloc(K,sizeof(double));

	rp_x = (double*) calloc(3,sizeof(double));

	mu = 0;
	sigma2 = 0;

	for (j=0;j<N;j++)
	{
		for(an=0;an<aN;an++)
		{
			p1 = 0;
			p2 = 0;
			zran = whrandom(); 
			for (i=0;i<K;i++) 
			{							
				p2 = 1.0*(i+1)/K;
				if((zran>p1) && (zran<=p2))
				{
					ZXPQ[an][j]=  i;
					i = K;
				}
				p1 = p2;
			}
		}
	}
	for (time=0;time<times;time++)
    {
		for(j=0;j<N;j++)
		{
		/*calculation m_k, and n_klj*/
		/*m_k: the number of allele copies in individual i 
		that orignated according to Z in population k */
			for(i=0;i<K;i++)
				mk[i] = 0;
			for (an=0;an<aN;an++)
			{
				for(i=0;i<K;i++)
				{
					if (ZXPQ[an][j]==i) 
					{
						mk[i] += 1;
					}
				}
			}
		/* draw q = Pr(Z=i) from q|X,Z */
			for(i=0;i<K;i++)
				az[i]=alpha_o+mk[i];
			dirichletrv(az,K-1,rp_z);
			for (i=0;i<K;i++) 
				pz_s[j][i] = rp_z[i];
		}

		/* draw p =Pr(X|Z) from p|X,Z */
		for (an = 0; an<aN; an++)
		{			
			for (i=0;i<K;i++) 
			{
				for (xm = 0; xm < xM[an]; xm++)
				{
					nk[i][xm] = 0; 
				}
			}			
			for(j=0;j<N;j++)
			{
				for (i=0;i<K;i++) 
				{
					if (ZXPQ[an][j]==i) 
					{
						for (xm = 0; xm < xM[an]; xm++)
						{
							if(fdata[j][1+an]==xm)
								nk[i][xm] += 1;
						}
					}
				}
			}
			for (i=0;i<K;i++) 
			{
				for (xm = 0; xm < xM[an]; xm++)
				{
					ax_an[xm] = ax[an][xm]+nk[i][xm];
				}
				dirichletrv(ax_an,(xM[an]-1),rp_x);
				for (xm = 0; xm<xM[an]; xm++)
					pxz_s[an][i][xm] = rp_x[xm];				
			}			
			for(j=0;j<N;j++)
			{
				for(i=0;i<K;i++)
					pxz_s_k[i] = pxz_s[an][i][fdata[j][1+an]];
				ZXPQ[an][j] = sampleZ(K, pz_s[j], pxz_s_k);
			}
		}
		/*
		alpha_s = whrandom()*10.0;
		lr = N*(log(gammafn(K*alpha_s))-log(K*gammafn(alpha_s))-(log(gammafn(K*alpha_o))-log(K*gammafn(alpha_o))));
		for(j=0;j<N;j++)
		{
			for(i=0;i<K;i++)
				lr += alpha_s*log(pz_s[j][i])-alpha_o*log(pz_s[j][i]);
		}
		prob = whrandom(); 
		if ((log(prob) <=0)&&(log(prob) < lr))
			alpha_o = alpha_s;
		else
			alpha_s = alpha_o;*/
		
		alpha_s = alpha_o =0.01;

		if (time>=times-Nt)
		{
			PXZ =0;
			for(j=0;j<N;j++)
			{
				for(an=0;an<aN;an++)
				{
					PXZ+=log(pxz_s[an][ZXPQ[an][j]][fdata[j][1+an]]);
				}

			}
			mu -= PXZ;
			sigma2 += 2*PXZ*2*PXZ;
		}
	}

	mu *=2.0/Nt;
	sigma2 = sigma2/Nt-mu*mu;

	//normality assumption
	PXK[1][K-1][tt]=-mu/2-sigma2/8;

	//gamma assumption 
	//PXK[1][K-1][tt]=mu*mu/sigma2*log(1-0.5*sigma2/mu);
	
	for (an=0;an<aN;an++)
	{
		for (i=0;i<K;i++)
			free(pxz_s[an][i]);
		free(pxz_s[an]);
	}
	free(pxz_s);
	free(pxz_s_k); 
	
	for (i=0;i<K;i++)
			free(nk[i]);
	free(nk);

	free(rp_x); 
	
	free(ax_an);

	for (j=0;j<N;j++)
		free(pz_s[j]);
	free(pz_s);free(rp_z);free(mk); free(az);


	for(i=0;i<aN;i++)
		free(ZXPQ[i]);
	free(ZXPQ);

 }
