// SPDX-FileCopyrightText: Copyright (c) 2016, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause

#pragma once

#include <thrust/detail/config.h>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <thrust/detail/type_deduction.h>
#include <thrust/system/cpp/detail/execution_policy.h>
#include <thrust/system/cuda/detail/execution_policy.h>

THRUST_NAMESPACE_BEGIN
namespace cuda_cub
{
template <class Sys1, class Sys2>
struct cross_system : execution_policy<cross_system<Sys1, Sys2>>
{
  using policy1 = thrust::execution_policy<Sys1>;
  using policy2 = thrust::execution_policy<Sys2>;

  policy1& sys1;
  policy2& sys2;

  inline _CCCL_HOST_DEVICE cross_system(policy1& sys1, policy2& sys2)
      : sys1(sys1)
      , sys2(sys2)
  {}

  inline _CCCL_HOST_DEVICE cross_system<Sys2, Sys1> rotate() const
  {
    return cross_system<Sys2, Sys1>(sys2, sys1);
  }
};

// Device to host.
template <class Sys1, class Sys2>
constexpr _CCCL_HOST_DEVICE auto
direction_of_copy(thrust::system::cuda::execution_policy<Sys1> const&, thrust::cpp::execution_policy<Sys2> const&)
  THRUST_DECLTYPE_RETURNS(thrust::detail::integral_constant<cudaMemcpyKind, cudaMemcpyDeviceToHost>{})

  // Host to device.
  template <class Sys1, class Sys2>
  constexpr _CCCL_HOST_DEVICE auto direction_of_copy(thrust::cpp::execution_policy<Sys1> const&,
                                                     thrust::system::cuda::execution_policy<Sys2> const&)
    THRUST_DECLTYPE_RETURNS(thrust::detail::integral_constant<cudaMemcpyKind, cudaMemcpyHostToDevice>{})

  // Device to device.
  template <class Sys1, class Sys2>
  constexpr _CCCL_HOST_DEVICE auto direction_of_copy(thrust::system::cuda::execution_policy<Sys1> const&,
                                                     thrust::system::cuda::execution_policy<Sys2> const&)
    THRUST_DECLTYPE_RETURNS(thrust::detail::integral_constant<cudaMemcpyKind, cudaMemcpyDeviceToDevice>{})

  // Device to device.
  template <class DerivedPolicy>
  constexpr _CCCL_HOST_DEVICE auto direction_of_copy(execution_policy<DerivedPolicy> const&)
    THRUST_DECLTYPE_RETURNS(thrust::detail::integral_constant<cudaMemcpyKind, cudaMemcpyDeviceToDevice>{})

      template <class Sys1, class Sys2>
      constexpr
  _CCCL_HOST_DEVICE auto direction_of_copy(execution_policy<cross_system<Sys1, Sys2>> const& systems)
    THRUST_DECLTYPE_RETURNS(
      direction_of_copy(derived_cast(derived_cast(systems).sys1), derived_cast(derived_cast(systems).sys2)))

      template <typename ExecutionPolicy0,
                typename ExecutionPolicy1,
                // MSVC2015 WAR: put decltype here instead of in trailing return type
                typename Direction =
                  decltype(direction_of_copy(std::declval<ExecutionPolicy0>(), std::declval<ExecutionPolicy1>()))>
      constexpr _CCCL_HOST_DEVICE thrust::detail::
        integral_constant<bool, cudaMemcpyDeviceToHost == Direction::value> is_device_to_host_copy(
          ExecutionPolicy0 const& exec0, ExecutionPolicy1 const& exec1) noexcept
{
  return {};
}

template <typename ExecutionPolicy,
          // MSVC2015 WAR: put decltype here instead of in trailing return type
          typename Direction = decltype(direction_of_copy(std::declval<ExecutionPolicy>()))>
constexpr _CCCL_HOST_DEVICE thrust::detail::integral_constant<bool, cudaMemcpyDeviceToHost == Direction::value>
is_device_to_host_copy(ExecutionPolicy const& exec) noexcept
{
  return {};
}

template <
  typename ExecutionPolicy0,
  typename ExecutionPolicy1,
  // MSVC2015 WAR: put decltype here instead of in trailing return type
  typename Direction = decltype(direction_of_copy(std::declval<ExecutionPolicy0>(), std::declval<ExecutionPolicy1>()))>
constexpr _CCCL_HOST_DEVICE thrust::detail::integral_constant<bool, cudaMemcpyHostToDevice == Direction::value>
is_host_to_device_copy(ExecutionPolicy0 const& exec0, ExecutionPolicy1 const& exec1) noexcept
{
  return {};
}

template <typename ExecutionPolicy,
          // MSVC2015 WAR: put decltype here instead of in trailing return type
          typename Direction = decltype(direction_of_copy(std::declval<ExecutionPolicy>()))>
constexpr _CCCL_HOST_DEVICE thrust::detail::integral_constant<bool, cudaMemcpyHostToDevice == Direction::value>
is_host_to_device_copy(ExecutionPolicy const& exec) noexcept
{
  return {};
}

template <
  typename ExecutionPolicy0,
  typename ExecutionPolicy1,
  // MSVC2015 WAR: put decltype here instead of in trailing return type
  typename Direction = decltype(direction_of_copy(std::declval<ExecutionPolicy0>(), std::declval<ExecutionPolicy1>()))>
constexpr _CCCL_HOST_DEVICE thrust::detail::integral_constant<bool, cudaMemcpyDeviceToDevice == Direction::value>
is_device_to_device_copy(ExecutionPolicy0 const& exec0, ExecutionPolicy1 const& exec1) noexcept
{
  return {};
}

template <typename ExecutionPolicy,
          // MSVC2015 WAR: put decltype here instead of in trailing return type
          typename Direction = decltype(direction_of_copy(std::declval<ExecutionPolicy>()))>
constexpr _CCCL_HOST_DEVICE thrust::detail::integral_constant<bool, cudaMemcpyDeviceToDevice == Direction::value>
is_device_to_device_copy(ExecutionPolicy const& exec) noexcept
{
  return {};
}

/////////////////////////////////////////////////////////////////////////////

// Device to host.
template <class Sys1, class Sys2>
_CCCL_HOST_DEVICE auto select_device_system(thrust::cuda::execution_policy<Sys1>& sys1, thrust::execution_policy<Sys2>&)
  THRUST_DECLTYPE_RETURNS(sys1)

