#include "THC.h"
#include "THCReduceApplyUtils.cuh"
#include "THCTensorCopy.h"
#include "THCTensorMath.h"
#include "THCAsmUtils.cuh"
#include "THCScanUtils.cuh"
#include "THCTensorTypeUtils.cuh"
#include <algorithm> // for std::min

#if CUDA_VERSION >= 7000
#include <thrust/system/cuda/execution_policy.h>
#endif

// Converts a float to an integer representation with the same
// sorting; i.e., for floats f1, f2:
// if f1 < f2 then convert(f1) < convert(f2)
// We use this to enable radix selection of floating-point values.
// This also gives a relative order for NaNs, but that's ok, as they
// will all be adjacent
struct FloatToSortedInt {
  inline __host__ __device__ FloatToSortedInt() {}

  inline __device__ unsigned int convert(float v) const {
    unsigned int x = __float_as_int(v);
    unsigned int mask = (x & 0x80000000) ? 0xffffffff : 0x80000000;

    return (x ^ mask);
  }

  inline __device__ float deconvert(unsigned int v) const {
    unsigned int mask = (v & 0x80000000) ? 0x80000000 : 0xffffffff;

    return __int_as_float(v ^ mask);
  }
};

// This function counts the distribution of all input values in a
// slice we are selecting by radix digit at `radixDigitPos`, but only
// those that pass the filter `((v & desiredMask) == desired)`.
// This produces and broadcasts the seen counts for a single block only.
// `smem` must have at least `RadixSize` elements.
template <typename DataType, typename BitDataType,
          typename IndexType, typename CountType,
          typename RadixConverter, int RadixSize, int RadixBits>
__device__ void countRadixUsingMask(const RadixConverter& conv,
                                    CountType counts[RadixSize],
                                    CountType* smem,
                                    BitDataType desired,
                                    BitDataType desiredMask,
                                    int radixDigitPos,
                                    IndexType sliceSize,
                                    IndexType withinSliceStride,
                                    DataType* data) {
  // Clear out per-thread counts from a previous round
#pragma unroll
  for (int i = 0; i < RadixSize; ++i) {
    counts[i] = 0;
  }

  if (threadIdx.x < RadixSize) {
    smem[threadIdx.x] = 0;
  }
  __syncthreads();

  // Scan over all the data. Upon a read, the warp will accumulate
  // counts per each digit in the radix using warp voting.
  for (IndexType i = threadIdx.x; i < sliceSize; i += blockDim.x) {
    BitDataType val = conv.convert(doLdg(&data[i * withinSliceStride]));

    bool hasVal = ((val & desiredMask) == desired);
    unsigned int digitInRadix = getBitfield(val, radixDigitPos, RadixBits);

#pragma unroll
    for (unsigned int j = 0; j < RadixSize; ++j) {
      bool vote = hasVal && (digitInRadix == j);
      counts[j] += __popc(__ballot(vote));
    }
  }

  // Now, for each warp, sum values
  if (getLaneId() == 0) {
#pragma unroll
    for (unsigned int i = 0; i < RadixSize; ++i) {
      atomicAdd(&smem[i], counts[i]);
    }
  }

  __syncthreads();

  // For each thread, read in the total counts
#pragma unroll
  for (unsigned int i = 0; i < RadixSize; ++i) {
    counts[i] = smem[i];
  }

  __syncthreads();
}

// Over what radix we are selecting values
#define RADIX_BITS 2 // digits are base-(2 ^ RADIX_BITS)
#define RADIX_SIZE 4 // 2 ^ RADIX_BITS
#define RADIX_MASK (RADIX_SIZE - 1)

