#pragma once

#include "reduce.hpp"

/**
@file taskflow/cuda/algorithm/scan.hpp
@brief CUDA scan algorithm include file
*/

namespace tf::detail {

// ----------------------------------------------------------------------------
// scan
// ----------------------------------------------------------------------------

/** @private */
inline constexpr unsigned cudaScanRecursionThreshold = 8;

/** @private */
enum class cudaScanType : int {
  EXCLUSIVE = 1,
  INCLUSIVE
};

/** @private */
template<typename T, unsigned vt = 0, bool is_array = (vt > 0)>
struct cudaScanResult {
  T scan;
  T reduction;
};

/** @private */
template<typename T, unsigned vt>
struct cudaScanResult<T, vt, true> {
  cudaArray<T, vt> scan;
  T reduction;
};

//-----------------------------------------------------------------------------

/** @private */
template<unsigned nt, typename T>
struct cudaBlockScan {

  const static unsigned num_warps  = nt / CUDA_WARP_SIZE;
  const static unsigned num_passes = log2(nt);
  const static unsigned capacity   = nt + num_warps;

  /** @private */
  union storage_t {
    T data[2 * nt];
    struct { T threads[nt], warps[num_warps]; };
  };

  // standard scan
  template<typename op_t>
  __device__ cudaScanResult<T> operator ()(
    unsigned tid,
    T x,
    storage_t& storage,
    unsigned count = nt,
    op_t op = op_t(),
    T init = T(),
    cudaScanType type = cudaScanType::EXCLUSIVE
  ) const;

