#include <thrust/device_vector.h>
#include <thrust/functional.h>
#include <thrust/transform.h>

#include <iostream>

// This example demonstrates the use of placeholders to implement
// the SAXPY operation (i.e. Y[i] = a * X[i] + Y[i]).
//
// Placeholders enable developers to write concise inline expressions
// instead of full functors for many simple operations.  For example,
// the placeholder expression "_1 + _2" means to add the first argument,
// represented by _1, to the second argument, represented by _2.
// The names _1, _2, _3, _4 ... _10 represent the first ten arguments
// to the function.
//
// In this example, the placeholder expression "a * _1 + _2" is used
// to implement the SAXPY operation.  Note that the placeholder
// implementation is considerably shorter and written inline.

// allows us to use "_1" instead of "thrust::placeholders::_1"
using namespace thrust::placeholders;

// implementing SAXPY with a functor is cumbersome and verbose
struct saxpy_functor
{
  float a;

  saxpy_functor(float a)
      : a(a)
  {}

  __host__ __device__ float operator()(float x, float y)
  {
    return a * x + y;
  }
};

int main()
{
  // input data
  float a                             = 2.0f;
  thrust::device_vector<float> x_data = {1, 2, 3, 4};
  thrust::device_vector<float> y_data = {1, 1, 1, 1};

  // SAXPY implemented with a functor (function object)
  {
    thrust::device_vector<float> X = x_data;
    thrust::device_vector<float> Y = y_data;

    thrust::transform(
      X.begin(),
      X.end(), // input range #1
      Y.begin(), // input range #2
      Y.begin(), // output range
      saxpy_functor(a)); // functor

    std::cout << "SAXPY (functor method)" << '\n';
    for (size_t i = 0; i < Y.size(); i++)
    {
      std::cout << a << " * " << x_data[i] << " + " << y_data[i] << " = " << Y[i] << '\n';
    }
  }

  // SAXPY implemented with a placeholders
  {
    thrust::device_vector<float> X = x_data;
    thrust::device_vector<float> Y = y_data;

    thrust::transform(
      X.begin(),
      X.end(), // input range #1
      Y.begin(), // input range #2
      Y.begin(), // output range
      a * _1 + _2); // placeholder expression

    std::cout << "SAXPY (placeholder method)" << '\n';
    for (size_t i = 0; i < Y.size(); i++)
    {
      std::cout << a << " * " << x_data[i] << " + " << y_data[i] << " = " << Y[i] << '\n';
    }
  }

  return 0;
}