#pragma once

#include "../cuda_flow.hpp"
#include "../cuda_capturer.hpp"
#include "../cuda_meta.hpp"

/**
@file cuda_sort.hpp
@brief CUDA sort algorithm include file
*/

namespace tf::detail {

// ----------------------------------------------------------------------------
// odd-even sort in register
// ----------------------------------------------------------------------------

/**
@private
@brief counts the number of leading zeros starting from the most significant bit
*/
constexpr int cuda_clz(int x) {
  for(int i = 31; i >= 0; --i) {
    if((1<< i) & x) {
      return 31 - i;
    }
  }
  return 32;
}

/**
@private
@brief finds log2(x) and optionally round up to the next integer logarithm.
*/
constexpr int cuda_find_log2(int x, bool round_up = false) {
  int a = 31 - cuda_clz(x);
  if(round_up) {
    a += !is_pow2(x);
  }
  return a;
}

/** @private */
template<typename T, unsigned vt, typename C>
__device__ auto cuda_odd_even_sort(
  cudaArray<T, vt> x, C comp, int flags = 0
) { 
  cuda_iterate<vt>([&](auto I) {
    #pragma unroll
    for(auto i = 1 & I; i < vt - 1; i += 2) {
      if((0 == ((2<< i) & flags)) && comp(x[i + 1], x[i]))
        cuda_swap(x[i], x[i + 1]);
    }
  });
  return x;
}

/** @private */
template<typename K, typename V, unsigned vt, typename C>
__device__ auto cuda_odd_even_sort(
  cudaKVArray<K, V, vt> x, C comp, int flags = 0
) { 
  cuda_iterate<vt>([&](auto I) {
    #pragma unroll
    for(auto i = 1 & I; i < vt - 1; i += 2) {
      if((0 == ((2<< i) & flags)) && comp(x.keys[i + 1], x.keys[i])) {
        cuda_swap(x.keys[i], x.keys[i + 1]);
        cuda_swap(x.vals[i], x.vals[i + 1]);
      }
    }
  });
  return x;
}

// ----------------------------------------------------------------------------
// range check
// ----------------------------------------------------------------------------

/** @private */
__device__ inline int cuda_out_of_range_flags(int first, int vt, int count) {
  int out_of_range = min(vt, first + vt - count);
  int head_flags = 0;
  if(out_of_range > 0) {
    const int mask = (1<< vt) - 1;
    head_flags = mask & (~mask>> out_of_range);
  }
  return head_flags;
}

/** @private */
__device__ inline auto cuda_compute_merge_sort_frame(
  unsigned partition, unsigned coop, unsigned spacing
) {

  unsigned size = spacing * (coop / 2);
  unsigned start = ~(coop - 1) & partition;
  unsigned a_begin = spacing * start;
  unsigned b_begin = spacing * start + size;

  return cudaMergeRange {
    a_begin,
    a_begin + size,
    b_begin,
    b_begin + size
  };
}

/** @private */
__device__ inline auto cuda_compute_merge_sort_range(
  unsigned count, unsigned partition, unsigned coop, unsigned spacing
) {

  auto frame = cuda_compute_merge_sort_frame(partition, coop, spacing);

  return cudaMergeRange {
    frame.a_begin,
    min(count, frame.a_end),
    min(count, frame.b_begin),
    min(count, frame.b_end)
  };
}

/** @private */
__device__ inline auto cuda_compute_merge_sort_range(
  unsigned count, unsigned partition, unsigned coop, unsigned spacing, 
  unsigned mp0, unsigned mp1
) {

  auto range = cuda_compute_merge_sort_range(count, partition, coop, spacing);

  // Locate the diagonal from the start of the A sublist.
  unsigned diag = spacing * partition - range.a_begin;

  // The end partition of the last cta for each merge operation is computed
  // and stored as the begin partition for the subsequent merge. i.e. it is
  // the same partition but in the wrong coordinate system, so its 0 when it
  // should be listSize. Correct that by checking if this is the last cta
  // in this merge operation.
  if(coop - 1 != ((coop - 1) & partition)) {
    range.a_end = range.a_begin + mp1;
    range.b_end = min(count, range.b_begin + diag + spacing - mp1);
  }

  range.a_begin = range.a_begin + mp0;
  range.b_begin = min(count, range.b_begin + diag - mp0);

  return range;
}

/** @private */
template<unsigned nt, unsigned vt, typename K, typename V>
struct cudaBlockSort {

