package learn.struct.scoring.search;

import java.io.BufferedWriter;
import java.io.IOException;
import java.util.Random;

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

/**
 * @author fschluter
 *
 */
public class HillClimbing {
	/**
	 * 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;

	/**
	 * An object for computing the score of structures. Possible instances are:
	 * ib.ibmap.score.IBScore, ib.ibmap.score.RIBScore,
	 * 
	 */
	private GraphScorable scoringFunction;

	/**
	 * Hill climbing search can use the greedy operator for selecting
	 * heuristically the neighbor (true) or selecting naively the best neighbor
	 * one edge away
	 */
	private boolean greedy = true;

	/**
	 * Represents the score of the best structure found by the optimization
	 */
	private double bestScore;

	/**
	 * Writes debug info about the optimization
	 */
	private BufferedWriter writer = null;

	/**
	 * The true graph is known when it is used with synthetic data
	 */
	private UndirectedGraph trueGraph = null;

	/**
	 * The scoring function can be computed incrementally
	 */
	private boolean incremental = false;

	/**
	 * Number of ascents made in the local search
	 */
	private int ascents = 0;

	/**
	 * Number of tests made during the search
	 */
	private int testCounts = 0;

	/**
	 * Weighted number of tests made during the search
	 */
	private int weightedTestCounts = 0;

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

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

	/**
	 * Class constructor
	 * @param n
	 *            Number of variables of the domain
	 * @param test
	 *            Statististical independence test instance
	 * @param dataset
	 *            Dataset instance
	 * @param greedy2
	 */
	public HillClimbing(int n, Dataset dataset, GraphScorable scoringFunction,
			boolean greedy) {
		this.n = n;
		this.dataset = dataset;
		this.scoringFunction = scoringFunction;
		this.greedy = greedy;
	}

	/**
	 * Class constructor
	 * @param n
	 *            Number of variables of the domain
	 * @param test
	 *            Statististical independence test instance
	 * @param dataset
	 *            Dataset instance
	 * @param greedy2
	 * 
	 * 
	 * @param useADTree Boolean value which specifies if the ADTree cache must be used 
	 * @param aDTree The ADTree cache 
	 */
	public HillClimbing(int n, Dataset dataset, GraphScorable scoringFunction,
			boolean greedy, boolean useADTree, ADNode aDTree) {
		this(n, dataset, scoringFunction, greedy);
		this.useADTree = useADTree;
		this.adTree = aDTree;
	}

	/**
	 * Simple implementation of pseudocode shown in (Schluter, et al., 2013)
	 * 
	 * @throws IOException
	 */
	public void run() throws IOException {
		// Computing the score of the empty structure
		UndirectedGraph g = new UndirectedGraph(n);
		run(g);
	}

	public void run(int randomRestarts, int seed) throws IOException {
		UndirectedGraphScore bestScoreInRandomRestart = new UndirectedGraphScore(
				new UndirectedGraph(n));
		bestScoreInRandomRestart.setTotalScore(Integer.MIN_VALUE);
		Random r = new Random(seed);
		for (int i = 0; i < randomRestarts; ++i) {
			UndirectedGraph initialStructure = UndirectedGraph.getRandomGraph(
					n, r);
			UndirectedGraphScore iScore = run(initialStructure);
			if (iScore.getTotalScore() >= bestScoreInRandomRestart
					.getTotalScore()) {
				bestScoreInRandomRestart = iScore;
			}
		}
		outputNetwork = bestScoreInRandomRestart.getUndirectedGraph();
		bestScore = bestScoreInRandomRestart.getTotalScore();
	}

	public UndirectedGraphScore run(UndirectedGraph initialStructure)
			throws IOException {
		ascents = 0;
		UndirectedGraphScore gScore = !useADTree ? scoringFunction.computeScore(
				initialStructure, dataset) : scoringFunction.computeScore(
						initialStructure, dataset, useADTree, adTree);
		testCounts += gScore.getTestCounts();
		weightedTestCounts += gScore.getWeightedTestCounts();
		double currentScore = gScore.getTotalScore();
		do {
			// Selecting the best neighbor structure
			UndirectedGraphScore neighbor = selectNeighbor(initialStructure,
					gScore);
			if (neighbor.getTotalScore() <= currentScore) {
				/* Local maxima reached. Termination condition. */
				if (writer != null) {
					debug(initialStructure, currentScore);
					writer.close();
				}
				outputNetwork = initialStructure;
				bestScore = currentScore;
				return neighbor;
			}
			initialStructure = neighbor.getUndirectedGraph();
			currentScore = neighbor.getTotalScore();
			gScore = neighbor;
			if (writer != null) {
				debug(initialStructure, currentScore);
			}

			++ascents;

		} while (true);
	}

