/**
 * 
 */
package edu.berkeley.nlp.discPCFG;

import edu.berkeley.nlp.PCFGLA.BinaryRule;
import edu.berkeley.nlp.PCFGLA.Grammar;
import edu.berkeley.nlp.PCFGLA.HierarchicalBinaryRule;
import edu.berkeley.nlp.PCFGLA.HierarchicalGrammar;
import edu.berkeley.nlp.PCFGLA.HierarchicalLexicon;
import edu.berkeley.nlp.PCFGLA.HierarchicalUnaryRule;
import edu.berkeley.nlp.PCFGLA.SimpleLexicon;
import edu.berkeley.nlp.PCFGLA.SpanPredictor;
import edu.berkeley.nlp.PCFGLA.UnaryRule;
import edu.berkeley.nlp.math.SloppyMath;
import edu.berkeley.nlp.syntax.StateSet;
import edu.berkeley.nlp.util.ArrayUtil;

/**
 * similar to cascading linearizer but doesnt compute the grammars explicitly
 * instead uses hierarchical rules and merges back unused splits
 * 
 * @author petrov
 * 
 */
public class HierarchicalLinearizer extends DefaultLinearizer {

	private static final long serialVersionUID = 1L;

	HierarchicalGrammar grammar;
	HierarchicalLexicon lexicon;

	int finalLevel;
	int[][] lexiconMapping;
	int[][][] unaryMapping;
	int[][][][] binaryMapping;

	public HierarchicalLinearizer() {
	}

	/**
	 * @param grammar
	 * @param lexicon
	 */
	public HierarchicalLinearizer(Grammar grammar, SimpleLexicon lexicon,
			SpanPredictor sp, int fLevel) {
		this.grammar = (HierarchicalGrammar) grammar;
		this.lexicon = (HierarchicalLexicon) lexicon;
		this.spanPredictor = sp;
		this.finalLevel = fLevel;
		this.nSubstates = (int) ArrayUtil.max(grammar.numSubStates);
		init();
		computeMappings();
	}

	protected void computeMappings() {
		lexiconMapping = new int[finalLevel + 1][nSubstates];
		unaryMapping = new int[finalLevel + 1][nSubstates][nSubstates];
		binaryMapping = new int[finalLevel + 1][nSubstates][nSubstates][nSubstates];

		int[] divisors = new int[finalLevel + 1];
		for (int i = 0; i <= finalLevel; i++) {
			divisors[i] = (int) Math.pow(2, finalLevel - i);
		}

		for (int level = 1; level <= finalLevel; level++) {
			int div = divisors[level];
			int l = (int) Math.pow(2, level);
			int[][] tmpU = new int[l][l];
			int[][][] tmpB = new int[l][l][l];
			int indU = 0, indB = 0;
			for (int i = 0; i < l; i++) {
				for (int j = 0; j < l; j++) {
					tmpU[i][j] = indU++;
					for (int k = 0; k < l; k++) {
						tmpB[i][j][k] = indB++;
					}
				}
			}
			for (int i = 0; i < nSubstates; i++) {
				lexiconMapping[level][i] = i / div;
				for (int j = 0; j < nSubstates; j++) {
					unaryMapping[level][i][j] = tmpU[i / div][j / div];
					for (int k = 0; k < nSubstates; k++) {
						binaryMapping[level][i][j][k] = tmpB[i / div][j / div][k
								/ div];
					}
				}
			}
		}
	}

	// public void delinearizeSpanPredictor(double[] logProbs) {
	//
	// }

	@Override
	public void delinearizeGrammar(double[] probs) {
		int nDangerous = 0;
		for (BinaryRule bRule : grammar.binaryRuleMap.keySet()) {
			HierarchicalBinaryRule hRule = (HierarchicalBinaryRule) bRule;
			int ind = hRule.identifier;// startIndex[ruleIndexer.indexOf(hRule)];
			double[][][] scores = hRule.getLastLevel();
			for (int j = 0; j < scores.length; j++) {
				for (int k = 0; k < scores[j].length; k++) {
					if (scores[j][k] != null) {
						for (int l = 0; l < scores[j][k].length; l++) {
							double val = probs[ind++];
							if (SloppyMath.isVeryDangerous(val)) {
								nDangerous++;
								continue;
							}
							scores[j][k][l] = val;
						}
					}
				}
			}
		}
		if (nDangerous > 0)
			System.out
					.println("Left "
							+ nDangerous
							+ " binary rule weights unchanged since the proposed weight was dangerous.");

		nDangerous = 0;
		for (UnaryRule uRule : grammar.unaryRuleMap.keySet()) {
			HierarchicalUnaryRule hRule = (HierarchicalUnaryRule) uRule;
			int ind = hRule.identifier;// startIndex[ruleIndexer.indexOf(hRule)];
			if (uRule.childState == uRule.parentState)
				continue;
			double[][] scores = hRule.getLastLevel();
			for (int j = 0; j < scores.length; j++) {
				if (scores[j] != null) {
					for (int k = 0; k < scores[j].length; k++) {
						double val = probs[ind++];
						if (SloppyMath.isVeryDangerous(val)) {
							nDangerous++;
							continue;
						}
						scores[j][k] = val;
					}
				}
			}
		}
		if (nDangerous > 0)
			System.out
					.println("Left "
							+ nDangerous
							+ " unary rule weights unchanged since the proposed weight was dangerous.");

		grammar.explicitlyComputeScores(finalLevel);
		grammar.closedSumRulesWithParent = grammar.closedViterbiRulesWithParent = grammar.unaryRulesWithParent;
		grammar.closedSumRulesWithChild = grammar.closedViterbiRulesWithChild = grammar.unaryRulesWithC;
		// computePairsOfUnaries();
		grammar.clearUnaryIntermediates();
		grammar.makeCRArrays();
		// return grammar;
	}