// This finds the unique value `v` that matches the pattern
// ((v & desired) == desiredMask) in our sorted int format
template <typename DataType, typename IndexType, typename RadixConverter>
__device__ float findPattern(const RadixConverter& conv,
                             DataType* smem,
                             DataType* data,
                             IndexType sliceSize,
                             IndexType withinSliceStride,
                             unsigned int desired,
                             unsigned int desiredMask) {
  if (threadIdx.x < 32) {
    smem[threadIdx.x] = (DataType) 0;
  }
  __syncthreads();

  // All threads participate in the loop, in order to sync on the flag
  IndexType numIterations = THCRoundUp(sliceSize, (IndexType) blockDim.x);
  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < sliceSize);
    DataType v = inRange ? doLdg(&data[i * withinSliceStride]) : (DataType) 0;

    if (inRange && ((conv.convert(v) & desiredMask) == desired)) {
      // There should not be conflicts if we are using findPattern,
      // since the result is unique
      smem[0] = (DataType) 1;
      smem[1] = v; // can't use val as the flag, since it could be 0
    }

    __syncthreads();

    DataType found = smem[0];
    DataType val = smem[1];

    __syncthreads();

    // Check to see if a thread found the value
    if (found != (DataType) 0) {
      // all threads return this value
      return val;
    }
  }

  // should not get here
  assert(false);
  return (DataType) 0;
}

// Returns the top-Kth element found in the data using radix selection
template <typename DataType, typename BitDataType, typename IndexType,
          typename RadixConverter, bool Order>
__device__ void radixSelect(const RadixConverter& conv,
                            DataType* data,
                            IndexType k,
                            IndexType sliceSize,
                            IndexType withinSliceStride,
                            int* smem,
                            DataType* topK) {
  // Per-thread buckets into which we accumulate digit counts in our
  // radix
  int counts[RADIX_SIZE];

  // We only consider elements x such that (x & desiredMask) == desired
  // Initially, we consider all elements of the array, so the above
  // statement is true regardless of input.
  unsigned int desired = 0;
  unsigned int desiredMask = 0;

  // We are looking for the top kToFind-th element when iterating over
  // digits; this count gets reduced by elimination when counting
  // successive digits
  int kToFind = k;

  // We start at the most significant digit in our radix, scanning
  // through to the least significant digit
#pragma unroll
  for (int digitPos = sizeof(BitDataType) * 8 - RADIX_BITS;
       digitPos >= 0;
       digitPos -= RADIX_BITS) {

    // Count radix distribution for the current position and reduce
    // across all threads
    countRadixUsingMask<DataType, BitDataType,
                        IndexType, int, RadixConverter,
                        RADIX_SIZE, RADIX_BITS>(
                          conv, counts, smem,
                          desired, desiredMask, digitPos,
                          sliceSize, withinSliceStride, data);

    // All threads participate in the comparisons below to know the
    // final result

#define CHECK_RADIX(i)                                                  \
    int count = counts[i];                                              \
                                                                        \
    /* All threads have the same value in counts here, so all */        \
    /* threads will return from the function. */                        \
    if (count == 1 && kToFind == 1) {                                   \
      /* There is a unique answer. */                                   \
      desired = setBitfield(desired, i, digitPos, RADIX_BITS);          \
      desiredMask =                                                     \
        setBitfield(desiredMask, RADIX_MASK, digitPos, RADIX_BITS);     \
                                                                        \
      /* The answer is now the unique element v such that: */           \
      /* (v & desiredMask) == desired */                                \
      /* However, we do not yet know what the actual element is. We */  \
      /* need to perform a search through the data to find the */       \
      /* element that matches this pattern. */                          \
      *topK = findPattern<DataType, IndexType, RadixConverter>(         \
        conv, (float*) smem, data, sliceSize,                           \
        withinSliceStride, desired, desiredMask);                       \
      return;                                                           \
    }                                                                   \
                                                                        \
    if (count >= kToFind) {                                             \
      desired = setBitfield(desired, i, digitPos, RADIX_BITS);          \
      desiredMask =                                                     \
        setBitfield(desiredMask, RADIX_MASK, digitPos, RADIX_BITS);     \
                                                                        \
      /* The top-Kth element v must now be one such that: */            \
      /* (v & desiredMask == desired) */                                \
      /* but we haven't narrowed it down; we must check the next */     \
      /* least-significant digit */                                     \
      break;                                                            \
    }                                                                   \
                                                                        \
    kToFind -= count                                                    \

    if (Order) {
      // Process in descending order
#pragma unroll
      for (int i = RADIX_SIZE - 1; i >= 0; --i) {
        CHECK_RADIX(i);
      }
    } else {
      // Process in ascending order
#pragma unroll
      for (int i = 0; i < RADIX_SIZE; ++i) {
        CHECK_RADIX(i);
      }
    }
#undef CHECK_RADIX
  } // end digitPos for

  // There is no unique result, but there is a non-unique result
  // matching `desired` exactly
  *topK = conv.deconvert(desired);
}