	/**
	 * Writes debug info about the optimization process
	 * 
	 * @param g
	 * @param currentScore
	 * @throws IOException
	 */
	private void debug(UndirectedGraph g, double currentScore)
			throws IOException {
		int hd = -1;
		int fp = -1;
		int fn = -1;

		if (trueGraph != null) {
			hd = UndirectedGraphsQualityMeasures.hammingDistance(trueGraph, g);
			fp = UndirectedGraphsQualityMeasures.falsePositives(trueGraph, g);
			fn = UndirectedGraphsQualityMeasures.falseNegatives(trueGraph, g);
		}
		writer.write(dataset.getList().size() + " "
				+ scoringFunction.getScoringFunctionName() + " "
				+ g.getBitsetString() + " " + currentScore + " " + hd + " "
				+ fp + " " + fn + "\n");
	}

	/**
	 * Returns a the neighbour structure flippng the edge that maximizes the
	 * IB-Score
	 * 
	 * @param g
	 * @param gScore
	 * @return
	 */
	protected UndirectedGraphScore selectNeighbor(UndirectedGraph g,
			UndirectedGraphScore gScore) {
		if (greedy)
			return selectGreedyNeighbor(g, gScore);
		else
			return selectLocalNeighbor(g, gScore);
	}

	/**
	 * @param g
	 * @return
	 */
	private UndirectedGraphScore selectGreedyNeighbor(UndirectedGraph g,
			UndirectedGraphScore gScore) {
		int bestX = 0, bestY = 0;
		double minProduct = Integer.MAX_VALUE;
		// Obtaining scores of g, computed in previous iteration
		double[][] markovBlanketClosureScores = gScore
				.getMarkovBlanketClosureScores();
		for (int i = 0; i < n - 1; ++i) {
			for (int j = i + 1; j < n; ++j) {
				double sum = markovBlanketClosureScores[i][j]
						+ markovBlanketClosureScores[j][i];
				// Obtaining the pair of tests whose product adds the lowest
				// value to the current score
				if (sum < minProduct) {
					minProduct = sum;
					bestX = i;
					bestY = j;
				}
			}
		}
		// The best neighbor is approximated by flipping the worst edge
		UndirectedGraph bestNeighbor = new UndirectedGraph(g);
		bestNeighbor.flipEdge(bestX, bestY);
		
		UndirectedGraphScore bestNeighborScore = null;
		if(!useADTree){
			bestNeighborScore = incremental ? scoringFunction
					.computeScore(bestNeighbor, gScore, dataset) : scoringFunction
					.computeScore(bestNeighbor, dataset);
		}else{
			bestNeighborScore = scoringFunction.computeScore(
					bestNeighbor, dataset, useADTree, adTree);
		}
		return bestNeighborScore;
	}

	/**
	 * @param g
	 * @param gScore
	 * @return
	 */
	private UndirectedGraphScore selectLocalNeighbor(UndirectedGraph g,
			UndirectedGraphScore gScore) {
		double maxScore = Integer.MIN_VALUE;
		UndirectedGraphScore neighborWithBestScore = null;
		for (int i = 0; i < n - 1; ++i) {
			for (int j = i + 1; j < n; ++j) {
				UndirectedGraph neighbor = new UndirectedGraph(g);
				neighbor.flipEdge(i, j);

				UndirectedGraphScore neighborScore = null;
				if(!useADTree){
					neighborScore = incremental ? scoringFunction
							.computeScore(neighbor, gScore, dataset) : scoringFunction
							.computeScore(neighbor, dataset);
				}else{
					neighborScore = scoringFunction.computeScore(
							neighbor, dataset, useADTree, adTree);
				}						
				if (neighborScore.getTotalScore() > maxScore) {
					maxScore = neighborScore.getTotalScore();
					neighborWithBestScore = neighborScore;
				}
			}
		}
		return neighborWithBestScore;
	}

	/**
	 * @return
	 */
	public UndirectedGraph getOutputNetwork() {
		return outputNetwork;
	}

	/**
	 * @return
	 */
	public double getBestScore() {
		return bestScore;
	}

	public void setDebugFile(BufferedWriter writer) {
		this.writer = writer;
	}

	public void setTrueGraph(UndirectedGraph trueGraph) {
		this.trueGraph = trueGraph;
	}

	public void setIncremental(boolean incremental) {
		this.incremental = incremental;
	}

	public int getNumberOfAscents() {
		return ascents;
	}

	public int getWeightedTestCounts() {
		return weightedTestCounts;
	}

	public int getTestCounts() {
		return testCounts;
	}
}
