// Copyright (c) 2020-2025 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

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

#include <cmath>
#include <iomanip>

#include "include/host_device.h"
#include <float.h>

// 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 thrust::tuple<IndexType, ValueType>;
  using OutputTuple = typename thrust::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 = (thrust::get<0>(t) % N) < n;
    return OutputTuple(is_valid, thrust::get<1>(t), thrust::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 thrust::tuple<bool, ValueType, ValueType>;

  __host__ __device__ Tuple operator()(const Tuple& t0, const Tuple& t1) const
  {
    if (thrust::get<0>(t0) && thrust::get<0>(t1)) // both valid
    {
      return Tuple(
        true, thrust::min(thrust::get<1>(t0), thrust::get<1>(t1)), thrust::max(thrust::get<2>(t0), thrust::get<2>(t1)));
    }
    else if (thrust::get<0>(t0))
    {
      return t0;
    }
    else if (thrust::get<0>(t1))
    {
      return t1;
    }
    else
    {
      return t1; // if neither is valid then it doesn't matter what we return
    }
  }
};

int main(void)
{
  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[i * N + j] = dist(rng);
    }
  }

  // print full grid
  std::cout << "padded grid" << std::endl;
  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[i * N + j] << " ";
    }
    std::cout << "\n";
  }
  std::cout << "\n";

  // compute min & max over valid region of the 2d grid
  using result_type = thrust::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::make_tuple(thrust::counting_iterator<int>(0), data.begin())),
    thrust::make_zip_iterator(thrust::make_tuple(thrust::counting_iterator<int>(0), data.begin())) + data.size(),
    unary_op,
    init,
    binary_op);

  std::cout << "minimum value: " << thrust::get<1>(result) << std::endl;
  std::cout << "maximum value: " << thrust::get<2>(result) << std::endl;

  return 0;
}