/*
 * Copyright (C) 2014 Indeed Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the
 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing permissions and
 * limitations under the License.
 */
 package com.indeed.imhotep.local;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimaps;
import com.google.common.primitives.Ints;
import com.google.common.primitives.Longs;
import com.google.protobuf.InvalidProtocolBufferException;
import com.indeed.flamdex.api.DocIdStream;
import com.indeed.flamdex.api.FlamdexOutOfMemoryException;
import com.indeed.flamdex.api.FlamdexReader;
import com.indeed.flamdex.api.IntTermIterator;
import com.indeed.flamdex.api.IntValueLookup;
import com.indeed.flamdex.api.RawFlamdexReader;
import com.indeed.flamdex.api.StringTermDocIterator;
import com.indeed.flamdex.api.StringTermIterator;
import com.indeed.flamdex.api.StringValueLookup;
import com.indeed.flamdex.datastruct.FastBitSet;
import com.indeed.flamdex.datastruct.FastBitSetPooler;
import com.indeed.flamdex.fieldcache.ByteArrayIntValueLookup;
import com.indeed.flamdex.fieldcache.IntArrayIntValueLookup;
import com.indeed.flamdex.query.Query;
import com.indeed.flamdex.query.Term;
import com.indeed.flamdex.reader.FlamdexMetadata;
import com.indeed.flamdex.search.FlamdexSearcher;
import com.indeed.flamdex.simple.SimpleFlamdexReader;
import com.indeed.flamdex.simple.SimpleFlamdexWriter;
import com.indeed.flamdex.utils.FlamdexUtils;
import com.indeed.imhotep.AbstractImhotepSession;
import com.indeed.imhotep.FTGSSplitter;
import com.indeed.imhotep.GroupMultiRemapRule;
import com.indeed.imhotep.GroupRemapRule;
import com.indeed.imhotep.ImhotepMemoryPool;
import com.indeed.imhotep.MemoryReservationContext;
import com.indeed.imhotep.MemoryReserver;
import com.indeed.imhotep.QueryRemapRule;
import com.indeed.imhotep.RegroupCondition;
import com.indeed.imhotep.TermCount;
import com.indeed.imhotep.api.DocIterator;
import com.indeed.imhotep.api.FTGSIterator;
import com.indeed.imhotep.api.ImhotepOutOfMemoryException;
import com.indeed.imhotep.api.RawFTGSIterator;
import com.indeed.imhotep.group.ImhotepChooser;
import com.indeed.imhotep.marshal.ImhotepDaemonMarshaller;
import com.indeed.imhotep.metrics.AbsoluteValue;
import com.indeed.imhotep.metrics.Addition;
import com.indeed.imhotep.metrics.CachedInterleavedMetrics;
import com.indeed.imhotep.metrics.CachedMetric;
import com.indeed.imhotep.metrics.Constant;
import com.indeed.imhotep.metrics.Count;
import com.indeed.imhotep.metrics.DelegatingMetric;
import com.indeed.imhotep.metrics.Division;
import com.indeed.imhotep.metrics.Equal;
import com.indeed.imhotep.metrics.Exponential;
import com.indeed.imhotep.metrics.GreaterThan;
import com.indeed.imhotep.metrics.GreaterThanOrEqual;
import com.indeed.imhotep.metrics.LessThan;
import com.indeed.imhotep.metrics.LessThanOrEqual;
import com.indeed.imhotep.metrics.Log;
import com.indeed.imhotep.metrics.Log1pExp;
import com.indeed.imhotep.metrics.Logistic;
import com.indeed.imhotep.metrics.Max;
import com.indeed.imhotep.metrics.Min;
import com.indeed.imhotep.metrics.Modulus;
import com.indeed.imhotep.metrics.Multiplication;
import com.indeed.imhotep.metrics.MultiplyAndShiftRight;
import com.indeed.imhotep.metrics.NotEqual;
import com.indeed.imhotep.metrics.ShiftLeftAndDivide;
import com.indeed.imhotep.metrics.Subtraction;
import com.indeed.imhotep.protobuf.QueryMessage;
import com.indeed.imhotep.service.CachedFlamdexReader;
import com.indeed.imhotep.service.RawCachedFlamdexReader;
import com.indeed.util.core.Pair;
import com.indeed.util.core.Throwables2;
import com.indeed.util.core.io.Closeables2;
import com.indeed.util.core.reference.SharedReference;
import com.indeed.util.core.threads.ThreadSafeBitSet;
import dk.brics.automaton.Automaton;
import dk.brics.automaton.RegExp;
import it.unimi.dsi.fastutil.PriorityQueue;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectHeapPriorityQueue;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.io.FileUtils;
import org.apache.log4j.Logger;

