/*
 * Copyright 2009-2010 by The Regents of the University of California
 * 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 from
 * 
 *     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 edu.uci.ics.hyracks.dataflow.std.group;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import edu.uci.ics.hyracks.api.comm.IFrameWriter;
import edu.uci.ics.hyracks.api.context.IHyracksStageletContext;
import edu.uci.ics.hyracks.api.dataflow.value.IBinaryComparator;
import edu.uci.ics.hyracks.api.dataflow.value.IBinaryComparatorFactory;
import edu.uci.ics.hyracks.api.dataflow.value.ISerializerDeserializer;
import edu.uci.ics.hyracks.api.dataflow.value.ITuplePartitionComputer;
import edu.uci.ics.hyracks.api.dataflow.value.ITuplePartitionComputerFactory;
import edu.uci.ics.hyracks.api.dataflow.value.RecordDescriptor;
import edu.uci.ics.hyracks.api.exceptions.HyracksDataException;
import edu.uci.ics.hyracks.dataflow.common.comm.io.FrameTupleAccessor;
import edu.uci.ics.hyracks.dataflow.common.comm.io.FrameTupleAppender;
import edu.uci.ics.hyracks.dataflow.common.comm.io.FrameTuplePairComparator;
import edu.uci.ics.hyracks.dataflow.common.comm.util.FrameUtils;

/**
 * An in-mem hash table for spillable grouping operations.
 * A table of {@link #Link}s are maintained in this object, and each row
 * of this table represents a hash partition.
 */
public class SpillableGroupingHashTable {

    /**
     * Context.
     */
    private final IHyracksStageletContext ctx;

    /**
     * Columns for group-by
     */
    private final int[] fields;

    /**
     * Key fields of records in the hash table (starting from 0
     * to the number of the key fields).
     * This is different from the key fields in the input records,
     * since these fields are extracted when being inserted into
     * the hash table.
     */
    private final int[] storedKeys;

    /**
     * Comparators: one for each column in {@link #groupFields}
     */
    private final IBinaryComparator[] comparators;

    /**
     * Record descriptor for the input tuple.
     */
    private final RecordDescriptor inRecordDescriptor;

    /**
     * Record descriptor for the partial aggregation result.
     */
    private final RecordDescriptor outputRecordDescriptor;

    /**
     * Accumulators in the main memory.
     */
    private ISpillableAccumulatingAggregator[] accumulators;

    /**
     * The hashing group table containing pointers to aggregators and also the
     * corresponding key tuples. So for each entry, there will be three integer
     * fields:
     * 1. The frame index containing the key tuple; 2. The tuple index inside of
     * the frame for the key tuple; 3. The index of the aggregator.
     * Note that each link in the table is a partition for the input records. Multiple
     * records in the same partition based on the {@link #tpc} are stored as
     * pointers.
     */
    private final Link[] table;

    /**
     * Number of accumulators.
     */
    private int accumulatorSize = 0;

    /**
     * Factory for the aggregators.
     */
    private final IAccumulatingAggregatorFactory aggregatorFactory;

    private final List<ByteBuffer> frames;

    private final ByteBuffer outFrame;

    /**
     * Frame appender for output frames in {@link #frames}.
     */
    private final FrameTupleAppender appender;

    /**
     * The count of used frames in the table.
     * Note that this cannot be replaced by {@link #frames} since frames will
     * not be removed after being created.
     */
    private int dataFrameCount;

    /**
     * Pointers for the sorted aggregators
     */
    private int[] tPointers;

    private static final int INIT_ACCUMULATORS_SIZE = 8;

    /**
     * The maximum number of frames available for this hashing group table.
     */
    private final int framesLimit;

    private final FrameTuplePairComparator ftpc;

    /**
     * A partition computer to partition the hashing group table.
     */
    private final ITuplePartitionComputer tpc;

    /**
     * Accessors for the tuples. Two accessors are necessary during the sort.
     */
    private final FrameTupleAccessor storedKeysAccessor1;
    private final FrameTupleAccessor storedKeysAccessor2;