template <typename IndexType, int Dim, bool Order>
__global__ void gatherTopK(TensorInfo<float, IndexType> input,
                           IndexType inputSliceSize,
                           IndexType outputSliceSize, // aka `k`

                           IndexType numInputSlices,
                           IndexType inputWithinSliceStride,

                           TensorInfo<float, IndexType> topK,
                           IndexType numTopKSlices,
                           IndexType topKWithinSliceStride,

                           TensorInfo<long, IndexType> indices,
                           IndexType indicesWithinSliceStride) {
  // Indices are limited to integer fp precision, so counts can fit in
  // int32, regardless of IndexType
  __shared__ int smem[32]; // one per each warp, up to warp limit

  IndexType slice = getLinearBlockId<IndexType>();
  if (slice >= numInputSlices) {
    return;
  }

  // Find the start offset for our slice
  IndexType sliceStartIndex =
    IndexToOffset<float, IndexType, Dim>::get(slice, input);
  IndexType topKSliceStartIndex =
    IndexToOffset<float, IndexType, Dim>::get(slice, topK);
  IndexType indicesSliceStartIndex =
    IndexToOffset<long, IndexType, Dim>::get(slice, indices);

  float* inputSliceStart = &input.data[sliceStartIndex];
  float* topKSliceStart = &topK.data[topKSliceStartIndex];
  long* indicesSliceStart = &indices.data[indicesSliceStartIndex];

  // Find the k-th highest element in our input
  float topKValue = -1.0f;
  radixSelect<float, unsigned int, IndexType, FloatToSortedInt, Order>(
    FloatToSortedInt(),
    inputSliceStart, outputSliceSize,
    inputSliceSize, inputWithinSliceStride,
    smem, &topKValue);

  // Every value that is strictly less/greater than `pattern`
  // (depending on sort dir) in sorted int format is in the top-K.
  // The top-K value itself might not be unique.
  //
  // Since there are a variable number of elements that we see that
  // are within the top-k, we don't know at what index to write out
  // the resulting values.
  // In order to get this, we perform an exclusive prefix sum of
  // `hasTopK`. This will return the resulting index into which we
  // need to write the result, if a thread has a result.

  // All threads need to participate in the loop and the prefix sum,
  // but not necessarily in the load; hence loop bounds being rounded
  // up to a multiple of the block dim.
  IndexType numIterations = THCRoundUp(inputSliceSize, (IndexType) blockDim.x);
  IndexType writeIndexStart = 0;

  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < inputSliceSize);
    float v =
      inRange ? doLdg(&inputSliceStart[i * inputWithinSliceStride]) : 0.0f;
    bool hasTopK;
    if (Order) {
      hasTopK = inRange && (v > topKValue);
    } else {
      hasTopK = inRange && (v < topKValue);
    }

    int index;
    int carry;
    exclusiveBinaryPrefixSum<int, true>(smem, hasTopK, &index, &carry);

    if (hasTopK) {
      int writeIndex = writeIndexStart + index;
      assert(writeIndex < outputSliceSize);

      IndexType topKOffset = writeIndex * topKWithinSliceStride;
      IndexType indexOffset = writeIndex * indicesWithinSliceStride;

      topKSliceStart[topKOffset] = v;
      indicesSliceStart[indexOffset] = i + TH_INDEX_BASE; // to Lua index
    }

    writeIndexStart += carry;
  }

  // We need to fill in the rest with actual == top-K values.
  // The number that we need is outputSliceSize -
  // writeIndexStart. There might be more than that number available,
  // in which case we have to choose the first seen set. We do this
  // via a prefix sum to calculate indices for writing results.
  assert(outputSliceSize >= writeIndexStart);
  IndexType topKRemaining = (outputSliceSize - writeIndexStart);

  for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
    bool inRange = (i < inputSliceSize);
    float v =
      inRange ? doLdg(&inputSliceStart[i * inputWithinSliceStride]) : 0.0f;
    bool hasTopK = inRange && (v == topKValue);

    int index;
    int carry;
    exclusiveBinaryPrefixSum<int, true>(smem, hasTopK, &index, &carry);

    if (hasTopK && index < topKRemaining) {
      int writeIndex = writeIndexStart + index;
      assert(writeIndex < outputSliceSize);

      IndexType topKOffset = writeIndex * topKWithinSliceStride;
      IndexType indexOffset = writeIndex * indicesWithinSliceStride;

      topKSliceStart[topKOffset] = v;
      indicesSliceStart[indexOffset] = i + TH_INDEX_BASE; // to Lua index
    }

    if (carry >= topKRemaining) {
      break;
    }

    topKRemaining -= carry;
    writeIndexStart += carry;
  }
}

