/*
  This file is part of the BJP package.

  BJP package is free software: you can redistribute it and/or modify it
  under the terms of the GNU Lesser General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  BJP package is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
  License for more details.

  You should have received a copy of the GNU Lesser General Public License
  along with BJP package.  If not, see <http://www.gnu.org/licenses/>.
 */
package learn.struct.it;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashMap;
import java.util.List;

import util.math.Gamma;
import util.math.Permutations;
import dataset.Dataset;
import dataset.adtree.ADNode;

/**
 * This is an open source implementation of the Bayesian statistical test of
 * conditional independence (Margaritis, 2005). The Bayesian statistical test of
 * conditional independence is a nonparametric statistical analyzing method that
 * can be used to assess the probability of association or independence of
 * facts. In contrast with the well-known Chi-square ($\chi^2$) test, the
 * Bayesian statistical test works by comparing two competing statistical
 * models: an independence model and a dependence model, which are obtained by
 * computing the posterior probability of each model on data. This source code
 * follows the pseudocode of the technical report
 * "An Open Source Implementation Of The Bayesian Statistical Test Of Independence"
 * , (Schl\"uter and Edera, 2013).
 * 
 * @author fschluter
 * 
 */
public class BayesianTest implements StatisticalIndependenceTest {

	/* Vectors of hyperparameters */
	public double[][] _gammaIJ;

	/* Uniform (homogeneous) values for hyperparameters */
	public double _uniformGamma = 1;
	public double _uniformAlpha = 1; // marginal over gamma
	public double _uniformBeta = 1; // marginal over gamma

	/* Prior Probability of conditional independence model */
	public double priorOfConditionalIndependenceModel = 0.5;

	/* Log of the posterior of the conditional dependence model */
	private double logPosteriorDependenceModel;
	/* Log of the posterior of the conditional independence model */
	private double logPosteriorIndependenceModel;
	/* Log of the likelihood of the conditional dependence model */
	private double loglikelihoodDependentModel;
	/* Log of the likelihood of the conditional independence model */
	private double loglikelihoodIndependentModel;

	/* Indicates if the ADTree cache is used */
	public boolean useADTree;

	/* Stores the ADTree cache */
	private ADNode aDTree;

	/**
	 * This method is used to test the independence of a set of variables x
	 * respect to a set of variables y, conditioning in a set of variables z.
	 * 
	 * @param dataset
	 *            An instance of the dataset class
	 * @param x
	 *            A set of random variables
	 * @param y
	 *            A set of random variables
	 * @param z
	 *            A set of random variables
	 * @param i
	 *            The cardinality of all the variables (i.e. 2 for binary
	 *            variables).
	 * @return A decision of Independences of the triplet, given data
	 */
	public boolean independent(Dataset dataset, List<Integer> x,
			List<Integer> y, List<Integer> z) {
		IndependenceTriplet triplet = new IndependenceTriplet(x, y, z);
		return independent(dataset, triplet);
	}

