/*!
 * \file
 * \brief   Bitonic sort CUDA implementation header
 *
 * \author
 *    Christos Choutouridis AEM:8997
 *    <cchoutou@ece.auth.gr>
 */

#ifndef BITONICSORTCUDA_H_
#define BITONICSORTCUDA_H_

#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#include <cstdint>
#include <utility>
#include <stdexcept>
#include <algorithm>

#include "utils.hpp"

/*
 * Exported timers
 */
extern Timing Timer_total, Timer_memory, Timer_sorting;

using threadId_t = size_t;


/*
 * ============================== Sort utilities ==============================
 */

/*!
 * Returns the ascending or descending configuration (up/down phase) of the thread id
 * depending on the current depth
 *
 * @param tid   [threadId_t] The current thread
 * @param stage [size_t]     The current stage of the sorting network (same for each step)
 * @return      [bool]       True if we need ascending configuration, false otherwise
 */
__device__ inline bool ascending(threadId_t tid, size_t stage) noexcept {
    return !(tid & (1 << stage));
}

/*!
 * Returns the thread's partner for data exchange during the sorting network iterations
 * of Bitonic
 *
 * @param tid   [threadId_t] The current node
 * @param step  [size_t]     The step of the sorting network
 * @return      [threadId_t] The node id of the partner for data exchange
 */
__device__ inline threadId_t partner(threadId_t tid, size_t step) noexcept {
    return (tid ^ (1 << step));
}


/*!
 * Predicate to check if a node keeps the small numbers during the bitonic sort network exchange.
 *
 * @param tid       [threadId_t] The node for which we check
 * @param partner   [threadId_t] The partner of the data exchange
 * @param stage     [size_t]     The current stage of the sorting network (same for each step)
 * @return          [bool]       True if the node should keep the small values, false otherwise
 */

__device__ inline bool keepSmall(threadId_t tid, threadId_t partner, size_t stage) {
    return ascending(tid, stage) == (tid < partner);
}



/*
 * ============================== Sort algorithms ==============================
 */

/*!
 * Each thread can handle 2 points in the array. For each of these 2 points it may
 * - compare and exchange if needed
 * - copy data to local and back if needed
 */
static constexpr size_t SizeToThreadsRatio = 2;

/*!
 * Calculates the blocks needed for the entire sorting process
 *
 * @note
 *  This "redundant" little trick makes sure blocks are allocated for arraySizes that are not exact
 *  multipliers of config.blockSize.
 *  Even if we don't need it, we keep it in case we experiment with weird sizes in the future!
 *
 * @param arraySize [ArraySize_t]   The size of the entire array (in points)
 * @return          [size_t]        The number of blocks
 */
inline size_t NBlocks(ArraySize_t arraySize) {
	return (((arraySize + config.blockSize - 1) / config.blockSize) / SizeToThreadsRatio);
}


/*!
 * Exchange utility
 *
 * @tparam  ValueT  The underlying data type of the array items
 *
 * @param data      [ValueT*]    Pointer to data array
 * @param tid       [threadId_t] Current thread's index to data
 * @param pid       [threadId_t] Parents's index to data
 * @param keepSmall [bool]       Flag to indicate if current threads is keeping the small
 */
template <typename ValueT>
__device__ void exchange(ValueT* data, threadId_t tid, threadId_t partner, bool keepSmall) {
    if (( keepSmall && (data[tid] > data[partner])) ||
        (!keepSmall && (data[tid] < data[partner])) ) {
        ValueT temp = data[tid];
        data[tid] = data[partner];
        data[partner] = temp;
    }
}
#if CODE_VERSION == SERIAL

template <typename DataT>
void bitonicSort(DataT& data) {
    std::sort(data.begin(), data.end());
}


#elif CODE_VERSION == V0

/*!
 * This is the body of each thread. This function compare and exchange data
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__global__ void bitonicStep(ValueT* data, size_t n, size_t step, size_t stage) {
    threadId_t tid = threadIdx.x + blockIdx.x * blockDim.x; // Keep contiguous addressing to the first half of the array
    threadId_t pid = partner(tid, step);
    if (tid > pid) {
        // Shift to the other half of the array for global data
        tid += n / SizeToThreadsRatio;
        pid += n / SizeToThreadsRatio;
    }
    if ((tid < n) && (pid < n)) {       // Boundary check
        bool keep = keepSmall(tid, pid, stage);
        exchange(data, tid, pid, keep);
    }
}


/*!
 * A CUDA version of the Bitonic sort algorithm.
 *
 * @tparam DataT    A container type to hold data array. Should have .data() and .size() methods
 * @param data      [DataT&]    Reference to the container to sort
 */