  // vectorized scan. accepts multiple values per thread and adds in
  // optional global carry-in
  template<unsigned vt, typename op_t>
  __device__ cudaScanResult<T, vt> operator()(
    unsigned tid,
    cudaArray<T, vt> x,
    storage_t& storage,
    T carry_in = T(),
    bool use_carry_in = false,
    unsigned count = nt,
    op_t op = op_t(),
    T init = T(),
    cudaScanType type = cudaScanType::EXCLUSIVE
  ) const;
};

// standard scan
template <unsigned nt, typename T>
template<typename op_t>
__device__ cudaScanResult<T> cudaBlockScan<nt, T>::operator () (
  unsigned tid, T x, storage_t& storage, unsigned count, op_t op,
  T init, cudaScanType type
) const {

  unsigned first = 0;
  storage.data[first + tid] = x;
  __syncthreads();

  cuda_iterate<num_passes>([&](auto pass) {
    if(auto offset = 1<<pass; tid >= offset) {
      x = op(storage.data[first + tid - offset], x);
    }
    first = nt - first;
    storage.data[first + tid] = x;
    __syncthreads();
  });

  cudaScanResult<T> result;
  result.reduction = storage.data[first + count - 1];
  result.scan = (tid < count) ?
    (cudaScanType::INCLUSIVE == type ? x :
      (tid ? storage.data[first + tid - 1] : init)) :
    result.reduction;
  __syncthreads();

  return result;
}

// vectorized scan block
template <unsigned nt, typename T>
template<unsigned vt, typename op_t>
__device__ cudaScanResult<T, vt> cudaBlockScan<nt, T>::operator()(
  unsigned tid,
  cudaArray<T, vt> x,
  storage_t& storage,
  T carry_in,
  bool use_carry_in,
  unsigned count, op_t op,
  T init,
  cudaScanType type
) const {

  // Start with an inclusive scan of the in-range elements.
  if(count >= nt * vt) {
    cuda_iterate<vt>([&](auto i) {
      x[i] = i ? op(x[i], x[i - 1]) : x[i];
    });
  } else {
    cuda_iterate<vt>([&](auto i) {
      auto index = vt * tid + i;
      x[i] = i ?
        ((index < count) ? op(x[i], x[i - 1]) : x[i - 1]) :
        (x[i] = (index < count) ? x[i] : init);
    });
  }

  // Scan the thread-local reductions for a carry-in for each thread.
  auto result = operator()(
    tid, x[vt - 1], storage,
    (count + vt - 1) / vt, op, init, cudaScanType::EXCLUSIVE
  );

  // Perform the scan downsweep and add both the global carry-in and the
  // thread carry-in to the values.
  if(use_carry_in) {
    result.reduction = op(carry_in, result.reduction);
    result.scan = tid ? op(carry_in, result.scan) : carry_in;
  } else {
    use_carry_in = tid > 0;
  }

  cudaArray<T, vt> y;
  cuda_iterate<vt>([&](auto i) {
    if(cudaScanType::EXCLUSIVE == type) {
      y[i] = i ? x[i - 1] : result.scan;
      if(use_carry_in && i > 0) y[i] = op(result.scan, y[i]);
    } else {
      y[i] = use_carry_in ? op(x[i], result.scan) : x[i];
    }
  });

  return cudaScanResult<T, vt> { y, result.reduction };
}

/**
@private
@brief single-pass scan for small input
 */
template <typename P, typename I, typename O, typename C>
void cuda_single_pass_scan(
  P&& p,
  cudaScanType scan_type,
  I input,
  unsigned count,
  O output,
  C op
  //reduction_it reduction,
) {

  using T = typename std::iterator_traits<O>::value_type;
  using E = std::decay_t<P>;

  // Small input specialization. This is the non-recursive branch.
  cuda_kernel<<<1, E::nt, 0, p.stream()>>>([=] __device__ (auto tid, auto bid) {

    using scan_t = cudaBlockScan<E::nt, T>;

    __shared__ union {
      typename scan_t::storage_t scan;
      T values[E::nv];
    } shared;

    auto carry_in = T();
    for(unsigned cur = 0; cur < count; cur += E::nv) {
      // Cooperatively load values into register.
      auto count2 = min(count - cur, E::nv);

      auto x = cuda_mem_to_reg_thread<E::nt, E::vt>(input + cur,
        tid, count2, shared.values);

      auto result = scan_t()(tid, x, shared.scan,
        carry_in, cur > 0, count2, op, T(), scan_type);

      // Store the scanned values back to global memory.
      cuda_reg_to_mem_thread<E::nt, E::vt>(result.scan, tid, count2,
        output + cur, shared.values);

      // Roll the reduction into carry_in.
      carry_in = result.reduction;
    }

    // Store the carry-out to the reduction pointer. This may be a
    // discard_iterator_t if no reduction is wanted.
    //if(!tid) *reduction = carry_in;
  });
}

/**
@private

@brief main scan loop
*/
template<typename P, typename I, typename O, typename C>
void cuda_scan_loop(
  P&& p,
  cudaScanType scan_type,
  I input,
  unsigned count,
  O output,
  C op,
  //reduction_it reduction,
  void* ptr
) {

  using E = std::decay_t<P>;
  using T = typename std::iterator_traits<O>::value_type;

  T* buffer = static_cast<T*>(ptr);

  //launch_t::cta_dim(context).B(count);
  unsigned B = (count + E::nv - 1) / E::nv;

  if(B > cudaScanRecursionThreshold) {

    //cudaDeviceVector<T> partials(B);
    //auto buffer = partials.data();

    // upsweep phase
    cuda_kernel<<<B, E::nt, 0, p.stream()>>>([=] __device__ (auto tid, auto bid) {

      __shared__ typename cudaBlockReduce<E::nt, T>::Storage shm;

      // Load the tile's data into register.
      auto tile = cuda_get_tile(bid, E::nv, count);
      auto x = cuda_mem_to_reg_strided<E::nt, E::vt>(
        input + tile.begin, tid, tile.count()
      );

      // Reduce the thread's values into a scalar.
      T scalar;
      cuda_strided_iterate<E::nt, E::vt>(
        [&] (auto i, auto j) { scalar = i ? op(scalar, x[i]) : x[0]; },
        tid, tile.count()
      );

      // Reduce across all threads.
      auto all_reduce = cudaBlockReduce<E::nt, T>()(
        tid, scalar, shm, tile.count(), op
      );

      // Store the final reduction to the partials.
      if(!tid) {
        buffer[bid] = all_reduce;
      }
    });

    // recursively call scan
    //cuda_scan_loop(p, cudaScanType::EXCLUSIVE, buffer, B, buffer, op, S);
    cuda_scan_loop(
      p, cudaScanType::EXCLUSIVE, buffer, B, buffer, op, buffer+B
    );

    // downsweep: perform an intra-tile scan and add the scan of the partials
    // as carry-in
    cuda_kernel<<<B, E::nt, 0, p.stream()>>>([=] __device__ (auto tid, auto bid) {

      using scan_t = cudaBlockScan<E::nt, T>;

      __shared__ union {
        typename scan_t::storage_t scan;
        T values[E::nv];
      } shared;

      // Load a tile to register in thread order.
      auto tile = cuda_get_tile(bid, E::nv, count);
      auto x = cuda_mem_to_reg_thread<E::nt, E::vt>(
        input + tile.begin, tid, tile.count(), shared.values
      );

      // Scan the array with carry-in from the partials.
      auto y = scan_t()(tid, x, shared.scan,
        buffer[bid], bid > 0, tile.count(), op, T(),
        scan_type).scan;

      // Store the scanned values to the output.
      cuda_reg_to_mem_thread<E::nt, E::vt>(
        y, tid, tile.count(), output + tile.begin, shared.values
      );
    });
  }
  // Small input specialization. This is the non-recursive branch.
  else {
    cuda_single_pass_scan(p, scan_type, input, count, output, op);
  }
}

}  // namespace tf::detail ----------------------------------------------------

