#include <vtkm/cont/ArrayHandle.h>
#include <vtkm/internal/ArrayPortalHelpers.h>

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

public:
  using ValueType = typename PortalType1::ValueType;

  VTKM_EXEC_CONT ArrayPortalInterlace()
    : Portal1()
    , Portal2()
    , Width(1)
  {
  }

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

  /// Copy constructor for any other ArrayPortalInterlace with a portal type
  /// that can be copied to this portal type. This allows us to do any type
  /// casting that the portals do (like the non-const to const cast).
  template<typename OtherP1, typename OtherP2>
  VTKM_EXEC_CONT ArrayPortalInterlace(
    const ArrayPortalInterlace<OtherP1, OtherP2>& src)
    : Portal1(src.GetPortal1())
    , Portal2(src.GetPortal2())
    , Width(src.GetWidth())
  {
  }

  VTKM_EXEC_CONT
  vtkm::Id GetNumberOfValues() const
  {
    return this->Portal1.GetNumberOfValues() + this->Portal2.GetNumberOfValues();
  }

  VTKM_EXEC_CONT
  ValueType Get(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));
    }
  }

  // The template is a trick to use SFINAE semantics to only define this Set
  // method if both Portal1 and Portal2 define a Set method.
  template<
    typename P1 = PortalType1,
    typename P2 = PortalType2,
    typename =
      typename std::enable_if<vtkm::internal::PortalSupportsSets<P1>::value>::type,
    typename =
      typename std::enable_if<vtkm::internal::PortalSupportsSets<P2>::value>::type>
  VTKM_EXEC_CONT void Set(vtkm::Id index, const ValueType& 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);
    }
  }

  VTKM_EXEC_CONT const PortalType1& GetPortal1() const { return this->Portal1; }
  VTKM_EXEC_CONT const PortalType2& GetPortal2() const { return this->Portal2; }

  VTKM_EXEC_CONT vtkm::Id GetWidth() const { return this->Width; }
};
////
//// END-EXAMPLE DerivedArrayPortal.cxx
////

////
//// BEGIN-EXAMPLE DerivedArrayStorage.cxx
////
template<typename StorageType1, typename StorageType2>
struct StorageTagInterlace
{
};

namespace vtkm
{
namespace cont
{
namespace internal
{

namespace detail
{

// We need to store some metadata to implement interlaced array storage.
struct InterlaceMetaData
{
  vtkm::Id Width;

  InterlaceMetaData(vtkm::Id width = 1)
    : Width(width)
  {
  }
};

} // namespace detail

template<typename T, typename StorageTag1, typename StorageTag2>
class Storage<T, StorageTagInterlace<StorageTag1, StorageTag2>>
{
  // We will be deriving the behavior of two other arrays, so we will be
  // using the Storage objects for those arrays to implement ours.
  using SourceStorage1 = vtkm::cont::internal::Storage<T, StorageTag1>;
  using SourceStorage2 = vtkm::cont::internal::Storage<T, StorageTag2>;

  // All storage objects store the actual data in an array of Buffer objects.
  // We will use the first object to store the metadata.
  // The next objects will be for the first array, the remaining for the
  // second array. These functions make it convenient to access these Buffers.
  template<typename Buff>
  VTKM_CONT static Buff* Buffers1(Buff* buffers)
  {
    return buffers + 1;
  }
  template<typename Buff>
  VTKM_CONT static Buff* Buffers2(Buff* buffers)
  {
    return buffers + 1 + SourceStorage1::GetNumberOfBuffers();
  }

public:
  using ReadPortalType =
    ArrayPortalInterlace<typename SourceStorage1::ReadPortalType,
                         typename SourceStorage2::ReadPortalType>;
  using WritePortalType =
    ArrayPortalInterlace<typename SourceStorage1::WritePortalType,
                         typename SourceStorage2::WritePortalType>;

  // Not necessary for Storage, but useful.
  VTKM_CONT static vtkm::Id GetWidth(const vtkm::cont::internal::Buffer* buffers)
  {
    return buffers[0].GetMetaData<detail::InterlaceMetaData>().Width;
  }

  VTKM_CONT static constexpr vtkm::IdComponent GetNumberOfBuffers()
  {
    return 1 + SourceStorage1::GetNumberOfBuffers() +
      SourceStorage2::GetNumberOfBuffers();
  }

  VTKM_CONT static vtkm::Id GetNumberOfValues(
    const vtkm::cont::internal::Buffer* buffers)
  {
    return (SourceStorage1::GetNumberOfValues(Buffers1(buffers)) +
            SourceStorage2::GetNumberOfValues(Buffers2(buffers)));
  }