    /**
     * Create a spillable grouping hash table.
     * 
     * @param ctx
     *            The context of the job.
     * @param fields
     *            Fields of keys for grouping.
     * @param comparatorFactories
     *            The comparators.
     * @param tpcf
     *            The partitioners. These are used to partition the incoming records into proper partition of the hash table.
     * @param aggregatorFactory
     *            The aggregators.
     * @param inRecordDescriptor
     *            Record descriptor for input data.
     * @param outputRecordDescriptor
     *            Record descriptor for output data.
     * @param framesLimit
     *            The maximum number of frames usable in the memory for hash table.
     * @param tableSize
     *            The size of the table, which specified the number of partitions of the table.
     */
    public SpillableGroupingHashTable(IHyracksStageletContext ctx, int[] fields,
            IBinaryComparatorFactory[] comparatorFactories, ITuplePartitionComputerFactory tpcf,
            IAccumulatingAggregatorFactory aggregatorFactory, RecordDescriptor inRecordDescriptor,
            RecordDescriptor outputRecordDescriptor, int framesLimit, int tableSize) {
        this.ctx = ctx;
        this.fields = fields;

        storedKeys = new int[fields.length];
        @SuppressWarnings("rawtypes")
        ISerializerDeserializer[] storedKeySerDeser = new ISerializerDeserializer[fields.length];

        // Note that after storing a record into the hash table, the index for the fields should
        // be updated. Here we assume that all these key fields are written at the beginning of 
        // the record, so their index should start from 0 and end at the length of the key fields.
        for (int i = 0; i < fields.length; ++i) {
            storedKeys[i] = i;
            storedKeySerDeser[i] = inRecordDescriptor.getFields()[fields[i]];
        }
        RecordDescriptor storedKeysRecordDescriptor = new RecordDescriptor(storedKeySerDeser);
        storedKeysAccessor1 = new FrameTupleAccessor(ctx.getFrameSize(), storedKeysRecordDescriptor);
        storedKeysAccessor2 = new FrameTupleAccessor(ctx.getFrameSize(), storedKeysRecordDescriptor);

        comparators = new IBinaryComparator[comparatorFactories.length];
        for (int i = 0; i < comparatorFactories.length; ++i) {
            comparators[i] = comparatorFactories[i].createBinaryComparator();
        }

        this.table = new Link[tableSize];

        this.aggregatorFactory = aggregatorFactory;
        accumulators = new ISpillableAccumulatingAggregator[INIT_ACCUMULATORS_SIZE];

        this.framesLimit = framesLimit;

        // Tuple pair comparator
        ftpc = new FrameTuplePairComparator(fields, storedKeys, comparators);

        // Partitioner
        tpc = tpcf.createPartitioner();

        this.inRecordDescriptor = inRecordDescriptor;
        this.outputRecordDescriptor = outputRecordDescriptor;
        frames = new ArrayList<ByteBuffer>();
        appender = new FrameTupleAppender(ctx.getFrameSize());

        dataFrameCount = -1;

        outFrame = ctx.allocateFrame();
    }

    public void reset() {
        dataFrameCount = -1;
        tPointers = null;
        // Reset the grouping hash table
        for (int i = 0; i < table.length; i++) {
            table[i] = new Link();
        }
    }

    public int getFrameCount() {
        return dataFrameCount;
    }

    /**
     * How to define pointers for the partial aggregation
     * 
     * @return
     */
    public int[] getTPointers() {
        return tPointers;
    }

    /**
     * Redefine the number of fields in the pointer.
     * Only two pointers are necessary for external grouping: one is to the
     * index of the hash table, and the other is to the row index inside of the
     * hash table.
     * 
     * @return
     */
    public int getPtrFields() {
        return 2;
    }

    public List<ByteBuffer> getFrames() {
        return frames;
    }

    /**
     * Set the working frame to the next available frame in the
     * frame list. There are two cases:<br>
     * 1) If the next frame is not initialized, allocate
     * a new frame.
     * 2) When frames are already created, they are recycled.
     * 
     * @return Whether a new frame is added successfully.
     */
    private boolean nextAvailableFrame() {
        // Return false if the number of frames is equal to the limit.
        if (dataFrameCount + 1 >= framesLimit)
            return false;

        if (frames.size() < framesLimit) {
            // Insert a new frame
            ByteBuffer frame = ctx.allocateFrame();
            frame.position(0);
            frame.limit(frame.capacity());
            frames.add(frame);
            appender.reset(frame, true);
            dataFrameCount++;
        } else {
            // Reuse an old frame
            dataFrameCount++;
            ByteBuffer frame = frames.get(dataFrameCount);
            frame.position(0);
            frame.limit(frame.capacity());
            appender.reset(frame, true);
        }
        return true;
    }

