package edu.berkeley.nlp.lm.map;

import java.io.Serializable;
import java.util.Iterator;
import java.util.List;

import edu.berkeley.nlp.lm.ConfigOptions;
import edu.berkeley.nlp.lm.array.CustomWidthArray;
import edu.berkeley.nlp.lm.array.LongArray;
import edu.berkeley.nlp.lm.bits.BitList;
import edu.berkeley.nlp.lm.bits.BitStream;
import edu.berkeley.nlp.lm.bits.VariableLengthBitCompressor;
import edu.berkeley.nlp.lm.util.Annotations.OutputParameter;
import edu.berkeley.nlp.lm.util.Logger;
import edu.berkeley.nlp.lm.values.CompressibleValueContainer;

public class CompressedNgramMap<T> extends AbstractNgramMap<T> implements Serializable
{

	/**
	 * 
	 */
	private static final long serialVersionUID = 1L;

	private final int compressedBlockSize;

	private static final int OFFSET_RADIX = 33;

	private static final int WORD_RADIX = 2;

	private final VariableLengthBitCompressor offsetCoder;

	private final VariableLengthBitCompressor wordCoder;

	private final VariableLengthBitCompressor suffixCoder;

	private double totalKeyBitsFinal = 0;

	private double totalValueBitsFinal = 0;

	private double totalBitsFinal = 0;

	private double totalSizeFinal = 0;

	private final int offsetDeltaRadix;

	private final CompressedMap[] maps;

	private final boolean reverseTrie = true;

	private final long[] numNgramsForEachOrder;

	public CompressedNgramMap(final CompressibleValueContainer<T> values, final long[] numNgramsForEachOrder, final ConfigOptions opts) {
		super(values, opts);
		offsetCoder = new VariableLengthBitCompressor(OFFSET_RADIX);
		wordCoder = new VariableLengthBitCompressor(WORD_RADIX);
		this.offsetDeltaRadix = opts.offsetDeltaRadix;
		suffixCoder = new VariableLengthBitCompressor(offsetDeltaRadix);
		this.compressedBlockSize = opts.compressedBlockSize;
		this.numNgramsForEachOrder = numNgramsForEachOrder;
		this.maps = new CompressedMap[numNgramsForEachOrder.length];
		values.setMap(this);

	}