	/**
	 * This method is used to test the independence truth value of a triplet
	 * conformed by a set of variables x, a set of variables y, and a
	 * conditioning set of variables z.
	 * 
	 * @param dataset
	 *            An instance of the dataset class
	 * @param x
	 *            A random variable
	 * @param y
	 *            A random variable
	 * @param z
	 *            A set of random variables
	 * @param i
	 *            The cardinality of all the variables (i.e. 2 for binary
	 *            variables).
	 * @return A decision of independence of the triplet, given data
	 */
	public boolean independent(Dataset dataset, IndependenceTriplet triplet) {
		List<Integer> x = triplet.getX();
		List<Integer> y = triplet.getY();
		List<Integer> z = triplet.getZ();
		initializeHyperparameters(x, y, dataset.getCardinalities());
		List<int[]> slices = Permutations.getPermutations(z,
				dataset.getCardinalities());

		/* Probability of unconditional dependence for each slice */
		double[] unconditionalDependenceLoglikelihoods = new double[slices
				.size()];
		/* probability of unconditional independence for each slice */
		double[] unconditionalIndependenceLoglikelihoods = new double[slices
				.size()];

		for (int i = 0; i < slices.size(); ++i) {
			int[] slice = slices.get(i);

			// System.out.println(" test x " + x+ " y "+y + " slice "+slice);
			unconditionalDependenceLoglikelihoods[i] = computeUnconditionalDependenceLogLikelihood(
					x, y, dataset, z, slice);
			if (unconditionalDependenceLoglikelihoods[i] == Double.NaN
					|| unconditionalDependenceLoglikelihoods[i] == Double.NEGATIVE_INFINITY
					|| unconditionalDependenceLoglikelihoods[i] == Double.POSITIVE_INFINITY) {
				throw new RuntimeException(
						"unconditional Dependence Loglikelihoods is "
								+ unconditionalDependenceLoglikelihoods[i]);
			}
			unconditionalIndependenceLoglikelihoods[i] = computeUnconditionalIndependenceLogLikelihood(
					x, y, dataset, z, slice);
			if (unconditionalIndependenceLoglikelihoods[i] == Double.NaN
					|| unconditionalIndependenceLoglikelihoods[i] == Double.NEGATIVE_INFINITY
					|| unconditionalIndependenceLoglikelihoods[i] == Double.POSITIVE_INFINITY) {
				throw new RuntimeException(
						"unconditional Dependence Loglikelihoods is "
								+ unconditionalIndependenceLoglikelihoods[i]);
			}
		}

		loglikelihoodIndependentModel = computeLoglikelihoodOfIndependentModel(unconditionalIndependenceLoglikelihoods);
		if (loglikelihoodIndependentModel == Double.NaN
				|| loglikelihoodIndependentModel == Double.NEGATIVE_INFINITY
				|| loglikelihoodIndependentModel == Double.POSITIVE_INFINITY) {
			throw new RuntimeException(
					"conditional Independence Loglikelihoods is "
							+ loglikelihoodIndependentModel);
		}
		loglikelihoodDependentModel = computeLoglikelihoodOfDependentModel(
				unconditionalIndependenceLoglikelihoods,
				unconditionalDependenceLoglikelihoods);
		if (loglikelihoodDependentModel == Double.NaN
				|| loglikelihoodDependentModel == Double.NEGATIVE_INFINITY
				|| loglikelihoodDependentModel == Double.POSITIVE_INFINITY) {
			throw new RuntimeException(
					"conditional Dependence Loglikelihoods is "
							+ loglikelihoodDependentModel);
		}

		// a is a name for log P(D|M_dep) - log P(D|M_ind) + log P(M_dep) - log
		// P(M_ind)
		double a = loglikelihoodDependentModel - loglikelihoodIndependentModel
				+ Math.log(1 - priorOfConditionalIndependenceModel)
				- Math.log(priorOfConditionalIndependenceModel);

		// applying a trick to compute log(exp()) without underflow
		// see
		// https://hips.seas.harvard.edu/blog/2013/01/09/computing-log-sum-exp/
		// -log (1+exp(a)) is now -a -log(1+exp(-a))
		logPosteriorIndependenceModel = -a - Math.log1p(Math.exp(-a));

		logPosteriorDependenceModel = -Math.log1p(Math
				.exp(loglikelihoodIndependentModel
						- loglikelihoodDependentModel));

		if (logPosteriorIndependenceModel > logPosteriorDependenceModel) {
			return true;
		} else {
			return false;
		}
	}

	/**
	 * @param dataset
	 * @param x
	 * @param y
	 * @param datasetForSlice
	 * @param z
	 * @param slice
	 * @return
	 */
	private double computeUnconditionalIndependenceLogLikelihood(
			List<Integer> x, List<Integer> y, Dataset dataset, List<Integer> z,
			int[] slice) {
		if (!useADTree) {
			return computeUnconditionalIndependenceLogLikelihoodFromData(x, y,
					dataset, z, slice);
		} else {
			return computeUnconditionalIndependenceLogLikelihoodFromADTree(x, y,
					dataset, z, slice);
		}
	}
	