    /**
     * Insert a new record from the input frame.
     * 
     * @param accessor
     * @param tIndex
     * @return
     * @throws HyracksDataException
     */
    public boolean insert(FrameTupleAccessor accessor, int tIndex) throws HyracksDataException {
        if (dataFrameCount < 0)
            nextAvailableFrame();
        // Get the partition for the inserting tuple
        int entry = tpc.partition(accessor, tIndex, table.length);
        Link link = table[entry];
        if (link == null) {
            link = table[entry] = new Link();
        }
        // Find the corresponding aggregator from existing aggregators
        ISpillableAccumulatingAggregator aggregator = null;
        for (int i = 0; i < link.size; i += 3) {
            int sbIndex = link.pointers[i];
            int stIndex = link.pointers[i + 1];
            int saIndex = link.pointers[i + 2];
            storedKeysAccessor1.reset(frames.get(sbIndex));
            int c = ftpc.compare(accessor, tIndex, storedKeysAccessor1, stIndex);
            if (c == 0) {
                aggregator = accumulators[saIndex];
                break;
            }
        }
        // Do insert
        if (aggregator == null) {
            // Did not find the aggregator. Insert a new aggregator entry
            if (!appender.appendProjection(accessor, tIndex, fields)) {
                if (!nextAvailableFrame()) {
                    // If buffer is full, return false to trigger a run file
                    // write
                    return false;
                } else {
                    // Try to do insert after adding a new frame.
                    if (!appender.appendProjection(accessor, tIndex, fields)) {
                        throw new IllegalStateException();
                    }
                }
            }
            int sbIndex = dataFrameCount;
            int stIndex = appender.getTupleCount() - 1;
            if (accumulatorSize >= accumulators.length) {
                accumulators = Arrays.copyOf(accumulators, accumulators.length * 2);
            }
            int saIndex = accumulatorSize++;
            aggregator = accumulators[saIndex] = aggregatorFactory.createSpillableAggregator(ctx, inRecordDescriptor,
                    outputRecordDescriptor);
            aggregator.init(accessor, tIndex);
            link.add(sbIndex, stIndex, saIndex);
        }
        aggregator.accumulate(accessor, tIndex);
        return true;
    }

    /**
     * Sort partial results
     */
    public void sortFrames() {
        int totalTCount = 0;
        // Get the number of records
        for (int i = 0; i < table.length; i++) {
            if (table[i] == null)
                continue;
            totalTCount += table[i].size / 3;
        }
        // Start sorting:
        /*
         * Based on the data structure for the partial aggregates, the
         * pointers should be initialized.
         */
        tPointers = new int[totalTCount * getPtrFields()];
        // Initialize pointers
        int ptr = 0;
        // Maintain two pointers to each entry of the hashing group table
        for (int i = 0; i < table.length; i++) {
            if (table[i] == null)
                continue;
            for (int j = 0; j < table[i].size; j = j + 3) {
                tPointers[ptr * getPtrFields()] = i;
                tPointers[ptr * getPtrFields() + 1] = j;
                ptr++;
            }
        }
        // Sort using quick sort
        if (tPointers.length > 0) {
            sort(tPointers, 0, totalTCount);
        }
    }

    /**
     * @param writer
     * @throws HyracksDataException
     */
    public void flushFrames(IFrameWriter writer, boolean sorted) throws HyracksDataException {
        FrameTupleAppender appender = new FrameTupleAppender(ctx.getFrameSize());

        ISpillableAccumulatingAggregator aggregator = null;
        writer.open();
        appender.reset(outFrame, true);
        if (sorted) {
            sortFrames();
        }
        if (tPointers == null) {
            // Not sorted
            for (int i = 0; i < table.length; ++i) {
                Link link = table[i];
                if (link != null) {
                    for (int j = 0; j < link.size; j += 3) {
                        int bIndex = link.pointers[j];
                        int tIndex = link.pointers[j + 1];
                        int aIndex = link.pointers[j + 2];
                        ByteBuffer keyBuffer = frames.get(bIndex);
                        storedKeysAccessor1.reset(keyBuffer);
                        aggregator = accumulators[aIndex];
                        while (!aggregator.output(appender, storedKeysAccessor1, tIndex, storedKeys)) {
                            FrameUtils.flushFrame(outFrame, writer);
                            appender.reset(outFrame, true);
                        }
                    }
                }
            }
            if (appender.getTupleCount() != 0) {
                FrameUtils.flushFrame(outFrame, writer);
            }
            return;
        }
        int n = tPointers.length / getPtrFields();
        for (int ptr = 0; ptr < n; ptr++) {
            int tableIndex = tPointers[ptr * 2];
            int rowIndex = tPointers[ptr * 2 + 1];
            int frameIndex = table[tableIndex].pointers[rowIndex];
            int tupleIndex = table[tableIndex].pointers[rowIndex + 1];
            int aggregatorIndex = table[tableIndex].pointers[rowIndex + 2];
            // Get the frame containing the value
            ByteBuffer buffer = frames.get(frameIndex);
            storedKeysAccessor1.reset(buffer);

            // Get the aggregator
            aggregator = accumulators[aggregatorIndex];
            // Insert
            if (!aggregator.output(appender, storedKeysAccessor1, tupleIndex, fields)) {
                FrameUtils.flushFrame(outFrame, writer);
                appender.reset(outFrame, true);
                if (!aggregator.output(appender, storedKeysAccessor1, tupleIndex, fields)) {
                    throw new IllegalStateException();
                } else {
                    accumulators[aggregatorIndex] = null;
                }
            } else {
                accumulators[aggregatorIndex] = null;
            }
        }
        if (appender.getTupleCount() > 0) {
            FrameUtils.flushFrame(outFrame, writer);
        }
    }