	@Override
	public void delinearizeLexicon(double[] logProbs) {
		int nDangerous = 0;
		for (short tag = 0; tag < lexicon.hierarchicalScores.length; tag++) {
			for (int word = 0; word < lexicon.hierarchicalScores[tag].length; word++) {
				int index = linearIndex[tag][word];
				double[] vals = lexicon.getLastLevel(tag, word);
				for (int substate = 0; substate < vals.length; substate++) {
					double val = logProbs[index++];
					if (SloppyMath.isVeryDangerous(val)) {
						nDangerous++;
						continue;
					}
					vals[substate] = val;
				}
			}
		}
		if (nDangerous > 0)
			System.out
					.println("Left "
							+ nDangerous
							+ " lexicon weights unchanged since the proposed weight was dangerous.");
		lexicon.explicitlyComputeScores(finalLevel);
		// System.out.println(lexicon);
		// return lexicon;
	}

	@Override
	public double[] getLinearizedGrammar(boolean update) {
		if (update) {
			// int nRules = grammar.binaryRuleMap.size() +
			// grammar.unaryRuleMap.size();
			// startIndex = new int[nRules];

			nGrammarWeights = 0;
			for (BinaryRule bRule : grammar.binaryRuleMap.keySet()) {
				HierarchicalBinaryRule hRule = (HierarchicalBinaryRule) bRule;
				// ruleIndexer.add(hRule);
				if (!grammar.isGrammarTag[bRule.parentState]) {
					System.out.println("Incorrect grammar tag");
				}
				bRule.identifier = nGrammarWeights;
				double[][][] scores = hRule.getLastLevel();
				for (int j = 0; j < scores.length; j++) {
					for (int k = 0; k < scores[j].length; k++) {
						if (scores[j][k] != null) {
							nGrammarWeights += scores[j][k].length;
						}
					}
				}
			}
			for (UnaryRule uRule : grammar.unaryRuleMap.keySet()) {
				HierarchicalUnaryRule hRule = (HierarchicalUnaryRule) uRule;
				// ruleIndexer.add(hRule);
				// startIndex[ruleIndexer.indexOf(uRule)] = nGrammarWeights;
				uRule.identifier = nGrammarWeights;
				double[][] scores = hRule.getLastLevel();
				for (int j = 0; j < scores.length; j++) {
					if (scores[j] != null) {
						nGrammarWeights += scores[j].length;
					}
				}
			}
		}
		double[] logProbs = new double[nGrammarWeights];

		for (BinaryRule bRule : grammar.binaryRuleMap.keySet()) {
			HierarchicalBinaryRule hRule = (HierarchicalBinaryRule) bRule;
			int ind = hRule.identifier;// startIndex[ruleIndexer.indexOf(hRule)];
			double[][][] scores = hRule.getLastLevel();
			for (int j = 0; j < scores.length; j++) {
				for (int k = 0; k < scores[j].length; k++) {
					if (scores[j][k] != null) {
						for (int l = 0; l < scores[j][k].length; l++) {
							double val = scores[j][k][l];
							logProbs[ind++] = val;
						}
					}
				}
			}
		}

		for (UnaryRule uRule : grammar.unaryRuleMap.keySet()) {
			HierarchicalUnaryRule hRule = (HierarchicalUnaryRule) uRule;
			int ind = hRule.identifier;// startIndex[ruleIndexer.indexOf(hRule)];
			if (uRule.childState == uRule.parentState)
				continue;
			double[][] scores = hRule.getLastLevel();
			for (int j = 0; j < scores.length; j++) {
				if (scores[j] != null) {
					for (int k = 0; k < scores[j].length; k++) {
						double val = scores[j][k];
						logProbs[ind++] = val;
					}
				}
			}
		}
		return logProbs;
	}

