#include <vtkm/cont/ArrayHandleDecorator.h>

namespace
{

////
//// BEGIN-EXAMPLE DecoratorFunctor.cxx
////
template<typename PortalType1, typename PortalType2>
class InterlaceFunctor
{
  PortalType1 Portal1;
  PortalType2 Portal2;
  vtkm::Id Width;

public:
  InterlaceFunctor(const PortalType1& portal1,
                   const PortalType2& portal2,
                   vtkm::Id width)
    : Portal1(portal1)
    , Portal2(portal2)
    , Width(width)
  {
  }

  VTKM_EXEC_CONT typename PortalType1::ValueType operator()(vtkm::Id index) const
  {
    vtkm::Id interleaveGroup = index / (this->Width * 2);
    vtkm::Id interleaveIndex = index % (this->Width * 2);
    if (interleaveIndex < this->Width)
    {
      return this->Portal1.Get(interleaveIndex + (interleaveGroup * this->Width));
    }
    else
    {
      return this->Portal2.Get((interleaveIndex - this->Width) +
                               (interleaveGroup * this->Width));
    }
  }
};
////
//// END-EXAMPLE DecoratorFunctor.cxx
////

////
//// BEGIN-EXAMPLE DecoratorInverseFunctor.cxx
////
template<typename PortalType1, typename PortalType2>
class InterlaceInverseFunctor
{
  PortalType1 Portal1;
  PortalType2 Portal2;
  vtkm::Id Width;

public:
  InterlaceInverseFunctor(const PortalType1& portal1,
                          const PortalType2& portal2,
                          vtkm::Id width)
    : Portal1(portal1)
    , Portal2(portal2)
    , Width(width)
  {
  }

  template<typename T>
  VTKM_EXEC_CONT void operator()(vtkm::Id index, const T& value) const
  {
    vtkm::Id interleaveGroup = index / (this->Width * 2);
    vtkm::Id interleaveIndex = index % (this->Width * 2);
    if (interleaveIndex < this->Width)
    {
      this->Portal1.Set(interleaveIndex + (interleaveGroup * this->Width), value);
    }
    else
    {
      this->Portal2.Set(
        (interleaveIndex - this->Width) + (interleaveGroup * this->Width), value);
    }
  }
};
////
//// END-EXAMPLE DecoratorInverseFunctor.cxx
////

////
//// BEGIN-EXAMPLE DecoratorImplementation.cxx
////
struct InterlaceImplementation
{
  vtkm::Id Width;

  InterlaceImplementation(vtkm::Id width = 1)
    : Width(width)
  {
  }

  template<typename PortalType1, typename PortalType2>
  VTKM_CONT InterlaceFunctor<PortalType1, PortalType2> CreateFunctor(
    const PortalType1& portal1,
    const PortalType2 portal2) const
  {
    return InterlaceFunctor<PortalType1, PortalType2>(portal1, portal2, this->Width);
  }

  template<typename PortalType1, typename PortalType2>
  VTKM_CONT InterlaceInverseFunctor<PortalType1, PortalType2> CreateInverseFunctor(
    const PortalType1& portal1,
    const PortalType2 portal2) const
  {
    return InterlaceInverseFunctor<PortalType1, PortalType2>(
      portal1, portal2, this->Width);
  }

  template<typename ArrayType1, typename ArrayType2>
  VTKM_CONT void AllocateSourceArrays(vtkm::Id numValues,
                                      vtkm::CopyFlag preserve,
                                      vtkm::cont::Token& token,
                                      ArrayType1& array1,
                                      ArrayType2& array2) const
  {
    vtkm::Id numInterleaveGroups = (numValues / (this->Width * 2));
    vtkm::Id remainder = numValues - (numInterleaveGroups * this->Width * 2);
    if (remainder < this->Width)
    {
      array1.Allocate(
        (numInterleaveGroups * this->Width) + remainder, preserve, token);
      array2.Allocate((numInterleaveGroups * this->Width), preserve, token);
    }
    else
    {
      array1.Allocate((numInterleaveGroups + 1) * this->Width, preserve, token);
      array2.Allocate((numInterleaveGroups * this->Width) + remainder - this->Width,
                      preserve,
                      token);
    }
  }
};
////
//// END-EXAMPLE DecoratorImplementation.cxx
////

////
//// BEGIN-EXAMPLE DecoratorArrayHandle.cxx
////
template<typename ArrayType1, typename ArrayType2>
class ArrayHandleInterlace
  : public vtkm::cont::
      ArrayHandleDecorator<InterlaceImplementation, ArrayType1, ArrayType2>
{
  VTKM_IS_ARRAY_HANDLE(ArrayType1);
  VTKM_IS_ARRAY_HANDLE(ArrayType2);

public:
  VTKM_ARRAY_HANDLE_SUBCLASS(
    ArrayHandleInterlace,
    (ArrayHandleInterlace<ArrayType1, ArrayType2>),
    (vtkm::cont::
       ArrayHandleDecorator<InterlaceImplementation, ArrayType1, ArrayType2>));

  VTKM_CONT ArrayHandleInterlace(const ArrayType1& array1,
                                 const ArrayType2& array2,
                                 vtkm::Id width = 1)
    : Superclass(vtkm::cont::make_ArrayHandleDecorator(
        array1.GetNumberOfValues() + array2.GetNumberOfValues(),
        InterlaceImplementation(width),
        array1,
        array2))
  {
  }
};

template<typename ArrayType1, typename ArrayType2>
VTKM_CONT ArrayHandleInterlace<ArrayType1, ArrayType2> make_ArrayHandleInterlace(
  const ArrayType1& array1,
  const ArrayType2& array2,
  vtkm::Id width = 1)
{
  return ArrayHandleInterlace<ArrayType1, ArrayType2>(array1, array2, width);
}
////
//// END-EXAMPLE DecoratorArrayHandle.cxx
////

} // anonymous namespace

