#include "MpUtil.h"

int N = 1000;		// Number of Data Points
int K = 3;			// Number of Parameters
int Burnin = 1000;	// Number of Burn in Iterations
int Gibbs  = 5000;	// Number of Gibbs Iterations
int Thin = 1;		// Thinning Factor

long RandomSeed = -23497234;	// RandomSeed

Vector<int>    Y(N);	// Dependent Variable
Matrix<double> X(N,K);	// Independent Variables

/* Globar Priol Parameters */
Vector<double> BetaPriorMu;
SymmetricMatrix<double> BetaPriorSigmaInv;

double LogPosterior(const Vector<double> &Beta)
{
	double SumLog = Zero;
	for(int n=0; n<N; n++)
	{
		double Temp = Zero;
		for(int k=0; k<K; k++) Temp += Beta(k) * X(n,k);

		double Prob;
		if(Y(n) == 0) Prob = One - NormalCdf(Temp);
		else          Prob = NormalCdf(Temp);

		SumLog += log(Prob);
	}

	/* Prior */
	double LogPrior = ProdXtAX(Beta - BetaPriorMu,BetaPriorSigmaInv);
	LogPrior *= -0.5;
	SumLog += LogPrior;
	return SumLog;
}

int main()
{
	Print("Begin");

	/* True Beta */
	Vector<double> Beta0 = LoadVector<double>("-1,0.5,-1.5");

	/* Generate the Data */

		for(int n=0; n<N; n++)
		{
			X(n,0) = One;
			X(n,1) = NormalSample(RandomSeed);
			X(n,2) = UniformSample(RandomSeed) - 0.5 * pow(Abs(X(n,1)),1.5);

			double Temp = Beta0(0) * X(n,0) + Beta0(1) * X(n,1) + Beta0(2) * X(n,2) + NormalSample(RandomSeed);

			if(Temp >= Zero) Y(n) = 1;
			else             Y(n) = 0;
		}

		Output("X.dat",X);

	/* Set the Priors */

		BetaPriorMu = Vector<double>(K,Zero);
		BetaPriorMu.Name = "BetaPriorMu";
		SymmetricMatrix<double> BetaPriorSigma(K,Zero);
		BetaPriorSigma.Name = "BetaPriorSigma";
		for(int k=0; k<K; k++) BetaPriorSigma(k,k) = 1000.0;
		BetaPriorSigmaInv = InverseCholesky(BetaPriorSigma);
		BetaPriorSigmaInv.Name = "BetaPriorSigmaInv";

	/* 1: Gibbs Sampler */

		/* Compute X^T X */

		SymmetricMatrix<double> XtX(K,Zero);
		XtX.Name = "XtX";
		for(int k1=0; k1<K; k1++) for(int k2=0; k2<=k1; k2++)
			for(int n=0; n<N; n++) XtX(k1,k2) += X(n,k1) * X(n,k2);

		SymmetricMatrix<double> Temp = BetaPriorSigmaInv + XtX;
		Temp.Name = "Temp";

		SymmetricMatrix<double> SigmaTilde = InverseCholesky(Temp);

		LowerMatrix<double> LTilde = CholeskyDecompL(SigmaTilde);

		Vector<double> BetaPriorSigmaInv_BetaPriorMu = BetaPriorSigmaInv * BetaPriorMu;

		Vector<double> YStarSamp(N);				YStarSamp.Name = "YStarSamp";
		Vector<double> XtYStar(K);					XtYStar.Name   = "XtYStar";
		Matrix<double> BetaSamp1(Burnin+Gibbs+1,K);	BetaSamp1.Name = "BetaSamp";
		Vector<double> s(K);						s.Name         = "s";

		/* Initialize Variables */

		Set(YStarSamp,Zero);
		Set(BetaSamp1,Zero);

		/* Iterate the Gibbs Sampler */

		for(int r=1; r<=Burnin+Gibbs; r++)
		{

			if(r % 100 == 0) Print("r",r);

			/* First Draw YStar | Y, X, Beta */
			for(int n=0; n<N; n++)
			{
				/* Form X(n) * Beta */
				double XBeta = Zero;
				for(int k=0; k<K; k++) XBeta += BetaSamp1(r-1,k) * X(n,k);

				/* Sample YStar(n) from the truncated normal distribution */
				if(Y(n) == 0) YStarSamp(n) = NormalTruncatedSample(XBeta,One,-Inf ,Zero,RandomSeed);
				else          YStarSamp(n) = NormalTruncatedSample(XBeta,One, Zero,Inf ,RandomSeed);
			}

			/* Compute XtYStar */
			Set(XtYStar,Zero);
			for(int n=0; n<N; n++) for(int k=0; k<K; k++) XtYStar(k) += X(n,k) * YStarSamp(n);

			/* Now draw Beta | Y, YStar, X */
			Vector<double> MuTilde = SigmaTilde * (BetaPriorSigmaInv_BetaPriorMu + XtYStar);
			for(int k=0; k<K; k++) s(k) = NormalSample(RandomSeed);
			for(int k=0; k<K; k++) BetaSamp1(r,k) = MuTilde(k);
			for(int k1=0; k1<K; k1++) for(int k2=0; k2<K; k2++) BetaSamp1(r,k1) += LTilde(k1,k2) * s(k2);
		}

		Output("BetaSamp1.dat",BetaSamp1);

		/* Compute Posterior Mean for the Last Gibbs Observations */

		Vector<double> Beta1(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++) for(int k=0; k<K; k++) Beta1(k) += BetaSamp1(r,k) / (double)Gibbs;

		SymmetricMatrix<double> BetaVar(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++)
			for(int k1=0; k1<K; k1++) for(int k2=0; k2<=k1; k2++)
				BetaVar(k1,k2) += (BetaSamp1(r,k1) - Beta1(k1)) * (BetaSamp1(r,k2) - Beta1(k2)) / (double)Gibbs;

		Vector<double> BetaSe1(K);
		for(int k=0; k<K; k++) BetaSe1(k) = sqrt(BetaVar(k,k));

		Print("Beta0",Beta0);
		Print("Beta1",Beta1);
		Print("BetaSe1",BetaSe1);

	/* 2: Metropolis Hastings Sampler */

		Matrix<double> BetaSamp2(Burnin+Gibbs+1,K,Zero);
		BetaSamp2.Name  = "BetaSamp2";
		Set(BetaSamp2,Zero);

		Vector<double> x2(K);	x2.Name = "x";
		Vector<double> y2(K);	y2.Name = "y";

		/* Iterate the Sampler */
		int Accept = 0;
		double SigmaCurr = 0.01;
		for(int r=1; r<=Burnin+Gibbs; r++)
		{
			if(r % 100 == 0)
			{
				Print("r",r);
				Print("x2",x2);
			}

			/* Current Point */
			for(int k=0; k<K; k++) x2(k) = BetaSamp2(r-1,k);

			/* Potential Point */
			for(int k=0; k<K; k++) y2(k) = x2(k) + SigmaCurr * NormalSample(RandomSeed);

			/* Compute the posterior */
			double LogPx = LogPosterior(x2);
			double LogPy = LogPosterior(y2);

			/* ... */
			double u = UniformSample(RandomSeed);
			if(u <= Min(One,exp(LogPy - LogPx)))
			{
				/* Accept MH Step */
				if(r > Burnin) Accept++;
				for(int k=0; k<K; k++) BetaSamp2(r,k) = y2(k);
			}
			else
			{
				/* Reject MH Step */
				for(int k=0; k<K; k++) BetaSamp2(r,k) = x2(k);
			}

			if(r > Burnin && r % 100 == 0) Print("Curr Accept",(double)Accept / (r - Burnin));
		}

		Print("% Accept",Accept / Gibbs);

		/* Compute posterior mean for last Gibbs observations */

		Vector<double> Beta2(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++) for(int k=0; k<K; k++)
			Beta2(k) += BetaSamp2(r,k) / (double)Gibbs;

		SymmetricMatrix<double> BetaVar2(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++)
			for(int k1=0; k1<K; k1++) for(int k2=0; k2<=k1; k2++)
				BetaVar2(k1,k2) += (BetaSamp2(r,k1) - Beta2(k1)) * (BetaSamp2(r,k2) - Beta2(k2)) / (double)Gibbs;

		Vector<double> BetaSe2(K);
		for(int k=0; k<K; k++) BetaSe2(k) = sqrt(BetaVar2(k,k));


	/* 3: Gibbs with MH */

		Matrix<double> BetaSamp3(Burnin+Gibbs+1,K,Zero);
		BetaSamp3.Name  = "BetaSamp2";
		Set(BetaSamp3,Zero);

		Vector<double> x3(K);	x3.Name = "x";
		Vector<double> y3(K);	y3.Name = "y";

		/* Iterate the Sampler */
		Vector<int> Accept3(K,0);
		Vector<double> AcceptRate3(K);
		Vector<double> SigmaCurr3(K,0.01);
		int kCurr = 0;
		for(int r=1; r<=Burnin+Gibbs; r++)
		{
			if(r % 100 == 0)
			{
				Print("r",r);
				Print("x2",x2);
			}

			if(kCurr >= K) kCurr = 0;

			/* Current Point */
			for(int k=0; k<K; k++) x3(k) = BetaSamp3(r-1,k);

			/* Potential Point */
			for(int k=0; k<K; k++) y3(k) = x3(k);
			y3(kCurr) += SigmaCurr3(kCurr) * NormalSample(RandomSeed);

			/* Compute the posterior */
			double LogPx = LogPosterior(x3);
			double LogPy = LogPosterior(y3);

			/* ... */
			double u = UniformSample(RandomSeed);
			if(u <= Min(One,exp(LogPy - LogPx)))
			{
				/* Accept MH Step */
				if(r > Burnin) Accept3(kCurr)++;
				for(int k=0; k<K; k++) BetaSamp3(r,k) = y3(k);
			}
			else
			{
				/* Reject MH Step */
				for(int k=0; k<K; k++) BetaSamp3(r,k) = x3(k);
			}

			if(r > Burnin && r % 100 == 0)
			{
				for(int k=0; k<K; k++) AcceptRate3(k) = (double)Accept3(k) / (r - Burnin);
				Print("Curr Accept",AcceptRate3);
			}

			kCurr++;
		}

		Print("% Accept",Accept / Gibbs);

		/* Compute posterior mean for last Gibbs observations */

		Vector<double> Beta3(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++) for(int k=0; k<K; k++)
			Beta3(k) += BetaSamp3(r,k) / (double)Gibbs;

		SymmetricMatrix<double> BetaVar3(K,Zero);
		for(int r=Burnin+1; r<=Burnin+Gibbs; r++)
			for(int k1=0; k1<K; k1++) for(int k2=0; k2<=k1; k2++)
				BetaVar3(k1,k2) += (BetaSamp3(r,k1) - Beta3(k1)) * (BetaSamp3(r,k2) - Beta3(k2)) / (double)Gibbs;

		Vector<double> BetaSe3(K);
		for(int k=0; k<K; k++) BetaSe3(k) = sqrt(BetaVar3(k,k));


		Print("Beta0",Beta0);
		Print("Beta1",Beta1);
		Print("BetaSe1",BetaSe1);
		Print("Beta2",Beta2);
		Print("BetaSe2",BetaSe2);
		Print("Beta3",Beta3);
		Print("BetaSe3",BetaSe3);

		Output("BetaSamp1.dat",BetaSamp1);
		Output("BetaSamp2.dat",BetaSamp2);
		Output("BetaSamp3.dat",BetaSamp3);


	Print("End");
	return 0;
}