  static constexpr bool has_values = !std::is_same<V, cudaEmpty>::value;
  static constexpr unsigned num_passes = log2(nt);

  /** @private */
  union Storage {
    K keys[nt * vt + 1];
    V vals[nt * vt];
  };

  static_assert(is_pow2(nt), "cudaBlockSort requires pow2 number of threads");

  template<typename C>
  __device__ auto merge_pass(
    cudaKVArray<K, V, vt> x, 
    unsigned tid, unsigned count, unsigned pass, 
    C comp, Storage& storage
  ) const {

    // Divide the CTA's keys into lists.
    unsigned coop = 2 << pass;
    auto range = cuda_compute_merge_sort_range(count, tid, coop, vt);
    unsigned diag = vt * tid - range.a_begin;

    // Store the keys into shared memory for searching.
    cuda_reg_to_shared_thread<nt, vt>(x.keys, tid, storage.keys);
    
    // Search for the merge path for this thread within its list.
    auto mp = cuda_merge_path<cudaMergeBoundType::LOWER>(
      storage.keys, range, diag, comp
    );

    // Run a serial merge and return.
    auto merge = cuda_serial_merge<cudaMergeBoundType::LOWER, vt>(
      storage.keys, range.partition(mp, diag), comp
    );
    x.keys = merge.keys;

    if(has_values) {
      // Reorder values through shared memory.
      cuda_reg_to_shared_thread<nt, vt>(x.vals, tid, storage.vals);
      x.vals = cuda_shared_gather<nt, vt>(storage.vals, merge.indices);
    }

    return x;
  }