namespace tf {

// Function: scan_bufsz
template <unsigned NT, unsigned VT>  
template <typename T>
unsigned cudaExecutionPolicy<NT, VT>::scan_bufsz(unsigned count) {
  unsigned B = num_blocks(count);
  unsigned n = 0;
  for(auto b=B; b>detail::cudaScanRecursionThreshold; b=num_blocks(b)) {
    n += b;
  }
  return n*sizeof(T);
}


/**
@brief performs asynchronous inclusive scan over a range of items

@tparam P execution policy type
@tparam I input iterator
@tparam O output iterator
@tparam C binary operator type

@param p execution policy
@param first iterator to the beginning of the input range
@param last iterator to the end of the input range
@param output iterator to the beginning of the output range
@param op binary operator to apply to scan
@param buf pointer to the temporary buffer

*/
template<typename P, typename I, typename O, typename C>
void cuda_inclusive_scan(
  P&& p, I first, I last, O output, C op, void* buf
) {

  unsigned count = std::distance(first, last);

  if(count == 0) {
    return;
  }

  // launch the scan loop
  detail::cuda_scan_loop(
    p, detail::cudaScanType::INCLUSIVE, first, count, output, op, buf
  );
}

/**
@brief performs asynchronous inclusive scan over a range of transformed items

@tparam P execution policy type
@tparam I input iterator
@tparam O output iterator
@tparam C binary operator type
@tparam U unary operator type

@param p execution policy
@param first iterator to the beginning of the input range
@param last iterator to the end of the input range
@param output iterator to the beginning of the output range
@param bop binary operator to apply to scan
@param uop unary operator to apply to transform each item before scan
@param buf pointer to the temporary buffer

*/
template<typename P, typename I, typename O, typename C, typename U>
void cuda_transform_inclusive_scan(
  P&& p, I first, I last, O output, C bop, U uop, void* buf
) {

  using T = typename std::iterator_traits<O>::value_type;

  unsigned count = std::distance(first, last);

  if(count == 0) {
    return;
  }

  // launch the scan loop
  detail::cuda_scan_loop(
    p, detail::cudaScanType::INCLUSIVE,
    cuda_make_load_iterator<T>([=]__device__(auto i){ return uop(*(first+i)); }),
    count, output, bop, buf
  );
}

/**
@brief performs asynchronous exclusive scan over a range of items

@tparam P execution policy type
@tparam I input iterator
@tparam O output iterator
@tparam C binary operator type

@param p execution policy
@param first iterator to the beginning of the input range
@param last iterator to the end of the input range
@param output iterator to the beginning of the output range
@param op binary operator to apply to scan
@param buf pointer to the temporary buffer

*/
template<typename P, typename I, typename O, typename C>
void cuda_exclusive_scan(
  P&& p, I first, I last, O output, C op, void* buf
) {

  unsigned count = std::distance(first, last);

  if(count == 0) {
    return;
  }

  // launch the scan loop
  detail::cuda_scan_loop(
    p, detail::cudaScanType::EXCLUSIVE, first, count, output, op, buf
  );
}

/**
@brief performs asynchronous exclusive scan over a range of items

@tparam P execution policy type
@tparam I input iterator
@tparam O output iterator
@tparam C binary operator type
@tparam U unary operator type

@param p execution policy
@param first iterator to the beginning of the input range
@param last iterator to the end of the input range
@param output iterator to the beginning of the output range
@param bop binary operator to apply to scan
@param uop unary operator to apply to transform each item before scan
@param buf pointer to the temporary buffer

*/
template<typename P, typename I, typename O, typename C, typename U>
void cuda_transform_exclusive_scan(
  P&& p, I first, I last, O output, C bop, U uop, void* buf
) {

  using T = typename std::iterator_traits<O>::value_type;

  unsigned count = std::distance(first, last);

  if(count == 0) {
    return;
  }

  // launch the scan loop
  detail::cuda_scan_loop(
    p, detail::cudaScanType::EXCLUSIVE,
    cuda_make_load_iterator<T>([=]__device__(auto i){ return uop(*(first+i)); }),
    count, output, bop, buf
  );
}


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