	/**
	 * @param dataset
	 * @param x
	 * @param y
	 * @param datasetForSlice
	 * @param z
	 * @param slice
	 * @return
	 */
	private double computeUnconditionalIndependenceLogLikelihoodFromData(
			List<Integer> x, List<Integer> y, Dataset dataset, List<Integer> z,
			int[] slice) {
		Dataset datasetForSlice = dataset.getSubdataset(z, slice);
		List<int[]> permutationsOfX = Permutations.getPermutations(x,
				datasetForSlice.getCardinalities());
		int[] countsOfX = new int[permutationsOfX.size()];
		for (int j = 0; j < permutationsOfX.size(); ++j) {
			int[] permutationOfX = permutationsOfX.get(j);
//			int[] values = new int[x.size()];
//			System.arraycopy(permutationOfX, 0, values, 0,
//					permutationOfX.length);
//			values[permutationOfX.length] = permutationOfX[0];

			countsOfX[j] = datasetForSlice.getCounts(x, permutationOfX);
			// System.out.println("counts of x" + x + " " +countsOfX[j]);
		}
		List<int[]> permutationsOfY = Permutations.getPermutations(y,
				datasetForSlice.getCardinalities());
		int[] countsOfY = new int[permutationsOfY.size()];
		for (int j = 0; j < permutationsOfY.size(); ++j) {
			int[] permutationOfY = permutationsOfY.get(j);
			
//			int[] values = new int[x.size()];
//			System.arraycopy(permutationOfY, 0, values, 0,
//					permutationOfY.length);
//			values[permutationOfY.length] = permutationOfY[0];

			countsOfY[j] = datasetForSlice.getCounts(y, permutationOfY);
			// System.out.println("counts of y " +y + " " +countsOfY[j]);
		}
		double alpha = permutationsOfX.size() * _uniformAlpha;
		double beta = permutationsOfY.size() * _uniformBeta;
		return dirichletFunction(countsOfX, alpha)
				+ dirichletFunction(countsOfY, beta);
	}
	
	/**
	 * @param dataset
	 * @param x
	 * @param y
	 * @param datasetForSlice
	 * @param z
	 * @param slice
	 * @return
	 */
	private double computeUnconditionalIndependenceLogLikelihoodFromADTree(
			List<Integer> x, List<Integer> y, Dataset dataset, List<Integer> z,
			int[] slice) {
		int n = dataset.getNumberOfVariables();
		
		HashMap<Integer,Integer> context = new HashMap<Integer, Integer>();
		for (int i = 0; i < z.size(); ++i){ 
			context.put(z.get(i), slice[i]);
		}
		
		int[] cards = new int[n];
		for (int i = 0; i < n; ++i) cards[i] = 2;

		BitSet vars = new BitSet(n);
		for (int i=0;i<x.size();i++) vars.set(x.get(i));
		for (int i=0;i<z.size();i++) vars.set(z.get(i));

		int[] countsOfX = aDTree.makeContab(vars, context, cards, dataset);

		vars = new BitSet(n);
		for (int i=0;i<y.size();i++) vars.set(y.get(i));
		for (int i=0;i<z.size();i++) vars.set(z.get(i));

		int[] countsOfY = aDTree.makeContab(vars, context, cards, dataset);
		
		double alpha = countsOfX.length * _uniformAlpha;
		double beta = countsOfY.length * _uniformBeta;
		return dirichletFunction(countsOfX, alpha)
				+ dirichletFunction(countsOfY, beta);
	}

	private double computeUnconditionalDependenceLogLikelihood(List<Integer> x,
			List<Integer> y, Dataset dataset, List<Integer> z, int[] slice) {
		if (!useADTree) {
			return computeUnconditionalDependenceLogLikelihoodFromData(x, y,
					dataset, z, slice);
		} else {
			return computeUnconditionalDependenceLogLikelihoodFromADTree(x, y,
					dataset, z, slice);
		}
	}