  template<typename C>
  __device__ auto block_sort(cudaKVArray<K, V, vt> x, 
    unsigned tid, unsigned count, C comp, Storage& storage
  ) const {

    // Sort the inputs within each thread. If any threads have fewer than
    // vt items, use the segmented sort network to prevent out-of-range
    // elements from contaminating the sort.
    if(count < nt * vt) {
      auto head_flags = cuda_out_of_range_flags(vt * tid, vt, count);
      x = cuda_odd_even_sort(x, comp, head_flags);
    } else {
      x = cuda_odd_even_sort(x, comp);
    }

    // Merge threads starting with a pair until all values are merged.
    for(unsigned pass = 0; pass < num_passes; ++pass) {
      x = merge_pass(x, tid, count, pass, comp, storage);
    }
    
    return x;
  }
};

/** @private */
template<typename P, typename K, typename C>
void cuda_merge_sort_partitions(
  P&& p, K keys, unsigned count, 
  unsigned coop, unsigned spacing, C comp, unsigned* buf
) {
  
  // bufer size is num_partitions + 1
  unsigned num_partitions = (count + spacing - 1) / spacing + 1;
  
  const unsigned nt = 128;
  const unsigned vt = 1;
  const unsigned nv = nt * vt;

  unsigned B = (num_partitions + nv - 1) / nv;  // nt = 128, vt = 1
  
  cuda_kernel<<<B, nt, 0, p.stream()>>>([=] __device__ (auto tid, auto bid) {
    auto range = cuda_get_tile(bid, nt * vt, num_partitions);
    cuda_strided_iterate<nt, vt>([=](auto, auto j) {
      auto index = j + range.begin;
      auto range = cuda_compute_merge_sort_range(count, index, coop, spacing);
      auto diag = min(spacing * index, count) - range.a_begin;
      buf[index] = cuda_merge_path<cudaMergeBoundType::LOWER>(
        keys + range.a_begin, range.a_count(), 
        keys + range.b_begin, range.b_count(), 
        diag, comp
      );
    }, tid, range.count());  
  });
}

/** @private */
template<typename P, typename K_it, typename V_it, typename C>
void merge_sort_loop(
  P&& p, K_it keys_input, V_it vals_input, unsigned count, C comp, void* buf
) {
  
  using K = typename std::iterator_traits<K_it>::value_type;
  using V = typename std::iterator_traits<V_it>::value_type;
  using E = std::decay_t<P>;

  const bool has_values = !std::is_same<V, cudaEmpty>::value;

  unsigned B = (count + E::nv - 1) / E::nv;
  unsigned R = cuda_find_log2(B, true);

  K* keys_output    {nullptr};
  V* vals_output    {nullptr};
  unsigned *mp_data {nullptr};
  
  if(R) {
    keys_output = (K*)(buf);
    if(has_values) {
      vals_output = (V*)(keys_output + count);
      mp_data = (unsigned*)(vals_output + count);
    }
    else {
      mp_data = (unsigned*)(keys_output + count);
    }
  }

  //cudaScopedDeviceMemory<K> keys_temp(R ? count : 0);
  //auto keys_output = keys_temp.data();
  ////std::cout << "keys_output = " << keys_temp.size()*sizeof(K) << std::endl;

  //cudaScopedDeviceMemory<V> vals_temp((has_values && R) ? count : 0);
  //auto vals_output = vals_temp.data();
  //std::cout << "vals_output = " << vals_temp.size()*sizeof(V) << std::endl;

  auto keys_blocksort = (1 & R) ? keys_output : keys_input;
  auto vals_blocksort = (1 & R) ? vals_output : vals_input;

  //printf("B=%u, R=%u\n", B, R);

  cuda_kernel<<<B, E::nt, 0, p.stream()>>>([=] __device__ (auto tid, auto bid) {
    
    using sort_t = cudaBlockSort<E::nt, E::vt, K, V>;

    __shared__ union {
      typename sort_t::Storage sort;
      K keys[E::nv];
      V vals[E::nv];
    } shared;

    auto tile = cuda_get_tile(bid, E::nv, count);

    // Load the keys and values.
    cudaKVArray<K, V, E::vt> unsorted;
    unsorted.keys = cuda_mem_to_reg_thread<E::nt, E::vt>(
      keys_input + tile.begin, tid, tile.count(), shared.keys
    );

    if(has_values) {
      unsorted.vals = cuda_mem_to_reg_thread<E::nt, E::vt>(
        vals_input + tile.begin, tid, tile.count(), shared.vals
      );
    }

    // Blocksort.
    auto sorted = sort_t().block_sort(unsorted, tid, tile.count(), comp, shared.sort);

    // Store the keys and values.
    cuda_reg_to_mem_thread<E::nt, E::vt>(
      sorted.keys, tid, tile.count(), keys_blocksort + tile.begin, shared.keys
    );

    if(has_values) {
      cuda_reg_to_mem_thread<E::nt, E::vt>(
        sorted.vals, tid, tile.count(), vals_blocksort + tile.begin, shared.vals
      );
    }
  });

  if(R == 0) {
    return;
  }

  // merge passes

  if(1 & R) {
    std::swap(keys_input, keys_output);
    std::swap(vals_input, vals_output);
  }
  
  // number of partitions
  //unsigned num_partitions = B + 1;
  //cudaScopedDeviceMemory<unsigned> mem(num_partitions);
  //auto mp_data = mem.data();
  //std::cout << "num_partitions = " << (B+1)*sizeof(unsigned) << std::endl;
  
  for(unsigned pass = 0; pass < R; ++pass) {

    unsigned coop = 2 << pass;

    cuda_merge_sort_partitions(
      p, keys_input, count, coop, E::nv, comp, mp_data
    );
    
    cuda_kernel<<<B, E::nt, 0, p.stream()>>>([=]__device__(auto tid, auto bid) {

      __shared__ union {
        K keys[E::nv + 1];
        unsigned indices[E::nv];
      } shared;

      auto tile = cuda_get_tile(bid, E::nv, count);

      // Load the range for this CTA and merge the values into register.
      auto range = cuda_compute_merge_sort_range(
        count, bid, coop, E::nv, mp_data[bid + 0], mp_data[bid + 1]
      );

      auto merge = block_merge_from_mem<cudaMergeBoundType::LOWER, E::nt, E::vt>(
        keys_input, keys_input, range, tid, comp, shared.keys
      );

      // Store merged values back out.
      cuda_reg_to_mem_thread<E::nt>(
        merge.keys, tid, tile.count(), keys_output + tile.begin, shared.keys
      );

      if(has_values) {
        // Transpose the indices from thread order to strided order.
        auto indices = cuda_reg_thread_to_strided<E::nt>(
          merge.indices, tid, shared.indices
        );

        // Gather the input values and merge into the output values.
        cuda_transfer_two_streams_strided<E::nt>(
          vals_input + range.a_begin, range.a_count(), 
          vals_input + range.b_begin, range.b_count(),
          indices, tid, vals_output + tile.begin
        );
      }
    });

    std::swap(keys_input, keys_output);
    std::swap(vals_input, vals_output);
  }
}

}  // end of namespace tf::detail ---------------------------------------------