#undef RADIX_BITS
#undef RADIX_SIZE
#undef RADIX_MASK

THC_API void THCudaTensor_topk(THCState* state,
                               THCudaTensor *topK,
                               THCudaLongTensor *indices,
                               THCudaTensor *input,
                               long k, int dim, int dir, int sorted) {
  THAssert(topK != NULL && indices != NULL && input != NULL);
  THAssert(THCudaTensor_checkGPU(state, 3, topK, indices, input));
  THCCheckTensorDims(state, topK, 2);
  long dims = THCudaLongTensor_nDimension(state, indices);
  THArgCheck(dims <= MAX_CUTORCH_DIMS, 2, CUTORCH_DIM_WARNING);
  THCCheckTensorDims(state, input, 2);

  int numDims = THCudaTensor_nDimension(state, input);
  THArgCheck(dim >= 0 && dim < numDims, 3, "dim not in range");

  long sliceSize = THCudaTensor_size(state, input, dim);
  THArgCheck(k > 0 && k <= sliceSize, 2, "k not in range for dimension");

  // Build the output size, which is the dim being selected set to
  // size k
  THLongStorage* topKSize = THCudaTensor_newSizeOf(state, input);
  THLongStorage_set(topKSize, dim, k);
  THCudaTensor_resize(state, topK, topKSize, NULL);
  THCudaLongTensor_resize(state, indices, topKSize, NULL);
  THLongStorage_free(topKSize);

#define RUN_K(INDEX_T, DIM, DIR)                                        \
  gatherTopK<INDEX_T, DIM, DIR>                                         \
    <<<grid, block, 0, THCState_getCurrentStream(state)>>>(             \
      inputInfo,                                                        \
      sliceSize,                                                        \
      k,                                                                \
      inputSlices,                                                      \
      /* The actual dimension that the k-selection is running in */     \
      /* may have changed from collapseDims() */                        \
      inputInfo.strides[collapseInputDim],                              \
      topKInfo,                                                         \
      topKSlices,                                                       \
      topKInfo.strides[collapseTopKDim],                                \
      indicesInfo,                                                      \
      indicesInfo.strides[collapseIndicesDim])

#define RUN_DIR(INDEX_T, DIM)                   \
  if (dir) {                                    \
    RUN_K(INDEX_T, DIM, true);                  \
  } else {                                      \
    RUN_K(INDEX_T, DIM, false);                 \
  }

#define RUN_DIM(INDEX_T)                        \
  if (allDims == 1) {                           \
    RUN_DIR(INDEX_T, 1);                        \
  } else if (allDims == 2) {                    \
    RUN_DIR(INDEX_T, 2);                        \
  } else if (allDims == 3) {                    \
    RUN_DIR(INDEX_T, 3);                        \
  } else {                                      \
    RUN_DIR(INDEX_T, -1);                       \
  }

#define RUN_T(INDEX_T)                                                  \
  TensorInfo<float, INDEX_T> inputInfo =                                \
    getTensorInfo<THCudaTensor, INDEX_T>(state, input);                 \
  TensorInfo<float, INDEX_T> topKInfo =                                 \
    getTensorInfo<THCudaTensor, INDEX_T>(state, topK);                  \
  TensorInfo<long, INDEX_T> indicesInfo =                               \
    getTensorInfo<THCudaLongTensor, INDEX_T>(state, indices);           \
                                                                        \
  /* We use these structures solely to find the offset to */            \
  /* each slice we are operating on */                                  \
  inputInfo.sizes[dim] = 1;                                             \
  topKInfo.sizes[dim] = 1;                                              \
  indicesInfo.sizes[dim] = 1;                                           \
                                                                        \
  /* Collapse all other dims */                                         \
  int collapseInputDim = inputInfo.collapseDims(dim);                   \
  int collapseTopKDim = topKInfo.collapseDims(dim);                     \
  int collapseIndicesDim = indicesInfo.collapseDims(dim);               \
                                                                        \
  long inputSlices = 1;                                                 \
  long topKSlices = 1;                                                  \
  for (int i = 0; i < numDims; ++i) {                                   \
    inputSlices *= inputInfo.sizes[i];                                  \
    topKSlices *= topKInfo.sizes[i];                                    \
  }                                                                     \
                                                                        \
  dim3 grid;                                                            \
  if (!THC_getGridFromTiles(inputSlices, grid)) {                       \
    THError("Slice to sort is too large");                              \
  }                                                                     \
                                                                        \
  dim3 block(std::min(THCRoundUp(sliceSize, 32L), 1024L));              \
                                                                        \
  /* This is used as a template parameter to calculate indices. */      \
  /* We only specialize it if all collapsed dim sizes are the */        \
  /* same; otherwise, we use -1 which is the specialization */          \
  /* parameter for arbitrary dimensions */                              \
  int allDims = inputInfo.dims;                                         \
  if (topKInfo.dims != allDims || indicesInfo.dims != allDims) {        \
    allDims = -1;                                                       \
  }                                                                     \
                                                                        \
  RUN_DIM(INDEX_T);

  // Based on required index size, run the algorithm with the
  // appropriate index type
  if (TensorUtils<THCudaTensor>::canUse32BitIndexMath(state, input) &&
      TensorUtils<THCudaTensor>::canUse32BitIndexMath(state, topK) &&
      TensorUtils<THCudaLongTensor>::canUse32BitIndexMath(state, indices)) {
    RUN_T(unsigned int);
  } else {
    RUN_T(unsigned long);
  }
#undef RUN_T
#undef RUN_DIM
#undef RUN_DIR
#undef RUN_K

  // Sort the results if the user wants them sorted, since our
  // selection routine does not ensure sorting
  if (sorted) {
    // FIXME: the k/v inplace sort along slice only works for size <=
    // 2048 at the moment
    if (sliceSize <= 2048) {
      // This avoids any memory allocations and performs all sorting
      // work inplace along the slice
      THCudaTensor_sortKeyValueInplace(state, topK, indices, dim, dir);
    } else {
      // Depend upon the backup sort that returns indices, which we
      // can use in conjunction with gather to produce the original
      // indices.
      // This is not the most efficient implementation, especially since
      // there are memory allocations performed here. If the user desires
      // greater performance, they should torch.gather() the results
      // themselves using the reported indices, providing previously
      // allocated tensors to receive the results.
      THCudaTensor* sortedTopK = THCudaTensor_new(state);
      THCudaLongTensor* sortedIndices = THCudaLongTensor_new(state);
      THCudaTensor_sort(state, sortedTopK, sortedIndices, topK, dim, dir);

      THCudaLongTensor* sortedTopKIndices = THCudaLongTensor_new(state);

      THCudaLongTensor_resizeAs(state, sortedTopKIndices, indices);
      THCudaLongTensor_gather(state, sortedTopKIndices, indices, dim, sortedIndices);

      THCudaTensor_freeCopyTo(state, sortedTopK, topK);
      THCudaLongTensor_freeCopyTo(state, sortedTopKIndices, indices);
      THCudaLongTensor_free(state, sortedIndices);
    }
  }

  THCudaCheck(cudaGetLastError());
}