  VTKM_CONT static void ResizeBuffers(vtkm::Id numValues,
                                      vtkm::cont::internal::Buffer* buffers,
                                      vtkm::CopyFlag preserve,
                                      vtkm::cont::Token& token)
  {
    vtkm::Id width = GetWidth(buffers);
    vtkm::Id numInterleaveGroups = (numValues / (width * 2));
    vtkm::Id remainder = numValues - (numInterleaveGroups * width * 2);
    if (remainder < width)
    {
      SourceStorage1::ResizeBuffers((numInterleaveGroups * width) + remainder,
                                    Buffers1(buffers),
                                    preserve,
                                    token);
      SourceStorage2::ResizeBuffers(
        (numInterleaveGroups * width), Buffers2(buffers), preserve, token);
    }
    else
    {
      SourceStorage1::ResizeBuffers(
        (numInterleaveGroups + 1) * width, Buffers1(buffers), preserve, token);
      SourceStorage2::ResizeBuffers((numInterleaveGroups * width) + remainder -
                                      width,
                                    Buffers2(buffers),
                                    preserve,
                                    token);
    }
  }

  VTKM_CONT static ReadPortalType CreateReadPortal(
    const vtkm::cont::internal::Buffer* buffers,
    vtkm::cont::DeviceAdapterId device,
    vtkm::cont::Token& token)
  {
    return ReadPortalType(
      SourceStorage1::CreateReadPortal(Buffers1(buffers), device, token),
      SourceStorage2::CreateReadPortal(Buffers2(buffers), device, token),
      GetWidth(buffers));
  }

  VTKM_CONT static WritePortalType CreateWritePortal(
    vtkm::cont::internal::Buffer* buffers,
    vtkm::cont::DeviceAdapterId device,
    vtkm::cont::Token& token)
  {
    return WritePortalType(
      SourceStorage1::CreateWritePortal(Buffers1(buffers), device, token),
      SourceStorage2::CreateWritePortal(Buffers2(buffers), device, token),
      GetWidth(buffers));
  }

  VTKM_CONT static auto CreateBuffers(
    const vtkm::cont::ArrayHandle<T, StorageTag1>& array1,
    const vtkm::cont::ArrayHandle<T, StorageTag2>& array2,
    vtkm::Id width) -> decltype(vtkm::cont::internal::CreateBuffers())
  {
    return vtkm::cont::internal::CreateBuffers(
      detail::InterlaceMetaData(width), array1, array2);
  }
};

} // namespace internal
} // namespace cont
} // namespace vtkm
////
//// END-EXAMPLE DerivedArrayStorage.cxx
////

////
//// BEGIN-EXAMPLE DerivedArrayHandle.cxx
////
template<typename ArrayHandleType1, typename ArrayHandleType2>
class ArrayHandleInterlace
  : public vtkm::cont::ArrayHandle<
      typename ArrayHandleType1::ValueType,
      StorageTagInterlace<typename ArrayHandleType1::StorageTag,
                          typename ArrayHandleType2::StorageTag>>
{
  VTKM_IS_ARRAY_HANDLE(ArrayHandleType1);
  VTKM_IS_ARRAY_HANDLE(ArrayHandleType2);

public:
  VTKM_ARRAY_HANDLE_SUBCLASS(
    ArrayHandleInterlace,
    (ArrayHandleInterlace<ArrayHandleType1, ArrayHandleType2>),
    (vtkm::cont::ArrayHandle<
      typename ArrayHandleType1::ValueType,
      StorageTagInterlace<typename ArrayHandleType1::StorageTag,
                          typename ArrayHandleType2::StorageTag>>));

private:
  using StorageType = vtkm::cont::internal::Storage<ValueType, StorageTag>;

public:
  VTKM_CONT
  ArrayHandleInterlace(const ArrayHandleType1& array1,
                       const ArrayHandleType2& array2,
                       vtkm::Id width = 1)
    : Superclass(StorageType::CreateBuffers(array1, array2, width))
  {
  }
};
////
//// END-EXAMPLE DerivedArrayHandle.cxx
////

////
//// BEGIN-EXAMPLE DerivedMakeArrayHandle.cxx
////
template<typename ArrayHandle1, typename ArrayHandle2>
VTKM_CONT ArrayHandleInterlace<ArrayHandle1, ArrayHandle2> make_ArrayHandleInterlace(
  const ArrayHandle1& array1,
  const ArrayHandle2& array2,
  vtkm::Id width = 1)
{
  VTKM_IS_ARRAY_HANDLE(ArrayHandle1);
  VTKM_IS_ARRAY_HANDLE(ArrayHandle2);

  return ArrayHandleInterlace<ArrayHandle1, ArrayHandle2>(array1, array2, width);
}
////
//// END-EXAMPLE DerivedMakeArrayHandle.cxx
////

#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 ArrayHandleDerived(int argc, char* argv[])
{
  return vtkm::cont::testing::Testing::Run(Test, argc, argv);
}