template <typename DataT>
void bitonicSort(DataT& data) {
    using value_t = typename DataT::value_type;

    value_t* dev_data;
    auto size = data.size();

    Timer_memory.start();
    if (cudaMalloc(&dev_data, size * sizeof(value_t)) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not allocate memory");
    if (cudaMemcpy(dev_data, data.data(), size * sizeof(value_t), cudaMemcpyHostToDevice) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory to device");
    Timer_memory.stop();

    size_t Nth = config.blockSize;
    size_t Nbl = NBlocks(size);

    size_t Stages = static_cast<size_t>(log2(size));
    Timer_sorting.start();
    for (size_t stage = 1; stage <= Stages; ++stage) {
        for (size_t step = stage; step > 0; ) {
            --step;
            bitonicStep<<<Nbl, Nth>>>(dev_data, size, step, stage);
            cudaDeviceSynchronize();
        }
    }
    Timer_sorting.stop();

    Timer_memory.start();
    if (cudaMemcpy(data.data(), dev_data, size * sizeof(value_t), cudaMemcpyDeviceToHost) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory from device");
    cudaFree(dev_data);
    Timer_memory.stop();
}

#elif CODE_VERSION == V1

/*!
 * This is the body of each thread. This function compare and exchange data
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__device__ void interBlockStep_(ValueT* data, size_t n, size_t step, size_t stage) {
    /*
     * Here we skip blocks every time (one for SizeToThreadsRatio = 2)
     * And we use the neighbor block address indices for the other half of the threads
     */
    threadId_t tid = threadIdx.x + SizeToThreadsRatio * blockIdx.x * blockDim.x;
    threadId_t pid = partner(tid, step);
    if (tid > pid) {
        // Shift to the other half of the array for global data
        tid += blockDim.x;
        pid += blockDim.x;
    }
    if ((tid < n) && (pid < n)) {       // Boundary check
        bool keep = keepSmall(tid, pid, stage);
        exchange(data, tid, pid, keep);
    }
}

/*!
 * This is the version of the body that is called outside of the loop unrolling
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__global__ void interBlockStep(ValueT* data, size_t n, size_t step, size_t stage) {
    interBlockStep_(data, n, step, stage);
}

/*!
 * This is unrolled part of the bitonic double loop.
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__global__ void inBlockStep(ValueT* data, size_t n, size_t innerSteps, size_t stage) {
    for (size_t step = innerSteps + 1; step > 0; ) {
        --step;
        interBlockStep_(data, n, step, stage);
        __syncthreads();
    }
}


/*!
 * This is unrolled part of the bitonic double loop for the init phase where the entire
 * double loop can fit in one block with shared memory access.
 *
 * First each thread caches its corresponding data point from the current and the following data block.
 * After that we execute the pre-phase on the local data and then we write back to global memory.
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param stages    [size_t]    The number of stages to pre execute
 * @param maxStages [size_t]    The maximum number of stages for the entire sort
 */
template <typename ValueT>
__global__ void prephase(ValueT* data, size_t n, size_t stages, size_t maxStages) {
    for (size_t stage = 1; (stage <= stages) && (stage <= maxStages); ++stage) {
        for (size_t step = stage; step > 0; ) {
            --step;
            interBlockStep_(data, n, step, stage);
            __syncthreads();
        }
    }
}

/*!
 * A CUDA version of the Bitonic sort algorithm.
 *
 * @tparam DataT    A container type to hold data array. Should have .data() and .size() methods
 * @param data      [DataT&]    Reference to the container to sort
 */
template <typename DataT>
void bitonicSort(DataT& data) {
    using value_t = typename DataT::value_type;

    value_t* dev_data;
    auto size = data.size();

    Timer_memory.start();
    if (cudaMalloc(&dev_data, size * sizeof(value_t)) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not allocate memory");
    if (cudaMemcpy(dev_data, data.data(), size * sizeof(value_t), cudaMemcpyHostToDevice) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory to device");
    Timer_memory.stop();

    size_t Nth = config.blockSize;
    size_t Nbl = NBlocks(size);

    auto Stages          = static_cast<size_t>(log2(size));
    auto InnerBlockSteps = static_cast<size_t>(log2(Nth));
    size_t PrephaseStages= InnerBlockSteps + 1;

    Timer_sorting.start();
    prephase<<<Nbl, Nth>>>(dev_data, size, PrephaseStages, Stages);
    cudaDeviceSynchronize();
    for (size_t stage = PrephaseStages + 1; stage <= Stages; ++stage) {
        size_t step = stage - 1;
        for ( ; step > InnerBlockSteps; --step) {
            interBlockStep<<<Nbl, Nth>>>(dev_data, size, step, stage);
            cudaDeviceSynchronize();
        }
        inBlockStep<<<Nbl, Nth>>>(dev_data, size, step, stage);
        cudaDeviceSynchronize();
    }
    Timer_sorting.stop();

    Timer_memory.start();
    if (cudaMemcpy(data.data(), dev_data, size * sizeof(value_t), cudaMemcpyDeviceToHost) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory from device");
    cudaFree(dev_data);
    Timer_memory.stop();
}



#elif CODE_VERSION == V2

/*!
 * @return The memory that each block local threads can affect.
 *
 * @note
 *  Each block thread collection can exchange twice the size of data points.
 */
inline constexpr size_t effectiveBlockSize(size_t blockSize) {
    return SizeToThreadsRatio * blockSize;
}



/*!
 * Converts the global address of the data to the local shared memory array which is used
 * as cached memory to the unrolled part of the bitonic sort loop.
 *
 * @note
 *  Each block's thread collection can exchange twice the size of data points.
 *  These points get copied (cached) in the shared memory location. We use contiguous blocks
 *  both in global data memory and the shared memory buffer.
 *
 * @param gIndex    The global array index
 * @param blockDim  The block size (threads per block)
 * @return          The equivalent local address of the shared memory
 */
__device__ inline size_t toLocal(size_t gIndex, size_t blockDim) {
    return gIndex % (SizeToThreadsRatio * blockDim);
}

/*!
 * This is the version of the body that is called outside of the loop unrolling
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__global__ void interBlockStep(ValueT* data, size_t n, size_t step, size_t stage) {
    threadId_t tid = threadIdx.x + blockIdx.x * blockDim.x;     // Keep contiguous addressing to the first half of the array
    threadId_t pid = partner(tid, step);
    if (tid > pid) {
        // Shift to the other half of the array for global data
        tid += n / SizeToThreadsRatio;
        pid += n / SizeToThreadsRatio;
    }
    if ((tid < n) && (pid < n)) {       // Boundary check
        bool keep = keepSmall(tid, pid, stage);
        exchange(data, tid, pid, keep);
    }
}

/*!
 * This is unrolled part of the bitonic double loop.
 *
 * First each thread caches its corresponding data point from the current and the following data block.
 * After that we execute the loop unrolling on the local data and then we write back to global memory.
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param step      [size_t]    The current step of the current stage of bitonic sort
 * @param stage     [size_t]    The current stage of bitonic sort
 */
template <typename ValueT>
__global__ void inBlockStep(ValueT* data, size_t n, size_t innerSteps, size_t stage) {
    extern __shared__ ValueT shared_data[];

    /*
     * Global and local(shared) memory indices (calculated once)
     * Here we skip blocks every time (one for SizeToThreadsRatio = 2)
     * And we cache the neighbor block address indexes in local (shared) memory
     */
    threadId_t gIdx0 = threadIdx.x + SizeToThreadsRatio * blockIdx.x * blockDim.x;
    threadId_t lIdx0 = toLocal(gIdx0, blockDim.x);

    if (gIdx0 + blockDim.x >= n)   // Boundary check
        return;

    // Fetch to local memory the entire effective block size (2 positions for each thread)
    shared_data[lIdx0]              = data[gIdx0];
    shared_data[lIdx0 + blockDim.x] = data[gIdx0 + blockDim.x];
    __syncthreads();

    for (size_t step = innerSteps + 1; step > 0; ) {
        --step;

        // Init thread global and local indices
        threadId_t gIdx = gIdx0;
        threadId_t lIdx = lIdx0;
        // Find partner and keep-small configuration based on the global data positions
        threadId_t pIdx  = partner(gIdx, step);
        if (gIdx > pIdx) {
            // Shift inside effective block
            gIdx += blockDim.x; // global
            pIdx += blockDim.x;
            lIdx += blockDim.x;	// local
        }
        bool keep = keepSmall(gIdx, pIdx, stage);

        // Exchange data on local(shared) copy
        threadId_t lpIdx = toLocal(pIdx, blockDim.x);
        exchange(shared_data, lIdx, lpIdx, keep);
        __syncthreads();
    }

    // Write back to global memory (no sync here, there will be sync from host)
    data[gIdx0]              = shared_data[lIdx0];
    data[gIdx0 + blockDim.x] = shared_data[lIdx0 + blockDim.x];
}

/*!
 * This is unrolled part of the bitonic double loop for the init phase where the entire
 * double loop can fit in one block with shared memory access.
 *
 * First each thread caches its corresponding data point from the current and the following data block.
 * After that we execute the pre-phase on the local data and then we write back to global memory.
 *
 * @tparam ValueT   The underlying data type of the array items
 * @param data      [ValueT*]   Pointer to data array
 * @param n         [size_t]    The total size of the array
 * @param stages    [size_t]    The number of stages to pre execute
 * @param maxStages [size_t]    The maximum number of stages for the entire sort
 */
template <typename ValueT>
__global__ void prephase(ValueT* data, size_t n, size_t stages, size_t maxStages) {
    extern __shared__ ValueT shared_data[];

    /*
     * Global and local(shared) memory indices (calculated once)
     * Here we skip blocks every time (one for SizeToThreadsRatio = 2)
     * And we cache the neighbor block address indexes in local (shared) memory
     */
    threadId_t gIdx0 = threadIdx.x + SizeToThreadsRatio * blockIdx.x * blockDim.x;
    threadId_t lIdx0 = toLocal(gIdx0, blockDim.x);

    if (gIdx0 + blockDim.x >= n)   // Boundary check
        return;

    // Fetch to local memory the entire effective block size (2 positions for each thread)
    shared_data[lIdx0]              = data[gIdx0];
    shared_data[lIdx0 + blockDim.x] = data[gIdx0 + blockDim.x];
    __syncthreads();
    for (size_t stage = 1; (stage <= stages) && (stage <= maxStages); ++stage) {
        for (size_t step = stage; step > 0; ) {
            --step;

            // Init thread global and local indices
            threadId_t gIdx = gIdx0;
            threadId_t lIdx = lIdx0;
            // Find partner and keep-small configuration based on the global data positions
            threadId_t pIdx  = partner(gIdx, step);
            if (gIdx > pIdx) {
                // Shift inside effective block
                gIdx += blockDim.x; // global
                pIdx += blockDim.x;
                lIdx += blockDim.x;	// local
            }
            bool keep = keepSmall(gIdx, pIdx, stage);

            // Exchange data on local(shared) copy
            threadId_t lpIdx = toLocal(pIdx, blockDim.x);
            exchange(shared_data, lIdx, lpIdx, keep);
            __syncthreads();
        }
    }

    // Write back to global memory (no sync here, there will be sync from host)
    data[gIdx0]              = shared_data[lIdx0];
    data[gIdx0 + blockDim.x] = shared_data[lIdx0 + blockDim.x];
}

/*!
 * A CUDA version of the Bitonic sort algorithm.
 *
 * @tparam DataT    A container type to hold data array. Should have .data() and .size() methods
 * @param data      [DataT&]    Reference to the container to sort
 */
template <typename DataT>
void bitonicSort(DataT& data) {
    using value_t = typename DataT::value_type;

    value_t* dev_data;
    auto size = data.size();

    Timer_memory.start();
    if (cudaMalloc(&dev_data, size * sizeof(value_t)) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not allocate memory");
    if (cudaMemcpy(dev_data, data.data(), size * sizeof(value_t), cudaMemcpyHostToDevice) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory to device");
    Timer_memory.stop();

    size_t Nth = config.blockSize;
    size_t Nbl = NBlocks(size);
    size_t kernelMemSize = effectiveBlockSize(config.blockSize) * sizeof(value_t);

    auto Stages          = static_cast<size_t>(log2(size));
    auto InnerBlockSteps = static_cast<size_t>(log2(Nth));
    size_t PrephaseStages= InnerBlockSteps + 1;

    Timer_sorting.start();
    prephase<<<Nbl, Nth, kernelMemSize>>>(dev_data, size, PrephaseStages, Stages);
    cudaDeviceSynchronize();
    for (size_t stage = PrephaseStages + 1; stage <= Stages; ++stage) {
        size_t step = stage - 1;
        for ( ; step > InnerBlockSteps; --step) {
            interBlockStep<<<Nbl, Nth>>>(dev_data, size, step, stage);
            cudaDeviceSynchronize();
        }
        inBlockStep<<<Nbl, Nth, kernelMemSize>>>(dev_data, size, step, stage);
        cudaDeviceSynchronize();
    }
    Timer_sorting.stop();

    Timer_memory.start();
    if (cudaMemcpy(data.data(), dev_data, size * sizeof(value_t), cudaMemcpyDeviceToHost) != cudaSuccess)
        throw std::runtime_error("[CUDA] - Can not copy memory from device");
    cudaFree(dev_data);
    Timer_memory.stop();
}

#endif

#endif //BITONICSORTCUDA_H_
