/*
  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.algs;

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

import util.math.Permutations;
import learn.struct.scoring.UndirectedGraphScore;
import learn.struct.scoring.GraphScorable;
import learn.struct.scoring.mpl.MPL;
import util.math.Gamma;
import rep.UndirectedGraph;
import dataset.Dataset;
import dataset.adtree.ADNode;

/**
 * This method optimizes the Marginal Pseudo Likelihood scoring function for
 * Markov networks structure discover, as proposed in (Pensar et al, 2014),
 * Section 5.
 * 
 * @author fschluter
 *
 */
public class MPLInterIAMBandHillClimbing {

	/**
	 * The number of variables of the domain
	 */
	private int n;

	/**
	 * The output network learned by IBMAP-HC
	 */
	private UndirectedGraph outputNetwork;

	/**
	 * The training dataset used in the algorithm to learn the independence
	 * structure
	 */
	private Dataset dataset;

	/**
	 * Equivalent sample size that adjust the strength of the prior
	 * */
	public double equivalentSampleSize = 1;

	/**
	 * Determines if the score for structures can be computed incrementally from
	 * other structures
	 */
	public boolean incremental = true;

	/**
	 * The ADTree cache
	 */
	private ADNode adTree;

	/**
	 * Boolean value which specifies if the ADTree cache must be used
	 */
	private boolean useADTree = true;

	public MPLInterIAMBandHillClimbing(int n, Dataset dataset,
			boolean incremental, boolean useADTree, ADNode adTree) {
		this.n = dataset.getNumberOfVariables();
		outputNetwork = new UndirectedGraph(n);
		this.dataset = dataset;
		this.adTree = adTree;
		this.useADTree = useADTree;
	}

	/**
	 * Run the algorithm
	 */
	public void run() {
		/*
		 * First phase: learning each Markov blanket with interleaved IAMB
		 * strategy and local MPL score
		 */
		List<List<Integer>> blankets = new ArrayList<List<Integer>>();
		for (int x = 0; x < n; x++) {
			blankets.add(interIAMB(x));
		}

		/*
		 * Second phase: learning a structure with hill climbing and global MPL
		 * score, only in the space of inconsistent edges
		 */
		outputNetwork = graphHillClimbing(blankets);
	}

	/**
	 * Learns the Markov blanket of a variable with interleaved IAMB strategy
	 * and local MPL score. Follows the pseudocode of Algorithm 1 of (Pensar, et
	 * al; 2014). Pensar, Johan, et al.
	 * "Marginal Pseudo-Likelihood Learning of Markov Network structures." arXiv
	 * preprint arXiv:1401.4988 (2014).
	 * 
	 * @param x
	 *            The target variable
	 * @return The Markov blanket of x
	 */
	private List<Integer> interIAMB(int x) {
		List<Integer> mb = new ArrayList<Integer>();
		List<Integer> potential_mb = new ArrayList<Integer>();

		boolean potentialMBHasChanged = true;
		while (potentialMBHasChanged) {

			List<Integer> complement = dataset.getDomainVariables();
			complement.remove((Integer) x);
			for (int var : mb)
				complement.remove((Integer) var);
			// complement.removeAll(mb);

			mb = new ArrayList<Integer>(potential_mb);

			potentialMBHasChanged = false;
			for (Integer i : complement) {
				if (!mb.contains(i)) {
					List<Integer> mbUnionI = new ArrayList<Integer>(mb);
					mbUnionI.add(i);

					if (logMPL(x, mbUnionI) > logMPL(x, potential_mb)) {
						potential_mb = new ArrayList<Integer>(mbUnionI);
						potentialMBHasChanged = true;
					}
				}
			}

			boolean removalIncreasesScore = true;
			while (removalIncreasesScore && potential_mb.size() > 2) {
				removalIncreasesScore = false;
				mb = new ArrayList<Integer>(potential_mb);
				for (Integer i : mb) {
					List<Integer> mbWithoutI = new ArrayList<Integer>(mb);
					mbWithoutI.remove(i);

					if (logMPL(x, mbWithoutI) > logMPL(x, potential_mb)) {
						potential_mb = new ArrayList<Integer>(mbWithoutI);
						removalIncreasesScore = true;
					}
				}
			}

		}

		return potential_mb;
	}

	/**
	 * @param i
	 * @param blanketOfX_i
	 * @return
	 */
	private double logMPL(int i, List<Integer> blanketOfX_i) {
		double variableScore = 0;

		if (!useADTree) {
			variableScore = computeFromData(i, blanketOfX_i);
		} else {
			variableScore = computeFromADTree(i, blanketOfX_i);
		}

		return variableScore;
	}

