#include <thrust/device_vector.h>
#include <thrust/extrema.h>
#include <thrust/functional.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/random.h>
#include <thrust/transform_reduce.h>

#include <cfloat>
#include <cmath>
#include <iomanip>
#include <iostream>

// This example computes the minimum and maximum values
// over a padded grid.  The padded values are not considered
// during the reduction operation.

// transform a tuple (int,value) into a tuple (bool,value,value)
// where the bool is true for valid grid values and false for
// values in the padded region of the grid
template <typename IndexType, typename ValueType>
struct transform_tuple
{
  using InputTuple  = typename cuda::std::tuple<IndexType, ValueType>;
  using OutputTuple = typename cuda::std::tuple<bool, ValueType, ValueType>;

  IndexType n, N;

  transform_tuple(IndexType n, IndexType N)
      : n(n)
      , N(N)
  {}

  __host__ __device__ OutputTuple operator()(const InputTuple& t) const
  {
    bool is_valid = (cuda::std::get<0>(t) % N) < n;
    return OutputTuple(is_valid, cuda::std::get<1>(t), cuda::std::get<1>(t));
  }
};

// reduce two tuples (bool,value,value) into a single tuple such that output
// contains the smallest and largest *valid* values.
template <typename IndexType, typename ValueType>
struct reduce_tuple
{
  using Tuple = typename cuda::std::tuple<bool, ValueType, ValueType>;

  __host__ __device__ Tuple operator()(const Tuple& t0, const Tuple& t1) const
  {
    if (cuda::std::get<0>(t0) && cuda::std::get<0>(t1)) // both valid
    {
      return Tuple(true,
                   thrust::min(cuda::std::get<1>(t0), cuda::std::get<1>(t1)),
                   thrust::max(cuda::std::get<2>(t0), cuda::std::get<2>(t1)));
    }
    else if (cuda::std::get<0>(t0))
    {
      return t0;
    }
    else
    {
      return t1; // if t0 is not valid, return t1 whether it is valid or not
    }
  }
};

int main()
{
  int M = 10; // number of rows
  int n = 11; // number of columns excluding padding
  int N = 16; // number of columns including padding

  thrust::default_random_engine rng(12345);
  thrust::uniform_real_distribution<float> dist(0.0f, 1.0f);

  thrust::device_vector<float> data(M * N, -1);

  // initialize valid values in grid
  for (int i = 0; i < M; i++)
  {
    for (int j = 0; j < n; j++)
    {
      data[static_cast<std::size_t>(i) * N + j] = dist(rng);
    }
  }

  // print full grid
  std::cout << "padded grid" << '\n';
  std::cout << std::fixed << std::setprecision(4);
  for (int i = 0; i < M; i++)
  {
    std::cout << " ";
    for (int j = 0; j < N; j++)
    {
      std::cout << data[(static_cast<std::size_t>(i) * N) + j] << " ";
    }
    std::cout << "\n";
  }
  std::cout << "\n";

  // compute min & max over valid region of the 2d grid
  using result_type = cuda::std::tuple<bool, float, float>;

  result_type init(true, FLT_MAX, -FLT_MAX); // initial value
  transform_tuple<int, float> unary_op(n, N); // transformation operator
  reduce_tuple<int, float> binary_op; // reduction operator

  result_type result = thrust::transform_reduce(
    thrust::make_zip_iterator(thrust::counting_iterator<int>(0), data.begin()),
    thrust::make_zip_iterator(cuda::std::tuple(thrust::counting_iterator<int>(0), data.begin()))
      + static_cast<std::ptrdiff_t>(data.size()),
    unary_op,
    init,
    binary_op);

  std::cout << "minimum value: " << cuda::std::get<1>(result) << '\n';
  std::cout << "maximum value: " << cuda::std::get<2>(result) << '\n';

  return 0;
}