	private double computeUnconditionalDependenceLogLikelihoodFromData(
			List<Integer> x, List<Integer> y, Dataset dataset, List<Integer> z,
			int[] slice) {
		Dataset datasetForSlice = dataset.getSubdataset(z, slice);
		List<Integer> xyList = new ArrayList<Integer>();
		xyList.addAll(x);
		xyList.addAll(y);
		List<int[]> permutationsOfXY = Permutations.getPermutations(xyList,
				datasetForSlice.getCardinalities());
		int[] countsOfXY = new int[permutationsOfXY.size()];
		for (int j = 0; j < permutationsOfXY.size(); ++j) {
			int[] permutationOfXY = permutationsOfXY.get(j);
//			int[] values = new int[xyList.size()];
//			System.arraycopy(permutationOfXY, 0, values, 0,
//					permutationOfXY.length);
//			values[permutationOfXY.length] = permutationOfXY[0];
			countsOfXY[j] = datasetForSlice.getCounts(xyList, permutationOfXY);
//			countsOfXY[j] = datasetForSlice.getCounts(xyList, permutationOfXY);
		}
		return dirichletFunction(countsOfXY, _uniformGamma);
	}

	private double computeUnconditionalDependenceLogLikelihoodFromADTree(
			List<Integer> x, List<Integer> y, Dataset dataset, List<Integer> z,
			int[] slice) {
		List<Integer> xyList = new ArrayList<Integer>();
		xyList.addAll(x);
		xyList.addAll(y);
		xyList.addAll(z);
		
		int n = dataset.getNumberOfVariables();
		BitSet vars = new BitSet(n);
		for (int i=0;i<xyList.size();i++) vars.set(xyList.get(i));
		
		int[] cards = new int[n];
		for (int i = 0; i < n; ++i) cards[i] = 2;

		HashMap<Integer,Integer> context = new HashMap<Integer, Integer>();
		for (int i = 0; i < z.size(); ++i){ 
			context.put(z.get(i), slice[i]);
		}
		
		int[] countsOfXY = aDTree.makeContab(vars, context, cards, dataset);

		return dirichletFunction(countsOfXY, _uniformGamma);
	}

	/**
	 * Initializing hyperparameters uniformly. This is a common decision to
	 * initialize uniformly all the hyperparameters with a value 1, which means
	 * that there is no prior information about the hyperparameters.
	 * 
	 * @param x
	 *            First variable of the test
	 * @param y
	 *            Second variable of the test
	 * @param cardinalities
	 *            Cardinalities of the variables of the test
	 */
	private void initializeHyperparameters(List<Integer> x, List<Integer> y,
			int[] cardinalities) {
		/* computing cardinality of sets of variables */
		int totalCardinalityOfX = 0, totalCardinalityOfY = 0;
		for (int i : x)
			totalCardinalityOfX += cardinalities[i];
		for (int i : y)
			totalCardinalityOfY += cardinalities[i];

		/* initializing gamma hyperparameters */
		_gammaIJ = new double[totalCardinalityOfX][totalCardinalityOfY];
		for (int i = 0; i < totalCardinalityOfX; ++i)
			for (int j = 0; j < totalCardinalityOfY; ++j)
				_gammaIJ[i][j] = _uniformGamma;

	}

	/**
	 * Compute the data likelihood of the independent model. Formula (19) in
	 * (Margaritis and Bromberg, 2009).
	 * 
	 * @param unconditionalDependenceLikelihoods
	 */
	private double computeLoglikelihoodOfIndependentModel(
			double[] unconditionalIndependenceLogikelihoods) {
		double logLikelihoodsDependenceModel = 0;
		for (double d : unconditionalIndependenceLogikelihoods) {
			logLikelihoodsDependenceModel += d;
		}
		return logLikelihoodsDependenceModel;
	}