import javax.annotation.Nonnull;
import java.io.EOFException;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * This class isn't even close to remotely thread safe, do not use it
 * simultaneously from multiple threads
 */
public final class ImhotepLocalSession extends AbstractImhotepSession {
    static final Logger log = Logger.getLogger(ImhotepLocalSession.class);

    static final boolean logTiming;

    static {
        logTiming =
                "true".equals(System.getProperty("com.indeed.imhotep.local.ImhotepLocalSession.logTiming"));
    }

    private static final int MAX_NUMBER_STATS = 64;
    static final int BUFFER_SIZE = 2048;
    private final AtomicLong tempFileSizeBytesLeft;

    private int numDocs;
    // buffers that will be reused to avoid excessive allocations
    final int[] docIdBuf = new int[BUFFER_SIZE];
    final long[] valBuf = new long[BUFFER_SIZE];
    final int[] docGroupBuffer = new int[BUFFER_SIZE];

    // do not close flamdexReader, it is separately refcounted
    private FlamdexReader flamdexReader;
    private SharedReference<FlamdexReader> flamdexReaderRef;
    private FlamdexReader originalReader;
    private SharedReference<FlamdexReader> originalReaderRef;

    private final String optimizedIndexesDir;

    final MemoryReservationContext memory;

    GroupLookup docIdToGroup;

    int[] groupDocCount;

    int numStats;
    private long[][] groupStats = new long[MAX_NUMBER_STATS][];
    final IntValueLookup[] statLookup = new IntValueLookup[MAX_NUMBER_STATS];
    private final List<String> statCommands;

    private final boolean[] needToReCalcGroupStats = new boolean[MAX_NUMBER_STATS];

    private boolean closed = false;
    @VisibleForTesting
    private Map<String, DynamicMetric> dynamicMetrics = Maps.newHashMap();

    Map<Pair<String, Boolean>, FastBitSet> fieldZeroDocBitsets;

    private final Exception constructorStackTrace;

    private final File optimizationLog;

    private FTGSSplitter ftgsIteratorSplits;

    public ImhotepLocalSession(final FlamdexReader flamdexReader) throws ImhotepOutOfMemoryException {
        this(flamdexReader, null,
                new MemoryReservationContext(new ImhotepMemoryPool(Long.MAX_VALUE)), false, null);
    }

    public ImhotepLocalSession(FlamdexReader flamdexReader, boolean optimizeGroupZeroLookups) throws ImhotepOutOfMemoryException {
        this(flamdexReader, null,
                new MemoryReservationContext(new ImhotepMemoryPool(Long.MAX_VALUE)),
                optimizeGroupZeroLookups, null);
    }

    public ImhotepLocalSession(final FlamdexReader flamdexReader,
                               String optimizedIndexDirectory,
                               final MemoryReservationContext memory,
                               boolean optimizeGroupZeroLookups,
                               AtomicLong tempFileSizeBytesLeft) throws ImhotepOutOfMemoryException {
        this.tempFileSizeBytesLeft = tempFileSizeBytesLeft;
        constructorStackTrace = new Exception();
        flamdexReaderRef = SharedReference.create(flamdexReader);
        this.flamdexReader = flamdexReader;
        this.memory = memory;
        this.numDocs = flamdexReader.getNumDocs();
        this.optimizedIndexesDir = optimizedIndexDirectory;

        if (!memory.claimMemory(BUFFER_SIZE * (4 + 4 + 4) + 12 * 2)) {
            throw new ImhotepOutOfMemoryException();
        }

        docIdToGroup = new ConstantGroupLookup(this, 1, numDocs);
        docIdToGroup.recalculateNumGroups();
        groupDocCount = clearAndResize((int[]) null, docIdToGroup.getNumGroups(), memory);
        groupDocCount[1] = numDocs;
        this.statCommands = new ArrayList<String>();

        this.optimizationLog =
                new File(this.optimizedIndexesDir, UUID.randomUUID().toString()
                        + ".optimization_log");

        if (optimizeGroupZeroLookups) {
            fieldZeroDocBitsets = Maps.newHashMap();
        } else {
            fieldZeroDocBitsets = null;
        }
    }

