#include <thrust/execution_policy.h>
#include <thrust/transform.h>

#include <unittest/unittest.h>

#ifdef THRUST_TEST_DEVICE_SIDE
template <typename ExecutionPolicy, typename Iterator1, typename Iterator2, typename Function, typename Iterator3>
__global__ void transform_kernel(
  ExecutionPolicy exec, Iterator1 first, Iterator1 last, Iterator2 result1, Function f, Iterator3 result2)
{
  *result2 = thrust::transform(exec, first, last, result1, f);
}

template <typename ExecutionPolicy>
void TestTransformUnaryDevice(ExecutionPolicy exec)
{
  typedef thrust::device_vector<int> Vector;
  typedef typename Vector::value_type T;

  typename Vector::iterator iter;

  Vector input(3);
  Vector output(3);
  Vector result(3);
  input[0]  = 1;
  input[1]  = -2;
  input[2]  = 3;
  result[0] = -1;
  result[1] = 2;
  result[2] = -3;

  thrust::device_vector<typename Vector::iterator> iter_vec(1);

  transform_kernel<<<1, 1>>>(exec, input.begin(), input.end(), output.begin(), thrust::negate<T>(), iter_vec.begin());
  cudaError_t const err = cudaDeviceSynchronize();
  ASSERT_EQUAL(cudaSuccess, err);

  iter = iter_vec[0];

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());
  ASSERT_EQUAL(output, result);
}

void TestTransformUnaryDeviceSeq()
{
  TestTransformUnaryDevice(thrust::seq);
}
DECLARE_UNITTEST(TestTransformUnaryDeviceSeq);

void TestTransformUnaryDeviceDevice()
{
  TestTransformUnaryDevice(thrust::device);
}
DECLARE_UNITTEST(TestTransformUnaryDeviceDevice);

template <typename ExecutionPolicy,
          typename Iterator1,
          typename Iterator2,
          typename Function,
          typename Predicate,
          typename Iterator3>
__global__ void transform_if_kernel(
  ExecutionPolicy exec,
  Iterator1 first,
  Iterator1 last,
  Iterator2 result1,
  Function f,
  Predicate pred,
  Iterator3 result2)
{
  *result2 = thrust::transform_if(exec, first, last, result1, f, pred);
}

template <typename ExecutionPolicy>
void TestTransformIfUnaryNoStencilDevice(ExecutionPolicy exec)
{
  typedef thrust::device_vector<int> Vector;
  typedef typename Vector::value_type T;

  typename Vector::iterator iter;

  Vector input(3);
  Vector output(3);
  Vector result(3);

  input[0]  = 0;
  input[1]  = -2;
  input[2]  = 0;
  output[0] = -1;
  output[1] = -2;
  output[2] = -3;
  result[0] = -1;
  result[1] = 2;
  result[2] = -3;

  thrust::device_vector<typename Vector::iterator> iter_vec(1);

  transform_if_kernel<<<1, 1>>>(
    exec, input.begin(), input.end(), output.begin(), thrust::negate<T>(), thrust::identity<T>(), iter_vec.begin());
  cudaError_t const err = cudaDeviceSynchronize();
  ASSERT_EQUAL(cudaSuccess, err);

  iter = iter_vec[0];

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());
  ASSERT_EQUAL(output, result);
}

void TestTransformIfUnaryNoStencilDeviceSeq()
{
  TestTransformIfUnaryNoStencilDevice(thrust::seq);
}
DECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceSeq);

void TestTransformIfUnaryNoStencilDeviceDevice()
{
  TestTransformIfUnaryNoStencilDevice(thrust::device);
}
DECLARE_UNITTEST(TestTransformIfUnaryNoStencilDeviceDevice);

template <typename ExecutionPolicy,
          typename Iterator1,
          typename Iterator2,
          typename Iterator3,
          typename Function,
          typename Predicate,
          typename Iterator4>
