#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 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;
    vtkm::cont::Token token;
    array.PrepareForInput(DeviceAdapterTag{}, token);
  }
}
////
//// 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.WritePortal());
  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()
{
  TryCheckArrayHandleDevice();
}

} // namespace DeviceAdapterTagExamples

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