namespace tf {

/** 
@brief queries the buffer size in bytes needed to call sort kernels 
       for the given number of elements

@tparam P execution policy type
@tparam K key type
@tparam V value type (default tf::cudaEmpty)

@param count number of keys/values to sort

The function is used to allocate a buffer for calling tf::cuda_sort.

*/
template <typename P, typename K, typename V = cudaEmpty>
unsigned cuda_sort_buffer_size(unsigned count) {

  using E = std::decay_t<P>;

  const bool has_values = !std::is_same<V, cudaEmpty>::value;

  unsigned B = (count + E::nv - 1) / E::nv;
  unsigned R = detail::cuda_find_log2(B, true);

  return R ? (count * sizeof(K) + (has_values ? count*sizeof(V) : 0) +
             (B+1)*sizeof(unsigned)) : 0;
}

// ----------------------------------------------------------------------------
// key-value sort
// ----------------------------------------------------------------------------

/**
@brief performs asynchronous key-value sort on a range of items

@tparam P execution policy type
@tparam K_it key iterator type
@tparam V_it value iterator type
@tparam C comparator type

@param p execution policy 
@param k_first iterator to the beginning of the key range
@param k_last iterator to the end of the key range
@param v_first iterator to the beginning of the value range
@param comp binary comparator
@param buf pointer to the temporary buffer

Sorts key-value elements in <tt>[k_first, k_last)</tt> and 
<tt>[v_first, v_first + (k_last - k_first))</tt> into ascending key order
using the given comparator @c comp.
If @c i and @c j are any two valid iterators in <tt>[k_first, k_last)</tt> 
such that @c i precedes @c j, and @c p and @c q are iterators in 
<tt>[v_first, v_first + (k_last - k_first))</tt> corresponding to 
@c i and @c j respectively, then <tt>comp(*j, *i)</tt> evaluates to @c false.

For example, assume:
  + @c keys are <tt>{1, 4, 2, 8, 5, 7}</tt>
  + @c values are <tt>{'a', 'b', 'c', 'd', 'e', 'f'}</tt>

After sort:
  + @c keys are <tt>{1, 2, 4, 5, 7, 8}</tt>
  + @c values are <tt>{'a', 'c', 'b', 'e', 'f', 'd'}</tt>

*/
template<typename P, typename K_it, typename V_it, typename C>
void cuda_sort_by_key(
  P&& p, K_it k_first, K_it k_last, V_it v_first, C comp, void* buf
) {

  unsigned N = std::distance(k_first, k_last);

  if(N <= 1) {
    return;
  }

  detail::merge_sort_loop(p, k_first, v_first, N, comp, buf);
}

// ----------------------------------------------------------------------------
// key sort
// ----------------------------------------------------------------------------

/**
@brief performs asynchronous key-only sort on a range of items

@tparam P execution policy type
@tparam K_it key iterator type
@tparam C comparator type

@param p execution policy 
@param k_first iterator to the beginning of the key range
@param k_last iterator to the end of the key range
@param comp binary comparator
@param buf pointer to the temporary buffer

This method is equivalent to tf::cuda_sort_by_key without values.

*/
template<typename P, typename K_it, typename C>
void cuda_sort(P&& p, K_it k_first, K_it k_last, C comp, void* buf) {
  cuda_sort_by_key(p, k_first, k_last, (cudaEmpty*)nullptr, comp, buf);
}

// ----------------------------------------------------------------------------
// cudaFlow
// ----------------------------------------------------------------------------

// Function: sort
template <typename I, typename C>
cudaTask cudaFlow::sort(I first, I last, C comp) {
  return capture([=](cudaFlowCapturer& cap){
    cap.make_optimizer<cudaLinearCapturing>();
    cap.sort(first, last, comp);
  });
}

// Function: sort
template <typename I, typename C>
void cudaFlow::sort(cudaTask task, I first, I last, C comp) {
  capture(task, [=](cudaFlowCapturer& cap){
    cap.make_optimizer<cudaLinearCapturing>();
    cap.sort(first, last, comp);
  });
}

// Function: sort_by_key
template <typename K_it, typename V_it, typename C>
cudaTask cudaFlow::sort_by_key(K_it k_first, K_it k_last, V_it v_first, C comp) {
  return capture([=](cudaFlowCapturer& cap){
    cap.make_optimizer<cudaLinearCapturing>();
    cap.sort_by_key(k_first, k_last, v_first, comp);
  });
}

// Function: sort_by_key
template <typename K_it, typename V_it, typename C>
void cudaFlow::sort_by_key(
  cudaTask task, K_it k_first, K_it k_last, V_it v_first, C comp
) {
  capture(task, [=](cudaFlowCapturer& cap){
    cap.make_optimizer<cudaLinearCapturing>();
    cap.sort_by_key(k_first, k_last, v_first, comp);
  });
}

// ----------------------------------------------------------------------------
// cudaFlowCapturer
// ----------------------------------------------------------------------------

// Function: sort
template <typename I, typename C>
cudaTask cudaFlowCapturer::sort(I first, I last, C comp) {

  using K = typename std::iterator_traits<I>::value_type;

  auto bufsz = cuda_sort_buffer_size<cudaDefaultExecutionPolicy, K>(
    std::distance(first, last)
  );

  return on([=, buf=MoC{cudaScopedDeviceMemory<std::byte>(bufsz)}] 
  (cudaStream_t stream) mutable {
    cuda_sort(
      cudaDefaultExecutionPolicy{stream}, first, last, comp, buf.get().data()
    );
  });
}

// Function: sort
template <typename I, typename C>
void cudaFlowCapturer::sort(cudaTask task, I first, I last, C comp) {
  
  using K = typename std::iterator_traits<I>::value_type;

  auto bufsz = cuda_sort_buffer_size<cudaDefaultExecutionPolicy, K>(
    std::distance(first, last)
  );

  on(task, [=, buf=MoC{cudaScopedDeviceMemory<std::byte>(bufsz)}] 
  (cudaStream_t stream) mutable {
    cuda_sort(
      cudaDefaultExecutionPolicy{stream}, first, last, comp, buf.get().data()
    );
  });
}

// Function: sort_by_key
template <typename K_it, typename V_it, typename C>
cudaTask cudaFlowCapturer::sort_by_key(
  K_it k_first, K_it k_last, V_it v_first, C comp
) {

  using K = typename std::iterator_traits<K_it>::value_type;
  using V = typename std::iterator_traits<V_it>::value_type;

  auto bufsz = cuda_sort_buffer_size<cudaDefaultExecutionPolicy, K, V>(
    std::distance(k_first, k_last)
  );

  return on([=, buf=MoC{cudaScopedDeviceMemory<std::byte>(bufsz)}] 
  (cudaStream_t stream) mutable {
    cuda_sort_by_key(cudaDefaultExecutionPolicy{stream}, 
      k_first, k_last, v_first, comp, buf.get().data()
    );
  });
}

// Function: sort_by_key
template <typename K_it, typename V_it, typename C>
void cudaFlowCapturer::sort_by_key(
  cudaTask task, K_it k_first, K_it k_last, V_it v_first, C comp
) {

  using K = typename std::iterator_traits<K_it>::value_type;
  using V = typename std::iterator_traits<V_it>::value_type;

  auto bufsz = cuda_sort_buffer_size<cudaDefaultExecutionPolicy, K, V>(
    std::distance(k_first, k_last)
  );

  on(task, [=, buf=MoC{cudaScopedDeviceMemory<std::byte>(bufsz)}] 
  (cudaStream_t stream) mutable {
    cuda_sort_by_key(cudaDefaultExecutionPolicy{stream}, 
      k_first, k_last, v_first, comp, buf.get().data()
    );
  });
}


}  // end of namespace tf -----------------------------------------------------