	private double computeFromData(int i, List<Integer> blanketOfX_i) {
		double variableScore = 0;
		// permutations of current variables X_i
		List<Integer> x_i = new ArrayList<Integer>();
		x_i.add(i);
		List<int[]> permutationsOfX_i = Permutations.getPermutations(x_i,
				dataset.getCardinalities());

		// permutations of the Markov blanket of the current variables X_i
		List<int[]> permutationsOfBlanketOfX_i = Permutations.getPermutations(
				blanketOfX_i, dataset.getCardinalities());

		for (int[] permutationOfBlanketOfX_i : permutationsOfBlanketOfX_i) {
			int countsOfPermutationOfBlanketOfX_i = 0;
			double alpha_i = 0;
			for (int[] permutationOfX_i : permutationsOfX_i) {

				List<Integer> variables = new ArrayList<Integer>();
				variables.addAll(blanketOfX_i);
				variables.addAll(x_i);
				int[] values = new int[variables.size()];
				System.arraycopy(permutationOfBlanketOfX_i, 0, values, 0,
						permutationOfBlanketOfX_i.length);
				values[permutationOfBlanketOfX_i.length] = permutationOfX_i[0];
				int counts = dataset.getCounts(variables, values);
				double alpha = equivalentSampleSize
						/ (permutationsOfBlanketOfX_i.size() * permutationsOfX_i
								.size());

				countsOfPermutationOfBlanketOfX_i += counts;
				alpha_i += alpha;

				variableScore += Gamma.logGamma(alpha + counts)
						- Gamma.logGamma(alpha);

			}
			variableScore += Gamma.logGamma(alpha_i)
					- Gamma.logGamma(alpha_i
							+ countsOfPermutationOfBlanketOfX_i);
		}
		return variableScore;
	}

	private double computeFromADTree(int i, List<Integer> blanketOfX_i) {
		double variableScore = 0;

		// permutations of current variables X_i
		List<Integer> x_i = new ArrayList<Integer>();
		x_i.add(i);
		List<int[]> permutationsOfX_i = Permutations.getPermutations(x_i,
				dataset.getCardinalities());

		long numberOfPermutations = (long) Math.pow(dataset.getCardinalities(0), blanketOfX_i.size());
		
		List<Integer> allVariables = new ArrayList<Integer>(x_i);
		allVariables.addAll(blanketOfX_i);

		int[] counts = getContingencyTable(allVariables);

		double alpha_i = 0;
		for(int j = 0; j<permutationsOfX_i.size();++j){
			alpha_i += equivalentSampleSize
					/ (numberOfPermutations * permutationsOfX_i
							.size());
		}

		for (int j = 0; j < counts.length; ++j) {
			double alpha = equivalentSampleSize
					/ (numberOfPermutations * permutationsOfX_i
							.size());
			variableScore += Gamma.logGamma(alpha + counts[j])
					- Gamma.logGamma(alpha);
		}

		counts = getContingencyTable(blanketOfX_i);
		for (int j = 0; j < counts.length; ++j) {
			variableScore += Gamma.logGamma(alpha_i)
					- Gamma.logGamma(alpha_i + counts[j]);

		}

		return variableScore;
	}

	/**
	 * Construct a contingency table from the ADTree cache
	 * @param variables The variables involved
	 * @return A contingency table from the ADTree cache
	 */
	private int[] getContingencyTable(List<Integer> variables) {
		BitSet vars = new BitSet(n);
		for (int x : variables)
			vars.set(x);
		int[] cards = new int[n];
		for (int x = 0; x < dataset.getNumberOfVariables(); ++x)
			cards[x] = 2;
		return adTree.makeContab(vars, new HashMap<Integer, Integer>(), cards,
				dataset);
	}

	/**
	 * Learns the independence structure with a hill climbing search, using as
	 * objective function the global MPL score, only in the space of
	 * inconsistent edges. Follows the pseudocode of Algorithm 2 of (Pensar, et
	 * al; 2014). Pensar, Johan, et al.
	 * "Marginal Pseudo-Likelihood Learning of Markov Network structures." arXiv
	 * preprint arXiv:1401.4988 (2014).
	 * 
	 * @param blankets
	 *            The Markov blankets learned previously
	 * @return A structure learned with a hill climbing search and MPL scoring
	 *         function
	 */
	private UndirectedGraph graphHillClimbing(List<List<Integer>> blankets) {
		UndirectedGraph g = new UndirectedGraph(n);
		UndirectedGraph potentialG = new UndirectedGraph(n);

		GraphScorable scoringFunction = new MPL();

		boolean potentialGHasChanged = true;
		while (potentialGHasChanged) {

			g = new UndirectedGraph(potentialG);
			potentialGHasChanged = false;

			UndirectedGraphScore potentialGScore = scoringFunction
					.computeScore(g, dataset, useADTree, adTree);

			for (int i = 0; i < n - 1; ++i) {
				for (int j = i + 1; j < n; ++j) {
					if (blankets.get(i).contains(j)
							|| blankets.get(j).contains(i)) {
						UndirectedGraph neighbor = new UndirectedGraph(g);
						neighbor.flipEdge(i, j);

						UndirectedGraphScore neighborScore = scoringFunction.computeScore(neighbor, dataset,
										useADTree, adTree);

						if (neighborScore.getTotalScore() > potentialGScore
								.getTotalScore()) {
							potentialG = new UndirectedGraph(neighbor);
							potentialGHasChanged = true;
						}
					}
				}
			}

		}

		return potentialG;
	}

	public UndirectedGraph getOutputNetwork() {
		return outputNetwork;
	}

}