    FlamdexReader getReader() {
        return this.flamdexReader;
    }

    public Map<String, DynamicMetric> getDynamicMetrics() {
        return dynamicMetrics;
    }

    int getNumDocs() {
        return this.numDocs;
    }

    /*
     * record structure to store the info necessary to rebuild the
     * DynamicMetrics after one or more optimizes and a reset
     */
    public static class OptimizationRecord implements Serializable {
        private static final long serialVersionUID = 1L;
        public long time;
        List<String> intFieldsMerged;
        List<String> stringFieldsMerged;
        String shardLocation;
        List<ShardMergeInfo> mergedShards;
    }

    public static class ShardMergeInfo implements Serializable {
        private static final long serialVersionUID = 1L;
        int numDocs;
        Map<String, DynamicMetric> dynamicMetrics;
        int[] newDocIdToOldDocId;
    }

    /*
     * Finds a good place to store the new, optimized shard and opens a
     * SimpleFlamdexWriter to it.
     */
    private SimpleFlamdexWriter createNewTempWriter(int maxDocs) throws IOException {
        final File tempIdxDir;
        final String newShardName;
        final File newShardDir;

        newShardName = "temp." + UUID.randomUUID().toString();
        tempIdxDir = new File(this.optimizedIndexesDir);
        newShardDir = new File(tempIdxDir, newShardName);
        newShardDir.mkdir();

        return new SimpleFlamdexWriter(newShardDir.getCanonicalPath(), maxDocs);
    }

    /* wrapper for SimpleFlamdexReader which deletes the on disk data on close() */
    private static class AutoDeletingReader extends SimpleFlamdexReader {

        public AutoDeletingReader(String directory,
                                  int numDocs,
                                  Collection<String> intFields,
                                  Collection<String> stringFields,
                                  boolean useMMapMetrics) {
            super(directory, numDocs, intFields, stringFields, useMMapMetrics);
        }

        public static AutoDeletingReader open(String directory) throws IOException {
            return open(directory, new Config());
        }

        public static AutoDeletingReader open(String directory, Config config) throws IOException {
            final FlamdexMetadata metadata = FlamdexMetadata.readMetadata(directory);
            final Collection<String> intFields = scan(directory, ".intterms");
            final Collection<String> stringFields = scan(directory, ".strterms");
            if (config.isWriteBTreesIfNotExisting()) {
                buildIntBTrees(directory, Lists.newArrayList(intFields));
                buildStringBTrees(directory, Lists.newArrayList(stringFields));
            }
            return new AutoDeletingReader(directory, metadata.numDocs, intFields, stringFields,
                                          config.isUseMMapMetrics());
        }

        @Override
        public void close() throws IOException {
            File dir = new File(this.directory);
            super.close();

            FileUtils.deleteDirectory(dir);
        }

    }

    /* Tweak to ObjectOutputStream which allows it to append to an existing file */
    public static class AppendingObjectOutputStream extends ObjectOutputStream {

        public AppendingObjectOutputStream(OutputStream out) throws IOException {
            super(out);
        }

        @Override
        protected void writeStreamHeader() throws IOException {
            // do not write a header, but reset:
            reset();
        }

    }