	@Override
	public long getValueAndOffset(final long contextOffset, final int contextNgramOrder, final int word, @OutputParameter final T outputVal) {
		if (word < 0) return -1L;
		final long hash = combineToKey(word, contextOffset);
		final int ngramOrder = contextNgramOrder + 1;
		final LongArray compressedKeys = (maps[ngramOrder]).compressedKeys;
		final long currIndex = decompressSearch(compressedKeys, hash, ngramOrder, outputVal);
		return currIndex;

	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see edu.berkeley.nlp.mt.lm.NgramMap#add(java.util.List, T)
	 */
	@Override
	public long put(final int[] ngram, final int startPos, final int endPos, final T val) {

		final int ngramOrder = endPos - startPos - 1;
		final int word = reverseTrie ? ngram[startPos] : ngram[endPos - 1];

		final long contextOffset = reverseTrie ? getContextOffset(ngram, startPos + 1, endPos, null) : getContextOffset(ngram, startPos, endPos - 1, null);
		if (contextOffset < 0) return -1;

		CompressedMap map = maps[ngramOrder];
		if (map == null) {
			map = maps[ngramOrder] = new CompressedMap();
			final long l = numNgramsForEachOrder[ngramOrder];
			maps[ngramOrder].init(l);
			values.setSizeAtLeast(l, ngramOrder);
		}
		final long oldSize = map.size();
		final long newOffset = map.add(combineToKey(word, contextOffset));
		final boolean addWorked = values.add(ngram, startPos, endPos, ngramOrder, map.size() - 1, contextOffset, word, val, -1, map.size() == oldSize);
		if (!addWorked) return -1;
		return newOffset;

	}

	private long getContextOffset(final int[] ngram, final int startPos, final int endPos, T val) {
		if (endPos == startPos) return 0;
		long hasValueSuffixIndex = 0;
		if (endPos > startPos) {
			long lastSuffix = 0L;
			for (int ngramOrder = 0; ngramOrder < endPos - startPos; ++ngramOrder) {
				final int firstWord = reverseTrie ? ngram[endPos - ngramOrder - 1] : ngram[startPos + ngramOrder];
				final long key = combineToKey(firstWord, lastSuffix);

				if (maps[ngramOrder] == null) return -1;

				final LongArray compressedKeys = (maps[ngramOrder]).compressedKeys;
				final long currIndex = decompressSearch(compressedKeys, key, ngramOrder, val);
				if (currIndex < 0) return -1;
				lastSuffix = currIndex;
			}
			hasValueSuffixIndex = lastSuffix;
		}
		return hasValueSuffixIndex;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see edu.berkeley.nlp.mt.lm.NgramMap#handleNgramsFinished(int)
	 */
	@Override
	public void handleNgramsFinished(final int justFinishedOrder) {
		final CompressedMap compressedMap = maps[justFinishedOrder - 1];
		if (compressedMap != null) {
			final LongArray currKeys = compressedMap.getUncompressedKeys();
			final long currSize = currKeys.size();
			sort(currKeys, 0, currSize - 1, justFinishedOrder - 1);
			compressedMap.trim();
			values.trimAfterNgram(justFinishedOrder - 1, currSize);
			compress(justFinishedOrder - 1);
		}
	}

	protected static int compareLongsRaw(final long a, final long b) {
		assert a >= 0;
		assert b >= 0;
		if (a > b) return +1;
		if (a < b) return -1;
		if (a == b) return 0;
		throw new RuntimeException();
	}

	private void compress(final int ngramOrder) {
		if (ngramOrder > 0) {
			(maps[ngramOrder]).compressedKeys = compress(maps[ngramOrder].getUncompressedKeys(), maps[ngramOrder].size(), ngramOrder);
			((CompressibleValueContainer<T>) values).clearStorageAfterCompression(ngramOrder);
			maps[ngramOrder].clearUncompressedKeys();
		}
	}

	private LongArray compress(final LongArray uncompressed, final long uncompressedSize, final int ngramOrder) {
		Logger.startTrack("Compressing");
		final LongArray compressedLongArray = LongArray.StaticMethods.newLongArray(Long.MAX_VALUE, uncompressedSize >>> 2);

		long uncompressedPos = 0;
		long totalNumKeyBits = 0;
		long totalNumValueBits = 0;
		long currBlock = 0;

		final CompressibleValueContainer<T> compressibleValues = (CompressibleValueContainer<T>) values;
		while (uncompressedPos < uncompressedSize) {
			final BitList currBlockBits = new BitList();
			final long firstKey = uncompressed.get(uncompressedPos);

			if (currBlock++ % 1000 == 0) Logger.logs("On block " + currBlock + " starting at pos " + uncompressedPos);

			currBlockBits.addLong(firstKey);
			final BitList offsetBits = offsetCoder.compress(uncompressedPos);

			final BitList firstValueBits = compressibleValues.getCompressed(uncompressedPos, ngramOrder);
			BitList headerBits = new BitList();
			BitList bodyBits = new BitList();
			long numKeyBits = 0;
			long numValueBits = 0;
			long currUncompressedPos = -1;

			// try compression assuming all words are the same (wordBitOn = false), and if that fails,
			// roll back and try with wordBitOn = true
			OUTER: for (boolean wordBitOn = false, done = false; !done; wordBitOn = true) {
				numKeyBits = 0;
				numValueBits = 0;
				long lastFirstWord = wordOf(firstKey);
				long lastSuffixPart = contextOffsetOf(firstKey);
				headerBits = makeHeader(offsetBits, firstValueBits, wordBitOn);
				bodyBits = new BitList();

				final BitList currBits = new BitList();
				for (currUncompressedPos = uncompressedPos + 1; currUncompressedPos < uncompressedSize; ++currUncompressedPos) {
					final long currKey = uncompressed.get(currUncompressedPos);
					final long currFirstWord = wordOf(currKey);
					final long currSuffixPart = contextOffsetOf(currKey);

					final long wordDelta = currFirstWord - lastFirstWord;
					final long suffixDelta = currSuffixPart - lastSuffixPart;
					currBits.clear();
					if (wordDelta > 0 && !wordBitOn) continue OUTER;
					if (wordBitOn) {
						final BitList keyBits = wordCoder.compress(wordDelta);
						currBits.addAll(keyBits);
						if (wordDelta > 0) {
							final BitList suffixBits = suffixCoder.compress(currSuffixPart);
							currBits.addAll(suffixBits);
						} else {
							final BitList suffixBits = suffixCoder.compress(suffixDelta);
							currBits.addAll(suffixBits);
						}
					} else {

						final BitList suffixBits = suffixCoder.compress(suffixDelta);
						currBits.addAll(suffixBits);
					}

					numKeyBits += currBits.size();
					lastFirstWord = currFirstWord;
					numValueBits += compressValue(ngramOrder, currUncompressedPos, currBits);

					lastSuffixPart = currSuffixPart;
					if (blockFull(currBlockBits, bodyBits, headerBits, currBits)) {
						break;
					}

					bodyBits.addAll(currBits);
				}
				done = true;
			}

			uncompressedPos = currUncompressedPos;

			totalNumKeyBits += numKeyBits;
			totalNumValueBits += numValueBits;

			final int bitLength = bodyBits.size() + headerBits.size();
			assert bitLength <= Short.MAX_VALUE;
			currBlockBits.addShort((short) bitLength);
			currBlockBits.addAll(headerBits);
			currBlockBits.addAll(bodyBits);

			assert currBlockBits.size() < Long.SIZE * compressedBlockSize;
			writeBlockToArray(currBlockBits, compressedLongArray);
		}
		compressedLongArray.trim();

		logCompressionInfo(uncompressedSize, compressedLongArray, totalNumKeyBits, totalNumValueBits);

		Logger.endTrack();
		return compressedLongArray;
	}

	/**
	 * @param blockBits
	 * @param array
	 */
	private void writeBlockToArray(final BitList blockBits, final LongArray array) {
		long curr = 0L;
		for (int i = 0; i <= Long.SIZE * compressedBlockSize; ++i) {
			if (i % Long.SIZE == 0 && i > 0) {
				array.add(curr);
				curr = 0;

			}
			curr = (curr << 1) | ((i >= blockBits.size() || !blockBits.get(i)) ? 0 : 1);
		}
		assert array.size() % compressedBlockSize == 0;
	}

	/**
	 * @param uncompressedSize
	 * @param compressedLongArray
	 * @param keyBits
	 * @param valueBits
	 */
	private void logCompressionInfo(final long uncompressedSize, final LongArray compressedLongArray, final long keyBits, final long valueBits) {
		final double keyAvg = (double) keyBits / uncompressedSize;
		Logger.logss("Key bits " + keyAvg);
		final double valueAvg = (double) valueBits / uncompressedSize;
		Logger.logss("Value bits " + valueAvg);
		final double avg = 64 * (double) compressedLongArray.size() / uncompressedSize;
		Logger.logss("Compressed bits " + avg);
		totalKeyBitsFinal += keyBits;
		totalValueBitsFinal += valueBits;
		totalBitsFinal += compressedLongArray.size();
		totalSizeFinal += uncompressedSize;
		Logger.logss("Total key bits " + totalKeyBitsFinal / totalSizeFinal);
		Logger.logss("Total value bits " + totalValueBitsFinal / totalSizeFinal);
		Logger.logss("Total bits " + 64.0 * totalBitsFinal / totalSizeFinal);
	}

	/**
	 * @param currBits
	 * @param restBits
	 * @param headerBits
	 * @param newBits
	 * @return
	 */
	private boolean blockFull(final BitList currBits, final BitList restBits, final BitList headerBits, final BitList newBits) {
		final int numTotalBitsSize = Short.SIZE;
		final int lengthSoFar = currBits.size() + numTotalBitsSize + headerBits.size() + restBits.size() + newBits.size();
		return lengthSoFar >= Long.SIZE * compressedBlockSize;
	}

	/**
	 * @param ngramOrder
	 * @param currPos
	 * @param newBits
	 * @return
	 */
	private long compressValue(final int ngramOrder, final long currPos, final BitList newBits) {
		final BitList valueBits = ((CompressibleValueContainer<T>) values).getCompressed(currPos, ngramOrder);
		newBits.addAll(valueBits);
		return valueBits.size();
	}

	/**
	 * @param offsetBits
	 * @param firstValueBits
	 * @param wordBitOn
	 * @return
	 */
	private BitList makeHeader(final BitList offsetBits, final BitList firstValueBits, final boolean wordBitOn) {
		BitList headerBits;
		headerBits = new BitList();

		headerBits.addAll(offsetBits);
		headerBits.add(wordBitOn);
		headerBits.addAll(firstValueBits);
		return headerBits;
	}

	/**
	 * searchOffset >= 0 means we are looking for a specific offset and ignore
	 * searchKey if searchOffset >= 0, we return the key, else we return the
	 * offset for searchKey
	 * 
	 * @param compressed
	 * @param pos
	 * @param searchKey
	 * @param ngramOrder
	 * @param outputVal
	 * @param searchOffset
	 * @return
	 */
	private long decompressLinearSearch(final LongArray compressed, final long pos, final long searchKey, final int ngramOrder, final T outputVal,
		final long searchOffset) {

		final long firstKey = compressed.get(pos);
		final BitStream bits = getCompressedBits(compressed, pos + 1);
		final long offset = offsetCoder.decompress(bits);
		final boolean wordBitOn = bits.nextBit();

		int currWord = wordOf(firstKey);
		long currSuffix = contextOffsetOf(firstKey);
		final boolean foundKeyFirst = searchOffset >= 0 ? searchOffset == offset : firstKey == searchKey;

		final CompressibleValueContainer<T> compressibleValues = (CompressibleValueContainer<T>) values;
		compressibleValues.decompress(bits, ngramOrder, !foundKeyFirst, outputVal);
		if (foundKeyFirst) return searchOffset >= 0 ? firstKey : offset;

		long currKey = -1;

		for (int k = 1; !bits.finished(); ++k) {
			int newWord = -1;
			long nextSuffix = -1;

			if (wordBitOn) {
				final int wordDelta = (int) wordCoder.decompress(bits);
				final boolean wordDeltaIsZero = wordDelta == 0;
				final long suffixDelta = suffixCoder.decompress(bits);
				newWord = currWord + wordDelta;
				nextSuffix = wordDeltaIsZero ? (currSuffix + suffixDelta) : suffixDelta;
			} else {
				final long suffixDelta = suffixCoder.decompress(bits);
				newWord = currWord;
				nextSuffix = (currSuffix + suffixDelta);
			}
			currKey = combineToKey(newWord, nextSuffix);
			currWord = newWord;
			currSuffix = nextSuffix;
			final long currOffset = offset + k;
			final boolean foundKey = searchOffset >= 0 ? searchOffset == currOffset : currKey == searchKey;
			compressibleValues.decompress(bits, ngramOrder, !foundKey, outputVal);
			if (foundKey) { return searchOffset >= 0 ? currKey : currOffset; }
			if (searchOffset >= 0) {
				if (currOffset > searchOffset) return -1;
			} else if (currKey > searchKey) return -1;

		}
		return -1;

	}

	/**
	 * @param compressed
	 * @param pos
	 * @return
	 */
	private BitStream getCompressedBits(final LongArray compressed, final long pos) {
		final short bitLength = readShort(compressed.get(pos));
		final BitStream bits = new BitStream(compressed, pos, Short.SIZE, bitLength);
		return bits;
	}

	private short readShort(final long l) {
		return (short) (l >>> (Long.SIZE - Short.SIZE));
	}

	private long decompressSearch(final LongArray compressed, final long searchKey, final int ngramOrder, final T outputVal) {

		return decompressSearch(compressed, searchKey, ngramOrder, outputVal, -1);
	}

	private long decompressSearch(final LongArray compressed, final long searchKey, final int ngramOrder, final T outputVal, final long searchOffset) {

		if (ngramOrder == 0) {
			final boolean lookingForOffset = searchKey >= 0;
			final int word = lookingForOffset ? wordOf(searchKey) : (int) searchOffset;
			if (word < 0 || word >= maps[0].size()) return -1;
			if (outputVal != null) values.getFromOffset(word, 0, outputVal);
			return lookingForOffset ? word : combineToKey(word, 0);
		} else {
			if (compressed == null) return -1;
			final long fromIndex = 0;
			final long toIndex = ((compressed.size() / compressedBlockSize) - 1);
			final long low = binarySearchBlocks(compressed, compressed.size(), searchKey, fromIndex, toIndex, searchOffset);
			if (low < 0) return -1;

			final long index = decompressLinearSearch(compressed, low, searchKey, ngramOrder, outputVal, searchOffset);
			return index;
		}
	}

	/**
	 * @param compressed
	 * @param searchKey
	 * @return
	 */
	private long binarySearchBlocks(final LongArray compressed, final long size, final long searchKey, final long low_, final long high_,
		final long searchOffset) {
		final long toFind = searchOffset >= 0 ? searchOffset : searchKey;
		long low = low_;
		long high = high_;
		assert size % compressedBlockSize == 0;

		while (low <= high) {
			final long mid = (low + high) >>> 1;
			final long currPos = mid * compressedBlockSize;
			final long midVal = searchOffset >= 0 ? offsetCoder.decompress(getCompressedBits(compressed, currPos + 1)) : compressed.get(currPos);
			final int compare = compareLongsRaw(midVal, toFind);
			if (compare < 0) //midVal < key
				low = mid + 1;
			else if (compare > 0) // midVal > key
				high = mid - 1;
			else {
				low = mid + 1;
				break;// key found
			}
		}
		if (low <= 0) return -1;
		final long i = (low - 1) * compressedBlockSize;
		return i;
	}

	protected void sort(final LongArray array, final long left0, final long right0, final int ngramOrder) {

		long left, right;
		long pivot;
		left = left0;
		right = right0 + 1;

		final long pivotIndex = (left0 + right0) >>> 1;

		pivot = array.get(pivotIndex);//[outerArrayPart(pivotIndex)][innerArrayPart(pivotIndex)];
		swap(pivotIndex, left0, array, ngramOrder);

		do {

			do
				left++;
			while (left <= right0 && compareLongsRaw(array.get(left), pivot) < 0);

			do
				right--;
			while (compareLongsRaw(array.get(right), pivot) > 0);

			if (left < right) {
				swap(left, right, array, ngramOrder);
			}

		} while (left <= right);

		swap(left0, right, array, ngramOrder);

		if (left0 < right) sort(array, left0, right, ngramOrder);
		if (left < right0) sort(array, left, right0, ngramOrder);

	}

	protected void swap(final long a, final long b, final LongArray array, final int ngramOrder) {
		swap(array, a, b);
		((CompressibleValueContainer<T>) values).swap(a, b, ngramOrder);
	}

	protected void swap(final LongArray array, final long a, final long b) {
		final long temp = array.get(a);
		array.set(a, array.get(b));
		array.set(b, temp);
	}

	@Override
	public void trim() {
		values.trim();

	}

	@Override
	public void initWithLengths(final List<Long> numNGrams) {
	}

	@Override
	public int getMaxNgramOrder() {
		return maps.length;
	}

	@Override
	public Iterable<Entry<T>> getNgramsForOrder(final int ngramOrder) {
		return new Iterable<Entry<T>>()
		{

			@Override
			public Iterator<edu.berkeley.nlp.lm.map.NgramMap.Entry<T>> iterator() {
				return new Iterator<edu.berkeley.nlp.lm.map.NgramMap.Entry<T>>()
				{

					long currOffset = 0;

					@Override
					public boolean hasNext() {
						return currOffset < maps[ngramOrder].size();
					}

					@Override
					public edu.berkeley.nlp.lm.map.NgramMap.Entry<T> next() {
						final T scratch_ = values.getScratchValue();
						long offset = currOffset;
						final int[] ngram = new int[ngramOrder + 1];
						for (int i = ngramOrder; i >= 0; --i) {
							final T scratch = i == ngramOrder ? scratch_ : null;
							final long foundKey = decompressSearch(maps[i].compressedKeys, -1, i, scratch, offset);
							assert foundKey >= 0;
							ngram[reverseTrie ? (ngramOrder - i) : i] = wordOf(foundKey);
							offset = contextOffsetOf(foundKey);
						}
						currOffset++;

						return new Entry<T>(ngram, scratch_);
					}

					@Override
					public void remove() {
						throw new UnsupportedOperationException("Method not yet implemented");
					}
				};
			}
		};
	}

	@Override
	public long getNumNgrams(final int ngramOrder) {
		return maps[ngramOrder].size();
	}

	@Override
	public boolean contains(final int[] ngram, final int startPos, final int endPos) {
		return getContextOffset(ngram, startPos, endPos, null) >= 0;

	}

	@Override
	public T get(int[] ngram, int startPos, int endPos) {
		T val = values.getScratchValue();
		final long offset = getContextOffset(ngram, startPos, endPos, val);
		if (offset < 0) { return null; }
		return val;

	}

	@Override
	public CustomWidthArray getValueStoringArray(final int ngramOrder) {
		return null;
	}

	@Override
	public void clearStorage() {
		for (int i = 0; i < maps.length; ++i)
			maps[i] = null;
	}

}