__global__ void transform_if_kernel(
  ExecutionPolicy exec,
  Iterator1 first,
  Iterator1 last,
  Iterator2 stencil_first,
  Iterator3 result1,
  Function f,
  Predicate pred,
  Iterator4 result2)
{
  *result2 = thrust::transform_if(exec, first, last, stencil_first, result1, f, pred);
}

template <typename ExecutionPolicy>
void TestTransformIfUnaryDevice(ExecutionPolicy exec)
{
  typedef thrust::device_vector<int> Vector;
  typedef typename Vector::value_type T;

  typename Vector::iterator iter;

  Vector input(3);
  Vector stencil(3);
  Vector output(3);
  Vector result(3);

  input[0]   = 1;
  input[1]   = -2;
  input[2]   = 3;
  output[0]  = 1;
  output[1]  = 2;
  output[2]  = 3;
  stencil[0] = 1;
  stencil[1] = 0;
  stencil[2] = 1;
  result[0]  = -1;
  result[1]  = 2;
  result[2]  = -3;

  thrust::device_vector<typename Vector::iterator> iter_vec(1);

  transform_if_kernel<<<1, 1>>>(
    exec,
    input.begin(),
    input.end(),
    stencil.begin(),
    output.begin(),
    thrust::negate<T>(),
    thrust::identity<T>(),
    iter_vec.begin());
  cudaError_t const err = cudaDeviceSynchronize();
  ASSERT_EQUAL(cudaSuccess, err);

  iter = iter_vec[0];

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());
  ASSERT_EQUAL(output, result);
}

void TestTransformIfUnaryDeviceSeq()
{
  TestTransformIfUnaryDevice(thrust::seq);
}
DECLARE_UNITTEST(TestTransformIfUnaryDeviceSeq);

void TestTransformIfUnaryDeviceDevice()
{
  TestTransformIfUnaryDevice(thrust::device);
}
DECLARE_UNITTEST(TestTransformIfUnaryDeviceDevice);

template <typename ExecutionPolicy,
          typename Iterator1,
          typename Iterator2,
          typename Iterator3,
          typename Function,
          typename Iterator4>
__global__ void transform_kernel(
  ExecutionPolicy exec,
  Iterator1 first1,
  Iterator1 last1,
  Iterator2 first2,
  Iterator3 result1,
  Function f,
  Iterator4 result2)
{
  *result2 = thrust::transform(exec, first1, last1, first2, result1, f);
}

template <typename ExecutionPolicy>
void TestTransformBinaryDevice(ExecutionPolicy exec)
{
  typedef thrust::device_vector<int> Vector;
  typedef typename Vector::value_type T;

  typename Vector::iterator iter;

  Vector input1(3);
  Vector input2(3);
  Vector output(3);
  Vector result(3);
  input1[0] = 1;
  input1[1] = -2;
  input1[2] = 3;
  input2[0] = -4;
  input2[1] = 5;
  input2[2] = 6;
  result[0] = 5;
  result[1] = -7;
  result[2] = -3;

  thrust::device_vector<typename Vector::iterator> iter_vec(1);

  transform_kernel<<<1, 1>>>(
    exec, input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus<T>(), iter_vec.begin());
  cudaError_t const err = cudaDeviceSynchronize();
  ASSERT_EQUAL(cudaSuccess, err);

  iter = iter_vec[0];

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());
  ASSERT_EQUAL(output, result);
}

void TestTransformBinaryDeviceSeq()
{
  TestTransformBinaryDevice(thrust::seq);
}
DECLARE_UNITTEST(TestTransformBinaryDeviceSeq);

void TestTransformBinaryDeviceDevice()
{
  TestTransformBinaryDevice(thrust::device);
}
DECLARE_UNITTEST(TestTransformBinaryDeviceDevice);

template <typename ExecutionPolicy,
          typename Iterator1,
          typename Iterator2,
          typename Iterator3,
          typename Iterator4,
          typename Function,
          typename Predicate,
          typename Iterator5>