    @Override
    public synchronized void rebuildAndFilterIndexes(@Nonnull final List<String> intFields,
                                                     @Nonnull final List<String> stringFields) throws ImhotepOutOfMemoryException {
        final IndexReWriter rewriter;
        final ObjectOutputStream oos;
        final SimpleFlamdexWriter w;
        final ArrayList<String> statsCopy;

        long time = System.currentTimeMillis();

        /* pop off all the stats, they will be repushed after the optimization */
        statsCopy = new ArrayList<String>(this.statCommands);
        while (this.numStats > 0) {
            this.popStat();
        }
        this.statCommands.clear();

        MemoryReservationContext rewriterMemory = new MemoryReservationContext(memory);
        rewriter = new IndexReWriter(Arrays.asList(this), this, rewriterMemory);
        try {
            w = createNewTempWriter(this.numDocs);
            rewriter.optimizeIndecies(intFields, stringFields, w);
            w.close();

            /*
             * save a record of the merge, so it can be unwound later if the
             * shards are reset
             */
            if (this.optimizationLog.exists()) {
                /* plain ObjectOutputStream does not append correctly */
                oos =
                        new AppendingObjectOutputStream(new FileOutputStream(this.optimizationLog,
                                                                             true));
            } else {
                oos = new ObjectOutputStream(new FileOutputStream(this.optimizationLog, true));
            }

            OptimizationRecord record = new OptimizationRecord();

            record.time = time;
            record.intFieldsMerged = intFields;
            record.stringFieldsMerged = stringFields;
            record.shardLocation = w.getOutputDirectory();
            record.mergedShards = new ArrayList<ShardMergeInfo>();

            ShardMergeInfo info = new ShardMergeInfo();
            info.numDocs = this.flamdexReader.getNumDocs();
            info.dynamicMetrics = this.getDynamicMetrics();
            info.newDocIdToOldDocId = rewriter.getPerSessionMappings().get(0);
            record.mergedShards.add(info);

            oos.writeObject(record);
            oos.close();

            /* use rebuilt structures */
            memory.releaseMemory(this.docIdToGroup.memoryUsed());
            rewriterMemory.hoist(rewriter.getNewGroupLookup().memoryUsed());
            this.docIdToGroup = rewriter.getNewGroupLookup();

            for (DynamicMetric dm : this.dynamicMetrics.values()) {
                memory.releaseMemory(dm.memoryUsed());
            }
            for (DynamicMetric dm : rewriter.getDynamicMetrics().values()) {
                rewriterMemory.hoist(dm.memoryUsed());
            }
            this.dynamicMetrics = rewriter.getDynamicMetrics();

            /* release memory used by the index rewriter */
            rewriterMemory.close();

            /*
             * replace flamdexReader pointers, but keep the originals in case
             * there is a reset() call
             */
            if (this.originalReader == null) {
                this.originalReader = this.flamdexReader;
            }
            if (this.originalReaderRef == null) {
                this.originalReaderRef = this.flamdexReaderRef;
            } else {
                /* close the unnecessary optimized index */
                this.flamdexReaderRef.close();
            }

            FlamdexReader flamdex = AutoDeletingReader.open(w.getOutputDirectory());
            if (flamdex instanceof RawFlamdexReader) {
                this.flamdexReader =
                        new RawCachedFlamdexReader(new MemoryReservationContext(memory),
                                                   (RawFlamdexReader) flamdex, null, null, null,
                                                   null);
            } else {
                this.flamdexReader =
                        new CachedFlamdexReader(new MemoryReservationContext(memory), flamdex,
                                                null, null, null, null);
            }
            this.flamdexReaderRef = SharedReference.create(this.flamdexReader);

            /* alter tracking fields to reflect the removal of group 0 docs */
            this.numDocs = this.flamdexReader.getNumDocs();
            this.groupDocCount[0] = 0;

            /* push the stats back on */
            for (String stat : statsCopy) {
                if ("pop".equals(stat)) {
                    this.popStat();
                } else {
                    this.pushStat(stat);
                }
            }

        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

    /*
     * Resets the Flamdex readers to the original un-optimized versions and
     * constructs the DynamicMetrics to match what they should be if no
     * optimization had taken place.
     * 
     * The GroupLookup does not have to be reconstructed since it will be set to
     * a constant value as a result of a reset() call
     */
    private synchronized void resetOptimizedReaders() throws ImhotepOutOfMemoryException {
        ObjectInputStream ois = null;
        ArrayList<OptimizationRecord> records = new ArrayList<OptimizationRecord>();
        final long memoryUse;
        final ArrayList<String> statsCopy;

        /* check if this session has been optimized */
        if (this.originalReader == null) {
            return;
        }

        /* check for space in memory */
        memoryUse = this.optimizationLog.length();
        if (!this.memory.claimMemory(memoryUse)) {
            throw new ImhotepOutOfMemoryException();
        }

        /* pop off all the stats, they will be repushed after the flamdex reset */
        statsCopy = new ArrayList<String>(this.statCommands);
        while (this.numStats > 0) {
            this.popStat();
        }
        this.statCommands.clear();

        /* read in all the optimization records */
        try {
            ois = new ObjectInputStream(new FileInputStream(this.optimizationLog));
            while (true) {
                /*
                 * adds the records so the last written record is first in the
                 * list
                 */
                records.add(0, (OptimizationRecord) ois.readObject());
            }
        } catch (EOFException e) {
            // read all the records
            try {
                if (ois != null) {
                    ois.close();
                    this.optimizationLog.delete();
                }
            } catch (IOException e1) {
                /* do nothing */
                e.printStackTrace();
            }
        } catch (ClassNotFoundException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            /* the log is no longer needed, so remove it */
            this.optimizationLog.delete();
        }

        /* reconstruct the dynamic metrics */
        Map<String, DynamicMetric> newMetrics;
        Map<String, DynamicMetric> oldMetrics;
        int numNewDocs;
        int numOldDocs;
        newMetrics = this.dynamicMetrics;
        numNewDocs = this.flamdexReader.getNumDocs();
        for (OptimizationRecord opRec : records) {
            int[] newToOldIdMapping = opRec.mergedShards.get(0).newDocIdToOldDocId;
            oldMetrics = opRec.mergedShards.get(0).dynamicMetrics;
            numOldDocs = opRec.mergedShards.get(0).numDocs;

            for (Map.Entry<String, DynamicMetric> e : newMetrics.entrySet()) {
                DynamicMetric oldMetric = oldMetrics.get(e.getKey());
                DynamicMetric newMetric = e.getValue();
                if (oldMetric == null) {
                    oldMetric = new DynamicMetric(numOldDocs);
                }
                for (int i = 0; i < numNewDocs; i++) {
                    int oldId = newToOldIdMapping[i];
                    int value = newMetric.lookupSingleVal(i);
                    oldMetric.set(oldId, value);
                }

                oldMetrics.put(e.getKey(), oldMetric);
            }
            numNewDocs = numOldDocs;
            newMetrics = oldMetrics;
        }

        /* adjust the memory tracking */
        for (DynamicMetric dm : this.dynamicMetrics.values()) {
            memory.releaseMemory(dm.memoryUsed());
        }
        for (DynamicMetric dm : newMetrics.values()) {
            memory.claimMemory(dm.memoryUsed());
        }
        this.dynamicMetrics = newMetrics;

        try {
            /* close temp index reader */
            this.flamdexReaderRef.close();
        } catch (IOException e) {
            log.error("Could not close optimized reader");
        }

        /* reopen the original flamdex readers */
        this.flamdexReader = this.originalReader;
        this.flamdexReaderRef = this.originalReaderRef;
        this.originalReader = null;
        this.originalReaderRef = null;

        this.numDocs = this.flamdexReader.getNumDocs();

        /* push the stats back on */
        for (String stat : statsCopy) {
            if ("pop".equals(stat)) {
                this.popStat();
            } else {
                this.pushStat(stat);
            }
        }

        /* release the memory used by the log reading */
        this.memory.releaseMemory(memoryUse);

    }

    /**
     * export the current docId -> group lookup into an array
     *
     * @param array
     *            the array to export docIdToGroup into
     */
    public synchronized void exportDocIdToGroupId(int[] array) {
        if (array.length != docIdToGroup.size()) {
            throw new IllegalArgumentException("array length is invalid");
        }
        for (int i = array.length - 1; i >= 0; --i) {
            array[i] = docIdToGroup.get(i);
        }
    }

    @Override
    public synchronized long getTotalDocFreq(String[] intFields, String[] stringFields) {
        long ret = 0L;

        for (final String intField : intFields) {
            ret += flamdexReader.getIntTotalDocFreq(intField);
        }

        for (final String stringField : stringFields) {
            ret += flamdexReader.getStringTotalDocFreq(stringField);
        }

        return ret;
    }

    @Override
    public synchronized FTGSIterator getFTGSIterator(String[] intFields, String[] stringFields) {
        if (fieldZeroDocBitsets != null) {
            for (String intField : intFields) {
                if (!fieldZeroDocBitsets.containsKey(Pair.of(intField, true))) {
                    if (memory.claimMemory(FastBitSet.calculateMemoryUsage(32))) {
                        fieldZeroDocBitsets.put(Pair.of(intField, true), new FastBitSet(32));
                    } else {
                        log.warn("Insufficient memory, not allocating ftgs group zero bitset for field "
                                + intField);
                    }
                }
            }
            for (String stringField : stringFields) {
                if (!fieldZeroDocBitsets.containsKey(Pair.of(stringField, true))) {
                    if (memory.claimMemory(FastBitSet.calculateMemoryUsage(32))) {
                        fieldZeroDocBitsets.put(Pair.of(stringField, false), new FastBitSet(32));
                    } else {
                        log.warn("Insufficient memory, not allocating ftgs group zero bitset for field "
                                + stringField);
                    }
                }
            }
        }
        if (flamdexReader instanceof RawFlamdexReader) {
            return new RawFlamdexFTGSIterator(this, flamdexReaderRef.copy(), intFields,
                                              stringFields);
        }
        return new FlamdexFTGSIterator(this, flamdexReaderRef.copy(), intFields, stringFields);
    }

    @Override
    public FTGSIterator getSubsetFTGSIterator(Map<String, long[]> intFields, Map<String, String[]> stringFields) {
        if (flamdexReader instanceof RawFlamdexReader) {
            return new RawFlamdexSubsetFTGSIterator(this, flamdexReaderRef.copy(), intFields,
                    stringFields);
        }
        return new FlamdexSubsetFTGSIterator(this, flamdexReaderRef.copy(), intFields, stringFields);
    }

    public DocIterator getDocIterator(final String[] intFields, final String[] stringFields) throws ImhotepOutOfMemoryException {
        boolean shardOnlyContainsGroupZero = true;
        for (int group = 1; group < groupDocCount.length; group++) {
            if (groupDocCount[group] != 0) {
                shardOnlyContainsGroupZero = false;
                break;
            }
        }
        if (shardOnlyContainsGroupZero) {
            return emptyDocIterator();
        }

        final IntValueLookup[] intValueLookups = new IntValueLookup[intFields.length];
        final StringValueLookup[] stringValueLookups = new StringValueLookup[stringFields.length];
        try {
            for (int i = 0; i < intFields.length; i++) {
                intValueLookups[i] = flamdexReader.getMetric(intFields[i]);
            }
            for (int i = 0; i < stringFields.length; i++) {
                stringValueLookups[i] = flamdexReader.getStringLookup(stringFields[i]);
            }
        } catch (FlamdexOutOfMemoryException e) {
            for (IntValueLookup lookup : intValueLookups) {
                if (lookup != null) {
                    lookup.close();
                }
            }
            for (StringValueLookup lookup : stringValueLookups) {
                if (lookup != null) {
                    lookup.close();
                }
            }
            throw new ImhotepOutOfMemoryException();
        }
        return new DocIterator() {

            int[] groups = new int[1024];
            int n = groups.length;
            int bufferStart = -groups.length;
            int docId = -1;
            boolean done = false;

            public boolean next() {
                if (done) {
                    return false;
                }
                while (true) {
                    docId++;
                    if (docId - bufferStart >= n) {
                        if (!readGroups()) {
                            return endOfData();
                        }
                    }
                    if (groups[docId - bufferStart] != 0) {
                        return true;
                    }
                }
            }

            boolean endOfData() {
                done = true;
                return false;
            }

            public boolean readGroups() {
                bufferStart += n;
                n = Math.min(numDocs - bufferStart, groups.length);
                if (n <= 0) {
                    return false;
                }
                docIdToGroup.fillDocGrpBufferSequential(bufferStart, groups, n);
                return true;
            }

            public int getGroup() {
                return groups[docId - bufferStart];
            }

            int[] docIdRef = new int[1];
            long[] valueRef = new long[1];

            public long getInt(final int index) {
                docIdRef[0] = docId;
                intValueLookups[index].lookup(docIdRef, valueRef, 1);
                return valueRef[0];
            }

            public String getString(final int index) {
                return stringValueLookups[index].getString(docId);
            }

            public void close() throws IOException {
                for (IntValueLookup lookup : intValueLookups) {
                    if (lookup != null) {
                        lookup.close();
                    }
                }
                for (StringValueLookup lookup : stringValueLookups) {
                    if (lookup != null) {
                        lookup.close();
                    }
                }
            }
        };
    }

    private static DocIterator emptyDocIterator() {
        return new DocIterator() {
            @Override
            public boolean next() {
                return false;
            }

            @Override
            public int getGroup() {
                return 0;
            }

            @Override
            public long getInt(int index) {
                return 0;
            }

            @Override
            public String getString(int index) {
                return null;
            }

            @Override
            public void close() {

            }
        };
    }

    public RawFTGSIterator[] getFTGSIteratorSplits(final String[] intFields, final String[] stringFields) {
        final int numSplits = 16;

        final RawFTGSIterator[] ret = new RawFTGSIterator[numSplits];
        for (int i = 0; i < numSplits; i++) {
            ret[i] = getFTGSIteratorSplit(intFields, stringFields, i, numSplits);
        }

        return ret;
    }

    public synchronized RawFTGSIterator getFTGSIteratorSplit(final String[] intFields,
                                                             final String[] stringFields,
                                                             final int splitIndex,
                                                             final int numSplits) {
        if (ftgsIteratorSplits == null || ftgsIteratorSplits.isClosed()) {
            try {
                ftgsIteratorSplits = new FTGSSplitter(getFTGSIterator(intFields, stringFields), numSplits, numStats, "getIteratorSplitsLocalSession", 969168349, tempFileSizeBytesLeft);
            } catch (IOException e) {
                throw Throwables.propagate(e);
            }
        }
        return ftgsIteratorSplits.getFtgsIterators()[splitIndex];
    }

    @Override
    public RawFTGSIterator[] getSubsetFTGSIteratorSplits(Map<String, long[]> intFields, Map<String, String[]> stringFields) {
        final int numSplits = 16;

        final RawFTGSIterator[] ret = new RawFTGSIterator[numSplits];
        for (int i = 0; i < numSplits; i++) {
            ret[i] = getSubsetFTGSIteratorSplit(intFields, stringFields, i, numSplits);
        }

        return ret;
    }

    @Override
    public synchronized RawFTGSIterator getSubsetFTGSIteratorSplit(Map<String, long[]> intFields, Map<String, String[]> stringFields, int splitIndex, int numSplits) {
        if (ftgsIteratorSplits == null || ftgsIteratorSplits.isClosed()) {
            try {
                ftgsIteratorSplits = new FTGSSplitter(getSubsetFTGSIterator(intFields, stringFields), numSplits, numStats, "getIteratorSplitsLocalSession", 969168349, tempFileSizeBytesLeft);
            } catch (IOException e) {
                throw Throwables.propagate(e);
            }
        }
        return ftgsIteratorSplits.getFtgsIterators()[splitIndex];
    }

    public RawFTGSIterator mergeFTGSSplit(final String[] intFields,
                                          final String[] stringFields,
                                          final String sessionId,
                                          final InetSocketAddress[] nodes,
                                          final int splitIndex) {
        throw new UnsupportedOperationException();
    }

    @Override
    public RawFTGSIterator mergeSubsetFTGSSplit(Map<String, long[]> intFields, Map<String, String[]> stringFields, String sessionId, InetSocketAddress[] nodes, int splitIndex) {
        throw new UnsupportedOperationException();
    }

    @Override
