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

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

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


/**
 * Implements the Marginal Pseudo Likelihood scoring function for
 * Markov networks structure discovery, proposed by (Pensar et al, 2014).
 * 
 * <!-- technical-plaintext-start --> Pensar, Johan, et al. 
 * "Marginal Pseudo-Likelihood Learning of Markov Network structures." 
 * arXiv preprint arXiv:1401.4988 (2014). <!-- * technical-plaintext-end -->
 * <p/>
 */
public class MPL implements GraphScorable {
	/* Equivalent sample size that adjust the strength of the prior */
	public double equivalentSampleSize = 1;
	
	/**
	 * Marginal Pseudo Likelihood scoring function proposed at (Pensar et. al., 2014) 
	 * 
	 * @param g
	 *            The structure we want to score
	 * @return The score of the structure
	 */
	public UndirectedGraphScore computeScore(UndirectedGraph g, Dataset dataset) {
		double score = 0;
		UndirectedGraphScore gScore = new UndirectedGraphScore(g);
		
		/* for each variable of the domain */ 
		for (int i = 0; i < g.getNumberOfNodes(); ++i) {
			double variableScore = computeVariableScoreFromData(g, dataset, i);
			gScore.setVariableScore(i, variableScore);
			score += variableScore;
		}

		gScore.setTotalScore(score);
		return gScore;
	}


	/* (non-Javadoc)
	 * @see learn.struct.scoring.GraphScorable#computeScore(rep.UndirectedGraph, learn.struct.scoring.UndirectedGraphScore, dataset.Dataset)
	 */
	@Override
	public UndirectedGraphScore computeScore(UndirectedGraph g,
			UndirectedGraphScore incremental, Dataset dataset) {
		UndirectedGraphScore gScore = new UndirectedGraphScore(g);

		boolean[] differences = g.getVariablesWithDifferentBlanket(incremental
				.getUndirectedGraph());
		double score = 0;
		for (int i = 0; i < differences.length; i++) {
			double variableScore = 0;
			if (differences[i]) {
				variableScore = computeVariableScoreFromData(g, dataset, i);
			} else {
				variableScore = incremental.getVariableScore(i);
			}
			gScore.setVariableScore(i, variableScore);
			score += variableScore; 
		}
		gScore.setTotalScore(score);
		return gScore;
	}

	/* (non-Javadoc)
	 * @see learn.struct.scoring.GraphScorable#computeScore(rep.UndirectedGraph, dataset.Dataset, boolean, dataset.adtree.ADNode)
	 */
	@Override
	public UndirectedGraphScore computeScore(UndirectedGraph g,
			Dataset dataset, boolean useADTree, ADNode adTree) {
		if(!useADTree){
			return computeScore(g, dataset);
		}else{
			double score = 0;
			UndirectedGraphScore gScore = new UndirectedGraphScore(g);
			
			/* for each variable of the domain */ 
			for (int i = 0; i < g.getNumberOfNodes(); ++i) {
				double variableScore = computeVariableScoreFromADTree(g, dataset, i,adTree);
				gScore.setVariableScore(i, variableScore);
				score += variableScore;
			}
			gScore.setTotalScore(score);
			return gScore;
		}
	}
	/**
	 * Computes the score by reading each time from the dataset
	 * @param g
	 * @param dataset
	 * @param i
	 * @return
	 */
	private double computeVariableScoreFromData(UndirectedGraph g, Dataset dataset, int 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<Integer> blanketOfX_i = g.getBlanketOf(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[permutationOfBlanketOfX_i.length + 1];
				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;
	}
	
	/**
	 * Computes the score by reading each time from an ADTree cache
	 * @param g
	 * @param dataset
	 * @param i
	 * @param adTree
	 * @return
	 */
	private double computeVariableScoreFromADTree(UndirectedGraph g, Dataset dataset, int i, ADNode adTree) {
		double variableScore = 0;
		List<Integer> blanketOfX_i = g.getBlanketOf(i);
		// 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());
		
		List<Integer> allVariables = new ArrayList<Integer>(x_i);
		allVariables.addAll(blanketOfX_i);
		
		int[] counts = getContingencyTable(allVariables, adTree,g,dataset);
		
		double alpha_i = 0;
			for (int j=0;j<permutationsOfX_i.size();++j) {
				alpha_i += equivalentSampleSize
						/ (permutationsOfBlanketOfX_i.size() * permutationsOfX_i
								.size());
			}
		
		for(int j = 0 ; j < counts.length ; ++j){
			double alpha = equivalentSampleSize
					/ (permutationsOfBlanketOfX_i.size() * permutationsOfX_i
							.size());
			variableScore += Gamma.logGamma(alpha + counts[j])
					- Gamma.logGamma(alpha);
		}
		
		counts = getContingencyTable(blanketOfX_i, adTree,g,dataset );
		for(int j = 0 ; j < counts.length ; ++j){
			variableScore += Gamma.logGamma(alpha_i)
					- Gamma.logGamma(alpha_i
							+ counts[j]);

		}
		
		return variableScore;
	}

	/**
	 * @param variables
	 * @param adTree
	 * @param g
	 * @param dataset
	 * @return
	 */
	private int[] getContingencyTable(List<Integer> variables, ADNode adTree, UndirectedGraph g, Dataset dataset) {
		int n = g.getNumberOfNodes();
		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);
	}
	
	
	
	/* (non-Javadoc)
	 * @see learn.struct.scoring.GraphScorable#getScoringFunctionName()
	 */
	@Override
	public String getScoringFunctionName() {
		return "mpl";
	}

}
