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

import java.util.ArrayList;
import java.util.List;

import learn.struct.it.BayesianTest;
import learn.struct.it.IndependenceTriplet;
import learn.struct.scoring.GraphScorable;
import learn.struct.scoring.UndirectedGraphScore;
import rep.UndirectedGraph;
import dataset.Dataset;
import dataset.adtree.ADNode;

/**
 * Implements the IBMAP scoring function for
 * Markov networks structure discovery, proposed by (Schluter et al, 2014).
 * 
 * <!-- technical-plaintext-start --> Schlüter, Federico, Facundo Bromberg, and Alejandro Edera. 
 * "The IBMAP approach for Markov network structure learning." 
 * Annals of Mathematics and Artificial Intelligence 72.3-4 (2014): 197-223. 
 * <!-- * technical-plaintext-end -->
 * <p/>
 */
public class IBScore implements GraphScorable {

	private boolean onlyIndependences = false;
	private BayesianTest test; 

	public IBScore() {
		this(false);
		test = new BayesianTest();
	}
	
	public IBScore(boolean onlyIndependences) {
		this.onlyIndependences = onlyIndependences;
		test = new BayesianTest();
	}  

	/**
	 * Computing the IB-score of a graph g. This is the product of the
	 * posteriors of all the assertions in the Markov blanket closure
	 * 
	 * @param g
	 *            The structure we want to score
	 * @return The score of the structure
	 */
	public UndirectedGraphScore computeScore(UndirectedGraph g, Dataset dataset) {
		
		UndirectedGraphScore gScore = new UndirectedGraphScore(g);
		// ArrayList<IndependenceTriplet> closure = getClosureOfStructure(g);
		int n = g.getNumberOfNodes();
		double[][] markovBlanketClosureScores = new double[n][n];
		double score = 0;

		long testCounts = 0;
		long weightedTestCounts = 0;
		for (int i = 0; i < g.getNumberOfNodes(); ++i) {
			ArrayList<IndependenceTriplet> closureOfX = getClosureOfVariable(g,
					i);
			for (IndependenceTriplet triplet : closureOfX) {
				test.independent(dataset, triplet);
				++testCounts;
				weightedTestCounts+=triplet.getX().size()+triplet.getY().size()+triplet.getZ().size();
				int x = triplet.getX().get(0);
				int y = triplet.getY().get(0);
				double[] logProbs = test.getLogProbs();
				if (g.existEdge(x, y)) {
					// log probability of dependence model
					score += logProbs[1];
					markovBlanketClosureScores[x][y] = logProbs[1];
				} else {
					// log probability of independence model
					score += logProbs[0];
					markovBlanketClosureScores[x][y] = logProbs[0];
				}
			}

		}

		gScore.setMarkovBlanketClosureScores(markovBlanketClosureScores);
		gScore.setTotalScore(score);
		gScore.setTestCounts(testCounts);
		gScore.setWeightedTestCounts(weightedTestCounts);
		return gScore;
	}

	private ArrayList<IndependenceTriplet> getClosureOfVariable(
			UndirectedGraph g, int x) {
		ArrayList<IndependenceTriplet> closureOfX = new ArrayList<IndependenceTriplet>();
		for (int y = 0; y < g.getNumberOfNodes(); ++y) {
			if (x != y) {
				List<Integer> blanketOfX = g.getBlanketOf(x);
				blanketOfX.remove((Integer) y);
				IndependenceTriplet triplet = new IndependenceTriplet(x, y,
						blanketOfX);

				/*
				 * when onlyIndependences is true, the closure only contains
				 * independence assertions when onlyIndependences is false, the
				 * closure contains independence and dependence assertions
				 */
				if (!onlyIndependences
						|| (onlyIndependences && !g.existEdge(x, y))) {
					closureOfX.add(triplet);
				}
			}
		}
		return closureOfX;
	}

	@Override
	public UndirectedGraphScore computeScore(UndirectedGraph g,
			UndirectedGraphScore incremental, Dataset dataset) {
		int n = g.getNumberOfNodes();
		UndirectedGraphScore gScore = new UndirectedGraphScore(g);
		boolean[] differences = g.getVariablesWithDifferentBlanket(incremental
				.getUndirectedGraph());
		double[][] markovBlanketClosureScores = new double[n][n];
		double score = 0;
		long testCounts = 0;
		long weightedTestCounts = 0;

		for (int i = 0; i < differences.length; i++) {
			if (differences[i]) {
				ArrayList<IndependenceTriplet> closureOfVariable = getClosureOfVariable(
						g, i);
				for (IndependenceTriplet triplet : closureOfVariable) {
					int x = triplet.getX().get(0);
					int y = triplet.getY().get(0);
					test.independent(dataset, triplet);
					++testCounts;
					weightedTestCounts+=triplet.getX().size()+triplet.getY().size()+triplet.getZ().size();

					double[] logProbs = test.getLogProbs();
					if (g.existEdge(x, y)) {
						// log probability of dependence model
						score += logProbs[1];
						markovBlanketClosureScores[x][y] = logProbs[1];
					} else {
						// log probability of independence model
						score += logProbs[0];
						markovBlanketClosureScores[x][y] = logProbs[0];
					}
				}
			} else {
				for (int j = 0; j < n; j++) {
					score += incremental.getMarkovBlanketClosureScores()[i][j];
					markovBlanketClosureScores[i][j] = incremental
							.getMarkovBlanketClosureScores()[i][j];
				}
			}
		}
		gScore.setMarkovBlanketClosureScores(markovBlanketClosureScores);
		gScore.setTotalScore(score);
		gScore.setTestCounts(testCounts);
		gScore.setWeightedTestCounts(weightedTestCounts);

		return gScore;
	}

	public void setTest(BayesianTest test) {
		this.test = test;
	}

	@Override
	public String getScoringFunctionName() {
		return "ibscore";
	}

	@Override
	public UndirectedGraphScore computeScore(UndirectedGraph g,
			Dataset dataset, boolean useADTree, ADNode aDTree) {
		this.test.setUseADTree(useADTree);
		this.test.setADTree(aDTree);
		return computeScore(g, dataset); //cache is implemented in the Bayesian test class

	}

	
}
