content: package edu.berkeley.nlp.lm.io;
  
  import java.io.File;
  import java.io.PrintWriter;
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Random;
  
  import edu.berkeley.nlp.lm.ArrayEncodedNgramLanguageModel;
  import edu.berkeley.nlp.lm.ConfigOptions;
  import edu.berkeley.nlp.lm.ContextEncodedNgramLanguageModel.LmContextInfo;
  import edu.berkeley.nlp.lm.WordIndexer;
  import edu.berkeley.nlp.lm.collections.Counter;
  import edu.berkeley.nlp.lm.map.HashNgramMap;
  import edu.berkeley.nlp.lm.map.NgramMap.Entry;
  import edu.berkeley.nlp.lm.util.Logger;
  import edu.berkeley.nlp.lm.util.LongRef;
  import edu.berkeley.nlp.lm.util.StrUtils;
  import edu.berkeley.nlp.lm.values.KneserNeyCountValueContainer;
  import edu.berkeley.nlp.lm.values.KneserNeyCountValueContainer.KneserNeyCounts;
  import edu.berkeley.nlp.lm.values.ProbBackoffPair;
  
  /**
   * Class for producing a Kneser-Ney language model in ARPA format from raw text.
   * 
   * Confusingly, this class is both a {@link LmReaderCallback} (called from
   * {@link TextReader}, which reads plain text), and a {@link LmReader}, which
   * "reads" counts and produces Kneser-Ney probabilities and backoffs and passes
   * them on an {@link ArpaLmReaderCallback}
   * 
   * @author adampauls
   * 
   * @param <W>
   */
  public class KneserNeyLmReaderCallback<W> implements NgramOrderedLmReaderCallback<LongRef>, LmReader<ProbBackoffPair, ArpaLmReaderCallback<ProbBackoffPair>>,
  	ArrayEncodedNgramLanguageModel<W>, Serializable
  {
  
  	//	from http://www-speech.sri.com/projects/srilm/manpages/ngram-discount.7.html
  
  	//	p(a_z) = g(a_z) + bow(a_) p(_z)  ; Eqn.4
  	//
  	//	Let Z1 be the set {z: c(a_z) > 0}. For highest order N-grams we have:
  	//
  	//		g(a_z)  = max(0, c(a_z) - D) / c(a_)
  	//		bow(a_) = 1 - Sum_Z1 g(a_z)
  	//		        = 1 - Sum_Z1 c(a_z) / c(a_) + Sum_Z1 D / c(a_)
  	//		        = D n(a_*) / c(a_)
  	//
  	//	Let Z2 be the set {z: n(*_z) > 0}. For lower order N-grams we have:
  	//
  	//		g(_z)  = max(0, n(*_z) - D) / n(*_*)
  	//		bow(_) = 1 - Sum_Z2 g(_z)
  	//		       = 1 - Sum_Z2 n(*_z) / n(*_*) + Sum_Z2 D / n(*_*)
  	//		       = D n(_*) / n(*_*)
  	//
  
  	/**
  	 * 
  	 */
  	protected static final long serialVersionUID = 1L;
  
  	protected static final float DEFAULT_DISCOUNT = 0.75f;
  
  	protected final int lmOrder;
  
  	/**
  	 * 
  	 */
  
  	/**
  	 * 
  	 * This array represents the discount used for each ngram order.
  	 * 
  	 * The original Kneser-Ney discounting (-ukndiscount) uses one discounting
  	 * constant for each N-gram order. These constants are estimated as
  	 * 
  	 * D = n1 / (n1 + 2*n2)
  	 * 
  	 * where n1 and n2 are the total number of N-grams with exactly one and two
  	 * counts, respectively.
  	 * 
  	 * For simplicity, our code just uses a constant discount for each order of
  	 * 0.75. However, other discounts can be specified.
  	 */
  
  	protected final WordIndexer<W> wordIndexer;
  
  	protected final HashNgramMap<KneserNeyCounts> ngrams;
  
  	protected final ConfigOptions opts;
  
  	protected final int startIndex;
  
  	/**
  	 * 
  	 * @param wordIndexer
  	 * @param maxOrder
  	 * @param inputIsSentences
  	 *            If true, input n-grams are assumed to be sentences, and all
  	 *            sub-ngrams of up to order <code>maxOrder</code> are added. If
  	 *            false, input n-grams are assumed to be atomic.
  	 */
  	public KneserNeyLmReaderCallback(final WordIndexer<W> wordIndexer, final int maxOrder) {
  		this(wordIndexer, maxOrder, new ConfigOptions());
  	}
  
  	public KneserNeyLmReaderCallback(final WordIndexer<W> wordIndexer, final int maxOrder, final ConfigOptions opts) {
  		this.lmOrder = maxOrder;
  		this.startIndex = wordIndexer.getIndexPossiblyUnk(wordIndexer.getStartSymbol());
  
  		this.opts = opts;
  		double last = Double.NEGATIVE_INFINITY;
  		for (final double c : opts.kneserNeyMinCounts) {
  			if (c < last)
  				throw new IllegalArgumentException("Please ensure that ConfigOptions.kneserNeyMinCounts is monotonic (value was "
  					+ Arrays.toString(opts.kneserNeyMinCounts) + ")");
  			last = c;
  		}
  		this.wordIndexer = wordIndexer;
  		final KneserNeyCountValueContainer values = new KneserNeyCountValueContainer(lmOrder, startIndex);//, justLastWord);
  		ngrams = HashNgramMap.createExplicitWordHashNgramMap(values, opts, lmOrder, false);
  
  	}
  
  	public void call(final W[] ngram, final LongRef value) {
  		final int[] ints = new int[ngram.length];
  		for (int i = 0; i < ngram.length; ++i)
  			ints[i] = wordIndexer.getOrAddIndex(ngram[i]);
  		call(ints, 0, ints.length, value, "");
  	}
  
  	public void callJustLast(final W[] ngram, final LongRef value, final long[][] scratch) {
  		final int[] ints = new int[ngram.length];
  		for (int i = 0; i < ngram.length; ++i)
  			ints[i] = wordIndexer.getOrAddIndex(ngram[i]);
  		addNgram(ints, 0, ints.length, value, "", true, scratch);
  	}
  
  	@Override
  	public void call(final int[] ngram, final int startPos, final int endPos, final LongRef value, final String words) {
  		final long[][] prevOffsets = new long[lmOrder][endPos - startPos];
  		addNgram(ngram, startPos, endPos, value, words, false, prevOffsets);
  	}
  
  	/**
  	 * @param ngram
  	 * @param startPos
  	 * @param endPos
  	 * @param value
  	 * @param words
  	 */
  	public void addNgram(final int[] ngram, final int startPos, final int endPos, final LongRef value, @SuppressWarnings("unused") final String words, final boolean justLastWord,
  		final long[][] scratch) {
  
  		final KneserNeyCounts scratchCounts = new KneserNeyCounts();
  		ngrams.rehashIfNecessary(endPos - startPos);
  		for (int ngramOrder = 0; ngramOrder < lmOrder; ++ngramOrder) {
  			for (int i = startPos; i < endPos; ++i) {
  				int j = i + ngramOrder + 1;
  				if (j > endPos) continue;
  				scratchCounts.tokenCounts = value.value;
  				final long prevOffset = ngramOrder == 0 ? 0 : scratch[ngramOrder - 1][i];
  				final long suffixOffset = ngramOrder == 0 ? 0 : scratch[ngramOrder - 1][i + 1];
  				assert prevOffset >= 0;
  				scratch[ngramOrder][i - startPos] = ngrams.putWithOffsetAndSuffix(ngram, i, j, prevOffset, suffixOffset, !justLastWord || j == endPos
  				/* || ngram[startPos] == startIndex */
  				? scratchCounts : null);
  			}
  		}
  
  	}
  
  	protected float interpolateProb(final int[] ngram, final int startPos, final int endPos) {
  		if (startPos == endPos) return 0.0f;
  		final float backoff = getLowerOrderBackoff(ngram, startPos, endPos - 1);
  		final float prob = getLowerOrderProb(ngram, startPos, endPos);
  		return prob + backoff * interpolateProb(ngram, startPos + 1, endPos);
  	}
  
  	protected float getHighestOrderProb(final int[] ngram, final int startPos, final int endPos) {
  		final KneserNeyCounts counts = getCounts(ngram, startPos, endPos, false);
  		final KneserNeyCounts rightDotCounts = getCounts(ngram, startPos, endPos - 1, true);
  		final int ngramOrder = endPos - startPos - 1;
  		final float D = getDiscountForOrder(ngramOrder);
  		final float prob = rightDotCounts.tokenCounts == 0 ? 0.0f : Math.max(0.0f, (counts.tokenCounts - D) / rightDotCounts.tokenCounts);
  		return prob;
  	}
  
  	protected float getLowerOrderProb(final int[] ngram, final int startPos, final int endPos) {
  		if (startPos == endPos) return 1.0f;
  		final KneserNeyCounts counts = getCounts(ngram, startPos, endPos, false);
  		final KneserNeyCounts prefixCounts = getCounts(ngram, startPos, endPos - 1, true);
  
  		final float probDiscount = (endPos - startPos == 1) ? 0.0f : getDiscountForOrder(endPos - startPos - 1);
  		final float prob = prefixCounts.dotdotTypeCounts == 0 ? 0.0f : Math.max(0.0f, counts.leftDotTypeCounts - probDiscount) / prefixCounts.dotdotTypeCounts;
  
  		return prob;
  	}
  
  	protected float getLowerOrderBackoff(final int[] ngram, final int startPos, final int endPos) {
  		if (startPos == endPos) return 1.0f;
  		final KneserNeyCounts counts = getCounts(ngram, startPos, endPos, true);
  		final long backoffDenom = (endPos - startPos == lmOrder - 1 || ngram[startPos] == startIndex) ? counts.tokenCounts : counts.dotdotTypeCounts;
  
  		assert backoffDenom >= 0;
  		//		final long backoffDenom = endPos - startPos == lmOrder - 1 ? counts.tokenCounts : counts.dotdotTypeCounts;
  		final float backoffDiscount = getDiscountForOrder(endPos - startPos);
  		final float backoff = backoffDenom == 0.0f ? 1.0f : backoffDiscount * counts.rightDotTypeCounts / backoffDenom;
  		return backoff;
  	}
  
  	protected float getDiscountForOrder(int ngramOrder) {
  		if (opts.kneserNeyDiscounts != null) return (float) opts.kneserNeyDiscounts[ngramOrder];
  		final int numOneCounters = ((KneserNeyCountValueContainer) ngrams.getValues()).getNumOneCountNgrams(ngramOrder);
  		final int numTwoCounters = ((KneserNeyCountValueContainer) ngrams.getValues()).getNumTwoCountNgrams(ngramOrder);
  		final float denom = (numOneCounters + 2 * (float) numTwoCounters);
  		return denom == 0.0f ? 1e-5f : numOneCounters / denom;
  	}
  
  	@Override
  	public void cleanup() {
  
  	}
  
  	/**
  	 * @param key
  	 * @param ngrams
  	 * @param startPos
  	 * @param endPos
  	 */
  	private KneserNeyCounts getCounts(final int[] key, final int startPos, final int endPos, final boolean isBackoff) {
  		final KneserNeyCounts value = new KneserNeyCounts();
  		if (startPos == endPos) {
  			//only happens when requesting number of bigrams
  			value.dotdotTypeCounts = ((KneserNeyCountValueContainer) ngrams.getValues()).getBigramTypeCounts();
  			return value;
  		}
  		final long offset = ngrams.getOffsetForNgramInModel(key, startPos, endPos);
  		if (offset < 0) return value;
  		ngrams.getValues().getFromOffset(offset, endPos - startPos - 1, value);
  
  		final boolean startsWithStartSym = key[startPos] == startIndex;
  		final boolean endsWithEndSym = key[endPos - 1] == wordIndexer.getIndexPossiblyUnk(wordIndexer.getEndSymbol());
  		if (startsWithStartSym) {
  			value.dotdotTypeCounts = value.rightDotTypeCounts;
  			if (endPos - startPos < lmOrder - 1 || (endPos - startPos == lmOrder - 1 && !isBackoff)) value.tokenCounts = value.leftDotTypeCounts;
  		}
  		if (endsWithEndSym) {
  			value.rightDotTypeCounts = 1;
  			value.dotdotTypeCounts = value.leftDotTypeCounts;
  		}
  		return value;
  	}
  
  	public static double[] defaultDiscounts() {
  		return constantArray(defaultMinCounts().length, DEFAULT_DISCOUNT);
  	}
  
  	public static double[] defaultMinCounts() {
  		//same as SRILM
  		return new double[] { 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2 };
  	}
  
  	private static double[] constantArray(final int n, final double f) {
  		final double[] ret = new double[n];
  		Arrays.fill(ret, f);
  		return ret;
  	}
  
  	@Override
  	public void parse(ArpaLmReaderCallback<ProbBackoffPair> callback) {
  		Logger.startTrack("Writing Kneser-Ney probabilities");
  
  		List<Long> lengths = new ArrayList<Long>();
  		for (int ngramOrder = 0; ngramOrder < lmOrder; ++ngramOrder) {
  			Logger.startTrack("Counting counts for order " + ngramOrder);
  			long numNgrams = 0; //ngrams.getNumNgrams(ngramOrder);
  			for (final Entry<KneserNeyCounts> entry : ngrams.getNgramsForOrder(ngramOrder)) {
  				final long relevantCount = entry.value.tokenCounts;
  				if (ngramOrder >= lmOrder - 2 && ngramOrder < opts.kneserNeyMinCounts.length && relevantCount < opts.kneserNeyMinCounts[ngramOrder]) continue;
  				numNgrams++;
  			}
  			lengths.add(numNgrams);
  			Logger.endTrack();
  		}
  		callback.initWithLengths(lengths);
  		for (int ngramOrder = 0; ngramOrder < lmOrder; ++ngramOrder) {
  			callback.handleNgramOrderStarted(ngramOrder + 1);
  
  			Logger.logss("On order " + (ngramOrder + 1));
  			int linenum = 0;
  			for (final Entry<KneserNeyCounts> entry : ngrams.getNgramsForOrder(ngramOrder)) {
  				if (linenum++ % 10000 == 0) Logger.logs("Writing line " + linenum);
  				final long relevantCount = entry.value.tokenCounts;
  				if (ngramOrder >= lmOrder - 2 && ngramOrder < opts.kneserNeyMinCounts.length && relevantCount < opts.kneserNeyMinCounts[ngramOrder]) continue;
  
  				final int[] ngram = entry.key;
  				final int endPos = ngram.length;
  				final int startPos = 0;
  				ProbBackoffPair value = getProbBackoff(ngram, startPos, endPos);
  				callback.call(ngram, startPos, endPos, value, "");
  
  			}
  			callback.handleNgramOrderFinished(ngramOrder + 1);
  
  		}
  		callback.cleanup();
  
  		Logger.endTrack();
  	}
  
  	/**
  	 * @param startIndex
  	 * @param ngramOrder
  	 * @param entry
  	 * @param ngram
  	 * @param endPos
  	 * @param startPos
  	 * @return
  	 */
  	private ProbBackoffPair getProbBackoff(final int[] ngram, final int startPos, final int endPos) {
  		final int ngramOrder = endPos - startPos - 1;
  		final boolean isHighestOrder = ngramOrder == lmOrder - 1;
  		final float val = isHighestOrder || ngram[startPos] == startIndex ? getHighestOrderProb(ngram, startPos, endPos) : getLowerOrderProb(ngram, startPos,
  			endPos);
  		int nextNonStart = startPos + 1;
  		while (nextNonStart < endPos && ngram[nextNonStart] == startIndex) {
  			nextNonStart++;
  		}
  		final float prob = val + getLowerOrderBackoff(ngram, startPos, endPos - 1) * interpolateProb(ngram, nextNonStart, endPos);
  		final boolean isStartEndSym = endPos - startPos == 1 && ngram[startPos] == startIndex;
  		final float logProb = isStartEndSym ? -99 : ((float) (Math.log10(prob)));
  		//		if (logProb == Float.NEGATIVE_INFINITY) {
  		//			System.out.println("here");
  		//		}
  		final float backoff = isHighestOrder ? 0.0f : (float) Math.log10(getLowerOrderBackoff(ngram, startPos, endPos));
  		final ProbBackoffPair ret = new ProbBackoffPair(logProb, backoff);
  		return ret;
  	}
  
  	public WordIndexer<W> getWordIndexer() {
  		return wordIndexer;
  	}
  
  	@Override
  	public void handleNgramOrderFinished(int order) {
  	}
  
  	@Override
  	public void handleNgramOrderStarted(int order) {
  	}
  
  	@Override
  	public int getLmOrder() {
  		return lmOrder;
  	}
  
  	@Override
  	public float scoreSentence(List<W> sentence) {
  		return ArrayEncodedNgramLanguageModel.DefaultImplementations.scoreSentence(sentence, this);
  	}
  
  	@Override
  	public float getLogProb(List<W> ngram) {
  		return ArrayEncodedNgramLanguageModel.DefaultImplementations.getLogProb(ngram, this);
  	}
  
  	@Override
  	public float getLogProb(int[] ngram, int startPos, int endPos) {
  		ProbBackoffPair probBackoff = getProbBackoff(ngram, startPos, endPos);
  		return probBackoff.prob;
  	}
  
  	@Override
  	public float getLogProb(int[] ngram) {
  		return ArrayEncodedNgramLanguageModel.DefaultImplementations.getLogProb(ngram, this);
  	}
  
  	public long getTotalSize() {
  		return ngrams.getTotalSize();
  	}
  
  	@Override
  	public void setOovWordLogProb(float logProb) {
  		throw new UnsupportedOperationException("Method not yet implemented");
  	}
  
  }
  