#include <vtkm/cont/Algorithm.h>
#include <vtkm/cont/DeviceAdapter.h>
#include <vtkm/cont/cuda/DeviceAdapterCuda.h>
#include <vtkm/cont/tbb/DeviceAdapterTBB.h>

#include <vtkm/exec/FunctorBase.h>

#include <vtkm/cont/testing/Testing.h>

namespace DeviceAdapterTagExamples
{

////
//// BEGIN-EXAMPLE DeviceTemplateArg.cxx
////
template<typename Device>
struct SetPortalFunctor : vtkm::exec::FunctorBase
{
  VTKM_IS_DEVICE_ADAPTER_TAG(Device);

  using ExecPortalType =
    typename vtkm::cont::ArrayHandle<vtkm::Id>::ExecutionTypes<Device>::Portal;
  ExecPortalType Portal;

  VTKM_CONT
  SetPortalFunctor(vtkm::cont::ArrayHandle<vtkm::Id> array, vtkm::Id size)
    : Portal(array.PrepareForOutput(size, Device()))
  {
  }

  //// PAUSE-EXAMPLE
  VTKM_CONT
  SetPortalFunctor(const ExecPortalType& portal)
    : Portal(portal)
  {
  }
  //// RESUME-EXAMPLE
  VTKM_EXEC
  void operator()(vtkm::Id index) const
  {
    //// PAUSE-EXAMPLE
    VTKM_ASSERT(index >= 0);
    VTKM_ASSERT(index < this->Portal.GetNumberOfValues());
    //// RESUME-EXAMPLE
    using ValueType = typename ExecPortalType::ValueType;
    this->Portal.Set(index, TestValue(index, ValueType()));
  }
};
////
//// END-EXAMPLE DeviceTemplateArg.cxx
////

template<typename ExecPortalType, typename ArrayHandleType, typename Device>
VTKM_CONT void TryUsingExecPortal(const ExecPortalType& execPortal,
                                  const ArrayHandleType& arrayHandle,
                                  Device)
{
  using ValueType = typename ArrayHandleType::ValueType;

  // In general, it is better to create a functor that takes an ExecObject as input
  // and the Schedule method in vtkm::cont::Algorithm can automatically call the
  // PrepareForExecution for the given device. However, the point of this is to test
  // a particular array portal for the execution environment, so we are going to
  // force the device.
  SetPortalFunctor<Device> functor(execPortal);

  vtkm::cont::Algorithm::Schedule(
    Device(), functor, arrayHandle.GetNumberOfValues());

  typename ArrayHandleType::PortalConstControl contPortal =
    arrayHandle.GetPortalConstControl();
  for (vtkm::Id index = 0; index < arrayHandle.GetNumberOfValues(); index++)
  {
    VTKM_TEST_ASSERT(contPortal.Get(index) == TestValue(index, ValueType()),
                     "Bad value set.");
  }
}

void UseDefaultDeviceAdapter()
{
  vtkm::cont::ArrayHandle<vtkm::Id> arrayHandle;

  SetPortalFunctor<vtkm::cont::DeviceAdapterTagSerial> functor(arrayHandle, 50);

  TryUsingExecPortal(
    functor.Portal, arrayHandle, vtkm::cont::DeviceAdapterTagSerial());
}

////
//// BEGIN-EXAMPLE DeviceAdapterTraits.cxx
////
template<typename ArrayHandleType, typename DeviceAdapterTag>
void CheckArrayHandleDevice(const ArrayHandleType& array, DeviceAdapterTag)
{
  VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
  VTKM_IS_DEVICE_ADAPTER_TAG(DeviceAdapterTag);

  vtkm::cont::DeviceAdapterId currentDevice = array.GetDeviceAdapterId();
  if (currentDevice == DeviceAdapterTag())
  {
    std::cout << "Array is already on device " << DeviceAdapterTag().GetName()
              << std::endl;
  }
  else
  {
    std::cout << "Copying array to device " << DeviceAdapterTag().GetName()
              << std::endl;
    array.PrepareForInput(DeviceAdapterTag());
  }
}
////
//// END-EXAMPLE DeviceAdapterTraits.cxx
////

////
//// BEGIN-EXAMPLE DeviceAdapterValid.cxx
////
namespace detail
{

template<typename ArrayHandleType, typename DeviceAdapterTag>
void SafeCheckArrayHandleDeviceImpl(const ArrayHandleType& array,
                                    DeviceAdapterTag,
                                    std::true_type)
{
  CheckArrayHandleDevice(array, DeviceAdapterTag());
}

template<typename ArrayHandleType, typename DeviceAdapterTag>
void SafeCheckArrayHandleDeviceImpl(const ArrayHandleType&,
                                    DeviceAdapterTag,
                                    std::false_type)
{
  std::cout << "Device " << DeviceAdapterTag().GetName() << " is not available"
            << std::endl;
}

} // namespace detail

template<typename ArrayHandleType, typename DeviceAdapterTag>
void SafeCheckArrayHandleDevice(const ArrayHandleType& array, DeviceAdapterTag)
{
  static const bool deviceValid = DeviceAdapterTag::IsEnabled;
  detail::SafeCheckArrayHandleDeviceImpl(
    array, DeviceAdapterTag(), std::integral_constant<bool, deviceValid>());
}
////
//// END-EXAMPLE DeviceAdapterValid.cxx
////

void TryCheckArrayHandleDevice()
{
  vtkm::cont::ArrayHandle<vtkm::Float32> array;
  array.Allocate(10);
  SetPortal(array.GetPortalControl());
  CheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagSerial());
  CheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagSerial());
  CheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagTBB());
  CheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagTBB());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagSerial());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagSerial());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagTBB());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagTBB());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagTBB());
  SafeCheckArrayHandleDevice(array, vtkm::cont::DeviceAdapterTagCuda());
}

void Test()
{
  UseDefaultDeviceAdapter();
  TryCheckArrayHandleDevice();
}

} // namespace DeviceAdapterTagExamples

int DeviceAdapterTag(int argc, char* argv[])
{
  return vtkm::cont::testing::Testing::Run(
    DeviceAdapterTagExamples::Test, argc, argv);
}