	/**
	 * Compute the data likelihood of the dependent model. Formula (20) in
	 * (Margaritis and Bromberg, 2009)
	 * 
	 * @param unconditionalDependenceLikelihoods2
	 * @param unconditionalIndependenceLoglikelihoods
	 * @return
	 */
	private double computeLoglikelihoodOfDependentModel(
			double[] unconditionalIndependenceLoglikelihoods,
			double[] unconditionalDependenceLoglikelihoods) {
		/* number of slices */
		int K = unconditionalDependenceLoglikelihoods.length;

		/* probability of the dependent model */
		double not_pMCI = 1 - priorOfConditionalIndependenceModel;

		/* prior of the independence model for each slice k */
		double pk = Math.log(priorOfConditionalIndependenceModel) / K;

		/* prior of the dependence model for each slice k */
		double qk = Math.log1p(-Math.pow(priorOfConditionalIndependenceModel,
				1.0 / K));

		if (qk == Double.NaN || qk == Double.NEGATIVE_INFINITY
				|| qk == Double.POSITIVE_INFINITY) {
			throw new RuntimeException("Math error with prior values, qk is  "
					+ qk + ". Prior used: "
					+ priorOfConditionalIndependenceModel);
		}
		if (pk == Double.NaN || pk == Double.NEGATIVE_INFINITY
				|| pk == Double.POSITIVE_INFINITY) {
			throw new RuntimeException("Math error with prior values, pk is  "
					+ pk + ". Prior used: "
					+ priorOfConditionalIndependenceModel);
		}

		double aux1, aux2 = 0, aux3 = 0;

		for (int k = 0; k < K; ++k) {
			double gk = unconditionalIndependenceLoglikelihoods[k];
			double hk = unconditionalDependenceLoglikelihoods[k];
			if (gk >= hk) { // preventing zero division
				aux1 = Math.log1p(Math.exp((hk + qk) - (gk + pk)));
				aux2 += gk + pk + aux1;
				aux3 += -aux1;
			} else {
				aux1 = Math.log1p(Math.exp((gk + pk) - (hk + qk)));
				aux2 += (hk + qk) + aux1;
				aux3 += (gk + pk) - (hk + qk) - aux1;
			}
		}
		double aux4 = Math.log(-Math.expm1(aux3));
		return aux2 + aux4 - Math.log(not_pMCI);
	}

	/**
	 * Computes the Dirichlet function for a given set of counts and the
	 * corresponding hyper-parameters.
	 * 
	 * @param counts
	 *            An array of counts of the variables of interest
	 * @param hyperparameter
	 *            The hyper-parameters of the variables of interest
	 * @return The Dirichlet value (factorial function).
	 */
	final private double dirichletFunction(int counts[],
			double hyperparameter) {
		double totalNumberOfCounts = 0;
		double marginalHyperparameter = 0;
		for (int i = 0; i < counts.length; i++) {
			totalNumberOfCounts += counts[i];
			marginalHyperparameter += hyperparameter;
		}

		double dirichletValue = 0;
		for (int i = 0; i < counts.length; i++) {
			dirichletValue += Gamma.logGamma(hyperparameter + counts[i])
					- Gamma.logGamma(hyperparameter);
		}
		dirichletValue += Gamma.logGamma(marginalHyperparameter)
				- Gamma.logGamma(marginalHyperparameter + totalNumberOfCounts);

		return dirichletValue;
	}

	/**
	 * @return A vector of double values, wit logprobs[0] the log probability of
	 *         the independent model, and logprobs[1] the log probability of the
	 *         dependent model, and
	 */
	public double[] getLogProbs() {
		double[] logProbs = new double[2];
		logProbs[0] = logPosteriorIndependenceModel;
		logProbs[1] = logPosteriorDependenceModel;
		return logProbs;
	}

	/**
	 * @return A vector of double values, wit loglikelihoods[0] the log
	 *         likelihood of the independent model, and logprobs[1] the log
	 *         likelihood of the dependent model, and
	 */
	public double[] getLoglikelihoods() {
		double[] logLikelihoods = new double[2];
		logLikelihoods[0] = loglikelihoodIndependentModel;
		logLikelihoods[1] = loglikelihoodDependentModel;
		return logLikelihoods;
	}

	public void setUniformParameters(int alpha) {
		_uniformGamma = alpha;
		_uniformAlpha = alpha;
		_uniformBeta = alpha;
	}

	public void setUseADTree(boolean useADTree) {
		this.useADTree = useADTree;
	}

	public void setADTree(ADNode aDTree) {
		this.aDTree = aDTree;
	}

}