  // Device to host.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_device_system(thrust::cuda::execution_policy<Sys1> const& sys1,
                                              thrust::execution_policy<Sys2> const&) THRUST_DECLTYPE_RETURNS(sys1)

  // Host to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_device_system(thrust::execution_policy<Sys1>&,
                                              thrust::cuda::execution_policy<Sys2>& sys2) THRUST_DECLTYPE_RETURNS(sys2)

  // Host to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_device_system(thrust::execution_policy<Sys1> const&,
                                              thrust::cuda::execution_policy<Sys2> const& sys2)
    THRUST_DECLTYPE_RETURNS(sys2)

  // Device to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_device_system(thrust::cuda::execution_policy<Sys1>& sys1,
                                              thrust::cuda::execution_policy<Sys2>&) THRUST_DECLTYPE_RETURNS(sys1)

  // Device to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_device_system(thrust::cuda::execution_policy<Sys1> const& sys1,
                                              thrust::cuda::execution_policy<Sys2> const&) THRUST_DECLTYPE_RETURNS(sys1)

  /////////////////////////////////////////////////////////////////////////////

  // Device to host.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::cuda::execution_policy<Sys1>&, thrust::execution_policy<Sys2>& sys2)
    THRUST_DECLTYPE_RETURNS(sys2)

  // Device to host.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::cuda::execution_policy<Sys1> const&,
                                            thrust::execution_policy<Sys2> const& sys2) THRUST_DECLTYPE_RETURNS(sys2)

  // Host to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::execution_policy<Sys1>& sys1, thrust::cuda::execution_policy<Sys2>&)
    THRUST_DECLTYPE_RETURNS(sys1)

  // Host to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::execution_policy<Sys1> const& sys1,
                                            thrust::cuda::execution_policy<Sys2> const&) THRUST_DECLTYPE_RETURNS(sys1)

  // Device to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::execution_policy<Sys1>& sys1, thrust::execution_policy<Sys2>&)
    THRUST_DECLTYPE_RETURNS(sys1)

  // Device to device.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE auto select_host_system(thrust::execution_policy<Sys1> const& sys1,
                                            thrust::execution_policy<Sys2> const&) THRUST_DECLTYPE_RETURNS(sys1)

  // Device to host.
  template <class Sys1, class Sys2>
  _CCCL_HOST_DEVICE cross_system<Sys1, Sys2> select_system(
    execution_policy<Sys1> const& sys1, thrust::cpp::execution_policy<Sys2> const& sys2)
{
  thrust::execution_policy<Sys1>& non_const_sys1      = const_cast<execution_policy<Sys1>&>(sys1);
  thrust::cpp::execution_policy<Sys2>& non_const_sys2 = const_cast<thrust::cpp::execution_policy<Sys2>&>(sys2);
  return cross_system<Sys1, Sys2>(non_const_sys1, non_const_sys2);
}

// Host to device.
template <class Sys1, class Sys2>
_CCCL_HOST_DEVICE cross_system<Sys1, Sys2>
select_system(thrust::cpp::execution_policy<Sys1> const& sys1, execution_policy<Sys2> const& sys2)
{
  thrust::cpp::execution_policy<Sys1>& non_const_sys1 = const_cast<thrust::cpp::execution_policy<Sys1>&>(sys1);
  thrust::execution_policy<Sys2>& non_const_sys2      = const_cast<execution_policy<Sys2>&>(sys2);
  return cross_system<Sys1, Sys2>(non_const_sys1, non_const_sys2);
}
} // namespace cuda_cub
THRUST_NAMESPACE_END