    private void sort(int[] tPointers, int offset, int length) {
        int m = offset + (length >> 1);
        // Get table index
        int mTable = tPointers[m * 2];
        int mRow = tPointers[m * 2 + 1];
        // Get frame and tuple index
        int mFrame = table[mTable].pointers[mRow];
        int mTuple = table[mTable].pointers[mRow + 1];
        storedKeysAccessor1.reset(frames.get(mFrame));

        int a = offset;
        int b = a;
        int c = offset + length - 1;
        int d = c;
        while (true) {
            while (b <= c) {
                int bTable = tPointers[b * 2];
                int bRow = tPointers[b * 2 + 1];
                int bFrame = table[bTable].pointers[bRow];
                int bTuple = table[bTable].pointers[bRow + 1];
                storedKeysAccessor2.reset(frames.get(bFrame));
                int cmp = ftpc.compare(storedKeysAccessor2, bTuple, storedKeysAccessor1, mTuple);
                // int cmp = compare(tPointers, b, mi, mj, mv);
                if (cmp > 0) {
                    break;
                }
                if (cmp == 0) {
                    swap(tPointers, a++, b);
                }
                ++b;
            }
            while (c >= b) {
                int cTable = tPointers[c * 2];
                int cRow = tPointers[c * 2 + 1];
                int cFrame = table[cTable].pointers[cRow];
                int cTuple = table[cTable].pointers[cRow + 1];
                storedKeysAccessor2.reset(frames.get(cFrame));
                int cmp = ftpc.compare(storedKeysAccessor2, cTuple, storedKeysAccessor1, mTuple);
                // int cmp = compare(tPointers, c, mi, mj, mv);
                if (cmp < 0) {
                    break;
                }
                if (cmp == 0) {
                    swap(tPointers, c, d--);
                }
                --c;
            }
            if (b > c)
                break;
            swap(tPointers, b++, c--);
        }

        int s;
        int n = offset + length;
        s = Math.min(a - offset, b - a);
        vecswap(tPointers, offset, b - s, s);
        s = Math.min(d - c, n - d - 1);
        vecswap(tPointers, b, n - s, s);

        if ((s = b - a) > 1) {
            sort(tPointers, offset, s);
        }
        if ((s = d - c) > 1) {
            sort(tPointers, n - s, s);
        }
    }

    private void swap(int x[], int a, int b) {
        for (int i = 0; i < 2; ++i) {
            int t = x[a * 2 + i];
            x[a * 2 + i] = x[b * 2 + i];
            x[b * 2 + i] = t;
        }
    }

    private void vecswap(int x[], int a, int b, int n) {
        for (int i = 0; i < n; i++, a++, b++) {
            swap(x, a, b);
        }
    }

    /**
     * The pointers in the link store 3 int values for each entry in the
     * hashtable: (bufferIdx, tIndex, accumulatorIdx).
     * 
     * @author vinayakb
     */
    private static class Link {
        private static final int INIT_POINTERS_SIZE = 9;

        int[] pointers;
        int size;

        Link() {
            pointers = new int[INIT_POINTERS_SIZE];
            size = 0;
        }

        void add(int bufferIdx, int tIndex, int accumulatorIdx) {
            while (size + 3 > pointers.length) {
                pointers = Arrays.copyOf(pointers, pointers.length * 2);
            }
            pointers[size++] = bufferIdx;
            pointers[size++] = tIndex;
            pointers[size++] = accumulatorIdx;
        }

        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("[Size=" + size + "]");
            for (int i = 0; i < pointers.length; i = i + 3) {
                sb.append(pointers[i] + ",");
                sb.append(pointers[i + 1] + ",");
                sb.append(pointers[i + 2] + "; ");
            }
            return sb.toString();
        }
    }
}