	@Override
	public double[] getLinearizedLexicon(boolean update) {
		if (update) {
			nLexiconWeights = 0;
			int[] substates = new int[finalLevel + 1];
			for (int i = 0; i <= finalLevel; i++)
				substates[i] = (int) Math.pow(2, i);
			for (short tag = 0; tag < lexicon.hierarchicalScores.length; tag++) {
				for (int word = 0; word < lexicon.hierarchicalScores[tag].length; word++) {
					nLexiconWeights += lexicon.getLastLevel(tag, word).length;
				}
			}
		}
		double[] logProbs = new double[nLexiconWeights];
		if (update)
			linearIndex = new int[lexicon.hierarchicalScores.length][];

		int index = 0;
		for (short tag = 0; tag < lexicon.hierarchicalScores.length; tag++) {
			if (update)
				linearIndex[tag] = new int[lexicon.hierarchicalScores[tag].length];
			for (int word = 0; word < lexicon.hierarchicalScores[tag].length; word++) {
				if (update)
					linearIndex[tag][word] = index + nGrammarWeights;
				double[] vals = lexicon.getLastLevel(tag, word);
				for (int substate = 0; substate < vals.length; substate++) {
					double val = vals[substate];
					logProbs[index++] = val;
				}
			}
		}
		if (index != logProbs.length)
			System.out.println("unequal length in lexicon");

		return logProbs;
	}

	@Override
	public int getLinearIndex(int globalWordIndex, int tag) {
		int tagSpecificWordIndex = lexicon.tagWordIndexer[tag]
				.indexOf(globalWordIndex);
		if (tagSpecificWordIndex == -1)
			return -1;
		return linearIndex[tag][tagSpecificWordIndex];
	}

	@Override
	public int dimension() {
		return nGrammarWeights + nLexiconWeights + nSpanWeights;
	}

	@Override
	public void increment(double[] counts, StateSet stateSet, int tag,
			double[] weights, boolean isGold) {
		int globalSigIndex = stateSet.sigIndex;
		if (globalSigIndex != -1) {
			int startIndexWord = getLinearIndex(globalSigIndex, tag);
			if (startIndexWord >= 0) {
				int finalLevel = lexicon.getFinalLevel(globalSigIndex, tag);
				for (int i = 0; i < nSubstates; i++) {
					if (isGold)
						counts[startIndexWord + lexiconMapping[finalLevel][i]] += weights[i];
					else
						counts[startIndexWord + lexiconMapping[finalLevel][i]] -= weights[i];
				}
			}
		}
		int globalWordIndex = stateSet.wordIndex;
		int startIndexWord = getLinearIndex(globalWordIndex, tag);
		if (startIndexWord >= 0) {
			int finalLevel = lexicon.getFinalLevel(globalWordIndex, tag);
			for (int i = 0; i < nSubstates; i++) {
				if (isGold)
					counts[startIndexWord + lexiconMapping[finalLevel][i]] += weights[i];
				else
					counts[startIndexWord + lexiconMapping[finalLevel][i]] -= weights[i];
				weights[i] = 0;
			}
		} else {
			for (int i = 0; i < nSubstates; i++) {
				weights[i] = 0;
			}
		}
	}

	@Override
	public void increment(double[] counts, UnaryRule rule, double[] weights,
			boolean isGold) {
		HierarchicalUnaryRule hr = (HierarchicalUnaryRule) rule;
		int thisStartIndex = hr.identifier;
		int finalLevel = hr.lastLevel;
		int curInd = 0;
		if (rule.parentState == 0) {
			for (int cp = 0; cp < nSubstates; cp++) {
				double val = weights[curInd];
				if (val > 0) {
					if (isGold)
						counts[thisStartIndex + lexiconMapping[finalLevel][cp]] += val;
					else
						counts[thisStartIndex + lexiconMapping[finalLevel][cp]] -= val;
					weights[curInd] = 0;
				}
				curInd++;
			}
			return;
		}

		for (int cp = 0; cp < nSubstates; cp++) {
			// if (scores[cp]==null) continue;
			for (int np = 0; np < nSubstates; np++) {
				double val = weights[curInd];
				if (val > 0) {
					if (isGold)
						counts[thisStartIndex
								+ unaryMapping[finalLevel][cp][np]] += val;
					else
						counts[thisStartIndex
								+ unaryMapping[finalLevel][cp][np]] -= val;
					weights[curInd] = 0;
				}
				curInd++;
			}
		}
	}

	@Override
	public void increment(double[] counts, BinaryRule rule, double[] weights,
			boolean isGold) {
		HierarchicalBinaryRule hr = (HierarchicalBinaryRule) rule;
		int thisStartIndex = hr.identifier;
		int finalLevel = hr.lastLevel;
		int curInd = 0;
		for (int lp = 0; lp < nSubstates; lp++) {
			for (int rp = 0; rp < nSubstates; rp++) {
				// if (scores[cp]==null) continue;
				for (int np = 0; np < nSubstates; np++) {
					double val = weights[curInd];
					if (val > 0) {
						if (isGold)
							counts[thisStartIndex
									+ binaryMapping[finalLevel][lp][rp][np]] += val;
						else
							counts[thisStartIndex
									+ binaryMapping[finalLevel][lp][rp][np]] -= val;
						weights[curInd] = 0;
					}
					curInd++;
				}
			}
		}
	}

	@Override
	public Grammar getGrammar() {
		return grammar;
	}

	@Override
	public SimpleLexicon getLexicon() {
		return lexicon;
	}

	@Override
	public SpanPredictor getSpanPredictor() {
		return spanPredictor;
	}

}
