/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.drill.exec.physical.impl.xsort;

import com.google.common.base.Preconditions;
import com.google.common.base.Stopwatch;
import com.google.common.collect.Queues;
import org.apache.drill.exec.exception.SchemaChangeException;
import org.apache.drill.exec.memory.BufferAllocator;
import org.apache.drill.exec.ops.FragmentContext;
import org.apache.drill.exec.record.RecordBatch;
import org.apache.drill.exec.record.VectorContainer;
import org.apache.drill.exec.record.selection.SelectionVector4;
import org.apache.hadoop.util.IndexedSortable;

import javax.inject.Named;
import java.util.Queue;

public abstract class MSortTemplate implements MSorter, IndexedSortable{
  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(MSortTemplate.class);

  private BufferAllocator allocator;
  private SelectionVector4 vector4;
  private SelectionVector4 aux;
  private long compares;
  private Queue<Integer> runStarts = Queues.newLinkedBlockingQueue();
  private Queue<Integer> newRunStarts;


  @Override
  public void setup(FragmentContext context, BufferAllocator allocator, SelectionVector4 vector4, VectorContainer hyperBatch) throws SchemaChangeException{
    this.allocator = allocator;
    // we pass in the local hyperBatch since that is where we'll be reading data.
    Preconditions.checkNotNull(vector4);
    this.vector4 = vector4.createNewWrapperCurrent();
    vector4.clear();
    doSetup(context, hyperBatch, null);
    runStarts.add(0);
    int batch = 0;
    for (int i = 0; i < this.vector4.getTotalCount(); i++) {
      int newBatch = this.vector4.get(i) >>> 16;
      if (newBatch == batch) {
        continue;
      } else if(newBatch == batch + 1) {
        runStarts.add(i);
        batch = newBatch;
      } else {
        throw new UnsupportedOperationException("Missing batch");
      }
    }
    BufferAllocator.PreAllocator preAlloc = allocator.getNewPreAllocator();
    preAlloc.preAllocate(4 * this.vector4.getTotalCount());
    aux = new SelectionVector4(preAlloc.getAllocation(), this.vector4.getTotalCount(), Character.MAX_VALUE);
  }

  private int merge(int leftStart, int rightStart, int rightEnd, int outStart) {
    int l = leftStart;
    int r = rightStart;
    int o = outStart;
    while (l < rightStart && r < rightEnd) {
      if (compare(l, r) <= 0) {
        aux.set(o++, vector4.get(l++));
      } else {
        aux.set(o++, vector4.get(r++));
      }
    }
    while (l < rightStart) {
      aux.set(o++, vector4.get(l++));
    }
    while (r < rightEnd){
      aux.set(o++, vector4.get(r++));
    }
    assert o == outStart + (rightEnd - leftStart);
    return o;
  }

  @Override
  public SelectionVector4 getSV4() {
    return vector4;
  }
  
  @Override
  public void sort(VectorContainer container){
    Stopwatch watch = new Stopwatch();
    watch.start();
    while (runStarts.size() > 1) {
      int outIndex = 0;
      newRunStarts = Queues.newLinkedBlockingQueue();
      newRunStarts.add(outIndex);
      int size = runStarts.size();
      for (int i = 0; i < size / 2; i++) {
        int left = runStarts.poll();
        int right = runStarts.poll();
        Integer end = runStarts.peek();
        if (end == null) end = vector4.getTotalCount();
        outIndex = merge(left, right, end, outIndex);
        if (outIndex < vector4.getTotalCount()) newRunStarts.add(outIndex);
      }
      SelectionVector4 tmp = aux.createNewWrapperCurrent();
      aux.clear();
      aux = this.vector4.createNewWrapperCurrent();
      vector4.clear();
      this.vector4 = tmp.createNewWrapperCurrent();
      tmp.clear();
      runStarts = newRunStarts;
    }
    aux.clear();
  }

  @Override
  public void swap(int sv0, int sv1) {
    int tmp = vector4.get(sv0);
    vector4.set(sv0, vector4.get(sv1));
    vector4.set(sv1, tmp);
  }
  
  @Override
  public int compare(int leftIndex, int rightIndex) {
    int sv1 = vector4.get(leftIndex);
    int sv2 = vector4.get(rightIndex);
    compares++;
    return doEval(sv1, sv2);
  }

  public abstract void doSetup(@Named("context") FragmentContext context, @Named("incoming") VectorContainer incoming, @Named("outgoing") RecordBatch outgoing);
  public abstract int doEval(@Named("leftIndex") int leftIndex, @Named("rightIndex") int rightIndex);

}

@@ -4166,24 +4166,132 @@
 x);%0A      %7D%0A
+      if (outIndex %3C vector4.getTotalCount()) %7B%0A        copyRun(outIndex, vector4.getTotalCount());%0A      %7D%0A
       Select
@@ -4559,24 +4559,154 @@
 ear();%0A  %7D%0A%0A
+  private void copyRun(int start, int end) %7B%0A    for (int i = start; i %3C end; i++) %7B%0A      aux.set(i, vector4.get(i));%0A    %7D%0A  %7D%0A%0A
   @Override%0A

copy unmerged records in merge sort to new buffer