code: package xyz.lejon.bayes.models.probit;
  
  import static java.lang.Math.abs;
  import static java.lang.Math.pow;
  import static java.lang.Math.sqrt;
  import static org.ejml.ops.CommonOps.addEquals;
  import static org.ejml.ops.CommonOps.invert;
  import static org.ejml.ops.CommonOps.multTransA;
  import static xyz.lejon.utils.MatrixOps.pow;
  import static xyz.lejon.utils.MatrixOps.rnorm;
  import static xyz.lejon.utils.MatrixOps.scalarDivide;
  import static xyz.lejon.utils.MatrixOps.scalarMultiply;
  import static xyz.lejon.utils.MatrixOps.scalarPlus;
  import static xyz.lejon.utils.MatrixOps.sqrt;
  import static xyz.lejon.utils.MatrixOps.sum;
  
  import java.io.IOException;
  import java.util.Arrays;
  
  import org.apache.commons.math3.distribution.MultivariateNormalDistribution;
  import org.ejml.data.DenseMatrix64F;
  import org.ejml.ops.CommonOps;
  
  import xyz.lejon.configuration.DOConfiguration;
  import xyz.lejon.sampling.Gamma;
  import xyz.lejon.utils.MatrixOps;
  
  public class HorseshoeDOProbitPX extends SerialDOSampler {
  
  	MatrixOps mo = new MatrixOps();
  	double [][] Theta;
  	double Sigma = 1.0;
  	double G = 1.0;
  	double Delta = 0.1;
  	double [] Tau;
  	double [][] Lambda;
  	double [][] precMtrx;
  	
  	public HorseshoeDOProbitPX(DOConfiguration config, double[][] xs, int[] ys) throws IOException {
  		super(config, xs, ys);
  		Theta = new double[noClasses][];
  		for (int i = 0; i < Theta.length; i++) {
  			Theta[i] = new double[noCovariates];
  			Arrays.fill(Theta[i], 1.0);
  		}
  		Lambda = new double[noClasses][];
  		for (int i = 0; i < Lambda.length; i++) {
  			Lambda[i] = new double[noCovariates];
  			Arrays.fill(Lambda[i], 1.0);
  		}
  		Tau = new double[noClasses];
  		for (int i = 0; i < Tau.length; i++) {
  			Tau[i] = Math.abs(Delta)*G;
  		}
  		precMtrx = new double[Xd.numCols][Xd.numCols];
  		for (int i = 0; i < noClasses; i++) {
  			for (int j = 0; j < noCovariates; j++) {
  				betas[i][j] = rnorm(0.0, 1.0);
  			}
  		}
  	}
  	
  	public double [] calcTheta(double[] Beta, double Sigma, double Delta, double[] Lambda) {
  		double [] Theta = scalarDivide(Beta,scalarMultiply(Lambda,Sigma*Delta));
  		return Theta;
  	}
  	
  	public double sampleSigma2(double [][] Res2, double Tau, double[] Lambda, double [][] Y) {
  		// Jeffreys prior is assumed
  		// for(i in 1:n)
  		//  {
  		//    Res2[,i] = {Y[,i]^2}/{1+(Tau^2)*{Lambda^2}}
  		//  }
  		  
  		for (int row = 0; row < Res2.length; row++) {			
  			for(int i=0; i < noCovariates; i++)
  			{
  				Res2[row][i] = pow(Y[row][i],2) / (1 + pow(Tau,2) * pow(Lambda[i],2) );
  			}
  		}
  		// RSS = sum(Res2)
  		double RSS = mo.sumPar(Res2);
  		// Sigma2 = 1/rgamma(1,n*p/2, rate = RSS/2)
  		double Sigma2 = 1/Gamma.rgamma(1,noRows*noCovariates/2.0, RSS/2);
  		// Sigma = sqrt(Sigma2)
  		
  		return Sigma2;
  	}
  
  	public double [] sampleTau(double Tau, double [] Theta, double [] Ybar, double Sigma, double [] Lambda) {
  		//G = 1/sqrt(rgamma(1,(p+1)/2, rate = (1+sum(Theta^2))/2))
  		double G = 1/sqrt(Gamma.rgamma(1, (noCovariates+1)/2, (1+sum(pow(Theta,2)))/2));
  		//Z = Ybar/(Sigma*Theta*Lambda)
  		double [] Z = scalarDivide(Ybar, scalarMultiply(scalarMultiply(Theta, Sigma),Lambda));
  		//a = n*(Lambda*Theta)^2
  		double [] a = scalarMultiply(pow(scalarMultiply(Lambda,Theta),2),noRows);
  		//b = sum(a)
  		double b = sum(a);
  		//s2 = 1/(1+b)
  		double s2 = 1/(1+b);
  		//m = {s2}*sum(a*Z)
  		double m = s2*sum(scalarMultiply(a,Z));
  		//Delta = rnorm(1,m,sqrt(s2))
  		double Delta = rnorm(m,sqrt(s2));
  		//Tau = abs(Delta)*G
  		double newTau = abs(Delta)*G;
  
  		double [] res = {newTau, Delta};
  		return res;
  	}
  
  	public double [] sampleLambda(double Tau, double [] Theta, double [] Ybar, 
  			double Sigma, double [] Lambda, double Delta) {
  		//Z = Ybar/(Sigma*Delta*Theta)
  		double [] Z = scalarDivide(Ybar, scalarMultiply(scalarMultiply(Theta, Sigma),Delta));
  		//V2 = 1/rgamma(p,1,rate=(Lambda^2+1)/2)
  		double [] V2 = scalarDivide(1,Gamma.rgamma(noCovariates,1.0,scalarDivide(scalarPlus(pow(Lambda,2),1),2)));
  		//num1 = n*V2*((Delta*Theta)^2)
  		double [] num1 = scalarMultiply(scalarMultiply(V2,(pow(scalarMultiply(Theta,Delta),2))),noRows);
  		//den = 1 + num1
  		double [] den = scalarPlus(num1,1);
  		//s = sqrt(V2/den)
  		double [] s = sqrt(scalarDivide(V2,den));
  		//m = {num1/den}*Z
  		double [] m = scalarMultiply(scalarDivide(num1,den),Z);
  		//Lambda = rnorm(p,m,s)
  		double [] newLambda = rnorm(noCovariates,m,s);
  
  		return newLambda;
  	}
  
  	public double [] sampleBeta(double Tau, double [] Lambda, double Sigma2, double [] Ybar) {
  		// a = (Tau^2)*(Lambda^2)
  		double [] a = scalarMultiply(pow(Lambda,2),pow(Tau,2));
  		// b = n*a
  		double [] b = scalarMultiply(a,noRows);
  		// s = sqrt(Sigma2*a/{1+b})
  		double [] s = sqrt(scalarDivide(scalarMultiply(a,Sigma2),scalarPlus(b,1.0)));
  		// m = {b/{1+b}}*Ybar
  		double [] m = scalarMultiply(scalarDivide(b,scalarPlus(b,1)),Ybar);
  		// Beta = rnorm(p, m, s)
  		double [] Beta = rnorm(noCovariates, m, s);
  		
  		return Beta;
  	}
  	
  	@Override
  	public double [] sampleBeta(int k) {
  		double [] zColk = zsT[k];
  		
  		Theta[k]  = calcTheta(betas[k], Sigma, Delta, Lambda[k]);
  		//System.out.println("Theta: " + MatrixOps.doubleArrayToPrintString(Theta));
  		double [] res = sampleTau(Tau[k], Theta[k], betas[k], Sigma, Lambda[k]);
  		Tau[k]    = res[0];
  		Delta     = res[1];
  		System.out.println(MatrixOps.arrToStr(Tau, "Tau"));
  		Lambda[k] = sampleLambda(Tau[k], Theta[k], betas[k], Sigma, Lambda[k], Delta);
  		//DenseMatrix64F td = new DenseMatrix64F(Lambda);
  		//System.out.println("New Lambda: " + td);
  
  		// Rebuild precision matrix
  		for (int i = 0; i < precMtrx.length; i++) {
  			for (int j = 0; j < precMtrx[i].length; j++) {
  				if(i==j) {
  					precMtrx[i][j] = Math.pow(Tau[k], 2) * Math.pow(Lambda[k][i],2);
  				}
  			}
  		}
  		//System.out.println(MatrixOps.doubleArrayToPrintString(precMtrx));
  		priorPrecision = new DenseMatrix64F(precMtrx);
  		
  		priorMean = new DenseMatrix64F(XtX);
  		addEquals(priorMean, priorPrecision);
  		invert(priorMean);
  		Stilde = MatrixOps.extractDoubleArray(priorMean);
  		
  		DenseMatrix64F zColKd = new DenseMatrix64F(zColk.length,1);
  		zColKd.setData(zColk);
  		DenseMatrix64F localTilde = new DenseMatrix64F(muTilde);
  		DenseMatrix64F localMu = new DenseMatrix64F(Xd.numCols, 1);
  		multTransA(Xd, zColKd, localTilde);
  		CommonOps.mult(priorMean, localTilde, localMu);
  		double [] mu_tile = localMu.getData();
  
  		MultivariateNormalDistribution mvn = new MultivariateNormalDistribution(mu_tile, Stilde);
  		betas[k] = mvn.sample();
  		return betas[k];
  	}
  	
  	public void sampleBetaOLD(int k) {
  		
  		Theta[k]  = calcTheta(betas[k], Sigma, Delta, Lambda[k]);
  		//System.out.println("Theta: " + MatrixOps.doubleArrayToPrintString(Theta));
  		double [] res = sampleTau(Tau[k], Theta[k], betas[k], Sigma, Lambda[k]);
  		Tau[k]    = res[0];
  		Delta     = res[1];
  		//System.out.println(MatrixOps.arrToStr(Tau, "Tau"));
  		Lambda[k] = sampleLambda(Tau[k], Theta[k], betas[k], Sigma, Lambda[k], Delta);
  		//DenseMatrix64F td = new DenseMatrix64F(Lambda);
  		//System.out.println("New Lambda: " + td);
  
  		// Rebuild precision matrix
  		for (int i = 0; i < precMtrx.length; i++) {
  			for (int j = 0; j < precMtrx[i].length; j++) {
  				if(i==j) {
  					precMtrx[i][j] = Math.pow(Tau[k], 2) * Math.pow(Lambda[k][i],2);
  				}
  			}
  		}
  		//System.out.println(MatrixOps.doubleArrayToPrintString(precMtrx));
  		priorPrecision = new DenseMatrix64F(precMtrx);
  		
  		DenseMatrix64F XtXA = new DenseMatrix64F(XtX);
  		addEquals(XtXA, priorPrecision);
  		invert(XtXA);
  		Stilde = MatrixOps.extractDoubleArray(XtXA);
  		//System.out.println(XtXA);
  		DenseMatrix64F mumuk = new DenseMatrix64F(Xd.numCols, 1);
  		
  		double [] zColk = zsT[k];
  		DenseMatrix64F ak = new DenseMatrix64F(zColk.length,1);
  		ak.setData(zColk);
  		CommonOps.multAddTransA(Xd, ak, mumuk);
  		DenseMatrix64F mud = new DenseMatrix64F(Xd.numCols,1);
  		//System.out.println("XtXA: " +XtXA);
  		//System.out.println("mumuk:" + mumuk);		
  		CommonOps.multAdd(XtXA, mumuk, mud);
  		//System.out.println(mud);
  		double [] mu = MatrixOps.extractDoubleVector(mud);
  
  		MultivariateNormalDistribution mvn = new MultivariateNormalDistribution(mu, Stilde);
  		betas[k] = mvn.sample();
  	}
  }
  