__global__ void transform_if_kernel(
  ExecutionPolicy exec,
  Iterator1 first1,
  Iterator1 last1,
  Iterator2 first2,
  Iterator3 stencil_first,
  Iterator4 result1,
  Function f,
  Predicate pred,
  Iterator5 result2)
{
  *result2 = thrust::transform_if(exec, first1, last1, first2, stencil_first, result1, f, pred);
}

template <typename ExecutionPolicy>
void TestTransformIfBinaryDevice(ExecutionPolicy exec)
{
  typedef thrust::device_vector<int> Vector;
  typedef typename Vector::value_type T;

  typename Vector::iterator iter;

  Vector input1(3);
  Vector input2(3);
  Vector stencil(3);
  Vector output(3);
  Vector result(3);

  input1[0]  = 1;
  input1[1]  = -2;
  input1[2]  = 3;
  input2[0]  = -4;
  input2[1]  = 5;
  input2[2]  = 6;
  stencil[0] = 0;
  stencil[1] = 1;
  stencil[2] = 0;
  output[0]  = 1;
  output[1]  = 2;
  output[2]  = 3;
  result[0]  = 5;
  result[1]  = 2;
  result[2]  = -3;

  thrust::identity<T> identity;

  thrust::device_vector<typename Vector::iterator> iter_vec(1);

  transform_if_kernel<<<1, 1>>>(
    exec,
    input1.begin(),
    input1.end(),
    input2.begin(),
    stencil.begin(),
    output.begin(),
    thrust::minus<T>(),
    thrust::not1(identity),
    iter_vec.begin());
  cudaError_t const err = cudaDeviceSynchronize();
  ASSERT_EQUAL(cudaSuccess, err);

  iter = iter_vec[0];

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());
  ASSERT_EQUAL(output, result);
}

void TestTransformIfBinaryDeviceSeq()
{
  TestTransformIfBinaryDevice(thrust::seq);
}
DECLARE_UNITTEST(TestTransformIfBinaryDeviceSeq);

void TestTransformIfBinaryDeviceDevice()
{
  TestTransformIfBinaryDevice(thrust::device);
}
DECLARE_UNITTEST(TestTransformIfBinaryDeviceDevice);
#endif

void TestTransformUnaryCudaStreams()
{
  typedef thrust::device_vector<int> Vector;
  typedef Vector::value_type T;

  Vector::iterator iter;

  Vector input(3);
  Vector output(3);
  Vector result(3);
  input[0]  = 1;
  input[1]  = -2;
  input[2]  = 3;
  result[0] = -1;
  result[1] = 2;
  result[2] = -3;

  cudaStream_t s;
  cudaStreamCreate(&s);

  iter = thrust::transform(thrust::cuda::par.on(s), input.begin(), input.end(), output.begin(), thrust::negate<T>());
  cudaStreamSynchronize(s);

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input.size());
  ASSERT_EQUAL(output, result);

  cudaStreamDestroy(s);
}
DECLARE_UNITTEST(TestTransformUnaryCudaStreams);

void TestTransformBinaryCudaStreams()
{
  typedef thrust::device_vector<int> Vector;
  typedef Vector::value_type T;

  Vector::iterator iter;

  Vector input1(3);
  Vector input2(3);
  Vector output(3);
  Vector result(3);
  input1[0] = 1;
  input1[1] = -2;
  input1[2] = 3;
  input2[0] = -4;
  input2[1] = 5;
  input2[2] = 6;
  result[0] = 5;
  result[1] = -7;
  result[2] = -3;

  cudaStream_t s;
  cudaStreamCreate(&s);

  iter = thrust::transform(
    thrust::cuda::par.on(s), input1.begin(), input1.end(), input2.begin(), output.begin(), thrust::minus<T>());
  cudaStreamSynchronize(s);

  ASSERT_EQUAL(std::size_t(iter - output.begin()), input1.size());
  ASSERT_EQUAL(output, result);

  cudaStreamDestroy(s);
}
DECLARE_UNITTEST(TestTransformBinaryCudaStreams);