#include <vtkm/cont/ArrayCopyDevice.h>
#include <vtkm/cont/ArrayHandleIndex.h>
#include <vtkm/cont/DeviceAdapter.h>

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

namespace
{

void Test()
{
  constexpr vtkm::Id WIDTH = 3;
  constexpr vtkm::Id HALF_ARRAY_SIZE = 10 * WIDTH;
  constexpr vtkm::Id ARRAY_SIZE = 2 * HALF_ARRAY_SIZE;
  vtkm::cont::ArrayHandleIndex inputArray(ARRAY_SIZE);

  using BaseArrayType = vtkm::cont::ArrayHandle<vtkm::Id>;
  BaseArrayType array1;
  BaseArrayType array2;

  ArrayHandleInterlace<BaseArrayType, BaseArrayType> interlaceArray(
    array1, array2, WIDTH);

  vtkm::cont::ArrayCopyDevice(inputArray, interlaceArray);

  {
    VTKM_TEST_ASSERT(array1.GetNumberOfValues() == HALF_ARRAY_SIZE, "Wrong size.");
    VTKM_TEST_ASSERT(array2.GetNumberOfValues() == HALF_ARRAY_SIZE, "Wrong size.");
    auto portal1 = array1.ReadPortal();
    auto portal2 = array2.ReadPortal();
    for (vtkm::Id group = 0; group < (HALF_ARRAY_SIZE / WIDTH); ++group)
    {
      vtkm::Id groupStartValue1 = group * (WIDTH * 2);
      vtkm::Id groupStartValue2 = groupStartValue1 + WIDTH;
      for (vtkm::Id groupIndex = 0; groupIndex < WIDTH; ++groupIndex)
      {
        VTKM_TEST_ASSERT(portal1.Get((group * WIDTH) + groupIndex) ==
                         (groupStartValue1 + groupIndex));
        VTKM_TEST_ASSERT(portal2.Get((group * WIDTH) + groupIndex) ==
                         (groupStartValue2 + groupIndex));
      }
    }
  }

  auto interlacedIndices =
    make_ArrayHandleInterlace(vtkm::cont::ArrayHandleIndex(HALF_ARRAY_SIZE),
                              vtkm::cont::ArrayHandleIndex(HALF_ARRAY_SIZE));
  BaseArrayType targetArray;
  vtkm::cont::ArrayCopyDevice(interlacedIndices, targetArray);

  {
    VTKM_TEST_ASSERT(targetArray.GetNumberOfValues() == ARRAY_SIZE);
    auto portal = targetArray.ReadPortal();
    for (vtkm::Id index = 0; index < ARRAY_SIZE; ++index)
    {
      VTKM_TEST_ASSERT(portal.Get(index) == index / 2);
    }
  }

  // Check all PrepareFor* methods.
  vtkm::cont::Token token;
  interlaceArray.ReleaseResourcesExecution();
  interlaceArray.PrepareForInput(vtkm::cont::DeviceAdapterTagSerial{}, token);
  interlaceArray.PrepareForInPlace(vtkm::cont::DeviceAdapterTagSerial{}, token);
  interlaceArray.PrepareForOutput(
    ARRAY_SIZE + 1, vtkm::cont::DeviceAdapterTagSerial{}, token);
}

} // anonymous namespace

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