//============================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//
//  Copyright 2014 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
//  Copyright 2014 UT-Battelle, LLC.
//  Copyright 2014 Los Alamos National Security.
//
//  Under the terms of Contract DE-NA0003525 with NTESS,
//  the U.S. Government retains certain rights in this software.
//
//  Under the terms of Contract DE-AC52-06NA25396 with Los Alamos National
//  Laboratory (LANL), the U.S. Government retains certain rights in
//  this software.
//============================================================================
#ifndef vtk_m_cont_DynamicArrayHandle_h
#define vtk_m_cont_DynamicArrayHandle_h

#include <vtkm/cont/vtkm_cont_export.h>

#include <vtkm/TypeListTag.h>
#include <vtkm/VecTraits.h>

#include <vtkm/cont/ArrayHandle.h>
#include <vtkm/cont/ErrorBadType.h>
#include <vtkm/cont/StorageListTag.h>

#include <vtkm/cont/internal/DynamicTransform.h>

#include <sstream>

namespace vtkm
{
namespace cont
{

// Forward declaration
template <typename TypeList, typename StorageList>
class DynamicArrayHandleBase;

namespace detail
{

/// \brief Base class for PolymorphicArrayHandleContainer
///
struct VTKM_CONT_EXPORT PolymorphicArrayHandleContainerBase
{
  PolymorphicArrayHandleContainerBase();

  // This must exist so that subclasses are destroyed correctly.
  virtual ~PolymorphicArrayHandleContainerBase();

  virtual vtkm::IdComponent GetNumberOfComponents() const = 0;
  virtual vtkm::Id GetNumberOfValues() const = 0;

  virtual void PrintSummary(std::ostream& out) const = 0;

  virtual std::shared_ptr<PolymorphicArrayHandleContainerBase> NewInstance() const = 0;
};

/// \brief ArrayHandle container that can use C++ run-time type information.
///
/// The \c PolymorphicArrayHandleContainer is similar to the
/// \c SimplePolymorphicContainer in that it can contain an object of an
/// unkown type. However, this class specifically holds ArrayHandle objects
/// (with different template parameters) so that it can polymorphically answer
/// simple questions about the object.
///
template <typename T, typename Storage>
struct VTKM_ALWAYS_EXPORT PolymorphicArrayHandleContainer
  : public PolymorphicArrayHandleContainerBase
{
  using ArrayHandleType = vtkm::cont::ArrayHandle<T, Storage>;

  ArrayHandleType Array;

  VTKM_CONT
  PolymorphicArrayHandleContainer()
    : Array()
  {
  }

  VTKM_CONT
  PolymorphicArrayHandleContainer(const ArrayHandleType& array)
    : Array(array)
  {
  }

  virtual vtkm::IdComponent GetNumberOfComponents() const
  {
    return vtkm::VecTraits<T>::NUM_COMPONENTS;
  }

  virtual vtkm::Id GetNumberOfValues() const { return this->Array.GetNumberOfValues(); }

  virtual void PrintSummary(std::ostream& out) const
  {
    vtkm::cont::printSummary_ArrayHandle(this->Array, out);
  }

  virtual std::shared_ptr<PolymorphicArrayHandleContainerBase> NewInstance() const
  {
    return std::shared_ptr<PolymorphicArrayHandleContainerBase>(
      new PolymorphicArrayHandleContainer<T, Storage>());
  }
};

// One instance of a template class cannot access the private members of
// another instance of a template class. However, I want to be able to copy
// construct a DynamicArrayHandle from another DynamicArrayHandle of any other
// type. Since you cannot partially specialize friendship, use this accessor
// class to get at the internals for the copy constructor.
struct DynamicArrayHandleCopyHelper
{
  template <typename TypeList, typename StorageList>
  VTKM_CONT static const std::shared_ptr<vtkm::cont::detail::PolymorphicArrayHandleContainerBase>&
  GetArrayHandleContainer(const vtkm::cont::DynamicArrayHandleBase<TypeList, StorageList>& src)
  {
    return src.ArrayContainer;
  }
};

// A simple function to downcast an ArrayHandle encapsulated in a
// PolymorphicArrayHandleContainerBase to the given type of ArrayHandle. If the
// conversion cannot be done, nullptr is returned.
template <typename Type, typename Storage>
VTKM_CONT vtkm::cont::ArrayHandle<Type, Storage>* DynamicArrayHandleTryCast(
  vtkm::cont::detail::PolymorphicArrayHandleContainerBase* arrayContainer)
{
  vtkm::cont::detail::PolymorphicArrayHandleContainer<Type, Storage>* downcastContainer =
    dynamic_cast<vtkm::cont::detail::PolymorphicArrayHandleContainer<Type, Storage>*>(
      arrayContainer);
  if (downcastContainer != nullptr)
  {
    return &downcastContainer->Array;
  }
  else
  {
    return nullptr;
  }
}

template <typename Type, typename Storage>
VTKM_CONT vtkm::cont::ArrayHandle<Type, Storage>* DynamicArrayHandleTryCast(
  const std::shared_ptr<vtkm::cont::detail::PolymorphicArrayHandleContainerBase>& arrayContainer)
{
  return detail::DynamicArrayHandleTryCast<Type, Storage>(arrayContainer.get());
}

} // namespace detail

/// \brief Holds an array handle without having to specify template parameters.
///
/// \c DynamicArrayHandle holds an \c ArrayHandle object using runtime
/// polymorphism to manage different value types and storage rather than
/// compile-time templates. This adds a programming convenience that helps
/// avoid a proliferation of templates. It also provides the management
/// necessary to interface VTK-m with data sources where types will not be
/// known until runtime.
///
/// To interface between the runtime polymorphism and the templated algorithms
/// in VTK-m, \c DynamicArrayHandle contains a method named \c CastAndCall that
/// will determine the correct type from some known list of types and storage
/// objects. This mechanism is used internally by VTK-m's worklet invocation
/// mechanism to determine the type when running algorithms.
///
/// By default, \c DynamicArrayHandle will assume that the value type in the
/// array matches one of the types specified by \c VTKM_DEFAULT_TYPE_LIST_TAG
/// and the storage matches one of the tags specified by \c
/// VTKM_DEFAULT_STORAGE_LIST_TAG. These lists can be changed by using the \c
/// ResetTypeList and \c ResetStorageList methods, respectively. It is
/// worthwhile to match these lists closely to the possible types that might be
/// used. If a type is missing you will get a runtime error. If there are more
/// types than necessary, then the template mechanism will create a lot of
/// object code that is never used, and keep in mind that the number of
/// combinations grows exponentially when using multiple \c DynamicArrayHandle
/// objects.
///
/// The actual implementation of \c DynamicArrayHandle is in a templated class
/// named \c DynamicArrayHandleBase, which is templated on the list of
/// component types and storage types. \c DynamicArrayHandle is really just a
/// typedef of \c DynamicArrayHandleBase with the default type and storage
/// lists.
///
template <typename TypeList, typename StorageList>
class VTKM_ALWAYS_EXPORT DynamicArrayHandleBase
{
public:
  VTKM_CONT
  DynamicArrayHandleBase() {}

  template <typename Type, typename Storage>
  VTKM_CONT DynamicArrayHandleBase(const vtkm::cont::ArrayHandle<Type, Storage>& array)
    : ArrayContainer(new vtkm::cont::detail::PolymorphicArrayHandleContainer<Type, Storage>(array))
  {
  }

  VTKM_CONT
  DynamicArrayHandleBase(const DynamicArrayHandleBase<TypeList, StorageList>& src)
    : ArrayContainer(src.ArrayContainer)
  {
  }

  template <typename OtherTypeList, typename OtherStorageList>
  VTKM_CONT explicit DynamicArrayHandleBase(
    const DynamicArrayHandleBase<OtherTypeList, OtherStorageList>& src)
    : ArrayContainer(detail::DynamicArrayHandleCopyHelper::GetArrayHandleContainer(src))
  {
  }

  VTKM_CONT
  ~DynamicArrayHandleBase() {}

  VTKM_CONT
  vtkm::cont::DynamicArrayHandleBase<TypeList, StorageList>& operator=(
    const vtkm::cont::DynamicArrayHandleBase<TypeList, StorageList>& src)
  {
    this->ArrayContainer = src.ArrayContainer;
    return *this;
  }

  /// Returns true if this array is of the provided type and uses the provided
  /// storage.
  ///
  template <typename Type, typename Storage>
  VTKM_CONT bool IsTypeAndStorage() const
  {
    return (detail::DynamicArrayHandleTryCast<Type, Storage>(this->ArrayContainer) != nullptr);
  }

  /// Returns true if this array matches the array handle type passed in.
  ///
  template <typename ArrayHandleType>
  VTKM_CONT bool IsType()
  {
    VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
    using ValueType = typename ArrayHandleType::ValueType;
    using StorageTag = typename ArrayHandleType::StorageTag;
    return this->IsTypeAndStorage<ValueType, StorageTag>();
  }

  /// Returns true if the array held in this object is the same (or equivalent)
  /// type as the object given.
  ///
  template <typename ArrayHandleType>
  VTKM_CONT bool IsSameType(const ArrayHandleType&)
  {
    VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
    return this->IsType<ArrayHandleType>();
  }

  /// Returns this array cast to an ArrayHandle object of the given type and
  /// storage. Throws \c ErrorBadType if the cast does not work. Use
  /// \c IsTypeAndStorage to check if the cast can happen.
  ///
  ///
  template <typename Type, typename Storage>
  VTKM_CONT vtkm::cont::ArrayHandle<Type, Storage> CastToTypeStorage() const
  {
    vtkm::cont::ArrayHandle<Type, Storage>* downcastArray =
      detail::DynamicArrayHandleTryCast<Type, Storage>(this->ArrayContainer);
    if (downcastArray == nullptr)
    {
      throw vtkm::cont::ErrorBadType("Bad cast of dynamic array.");
    }
    // Technically, this method returns a copy of the \c ArrayHandle. But
    // because \c ArrayHandle acts like a shared pointer, it is valid to
    // do the copy.
    return *downcastArray;
  }

  /// Returns this array cast to the given \c ArrayHandle type. Throws \c
  /// ErrorBadType if the cast does not work. Use \c IsType
  /// to check if the cast can happen.
  ///
  template <typename ArrayHandleType>
  VTKM_CONT ArrayHandleType Cast() const
  {
    VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
    using ValueType = typename ArrayHandleType::ValueType;
    using StorageTag = typename ArrayHandleType::StorageTag;
    // Technically, this method returns a copy of the \c ArrayHandle. But
    // because \c ArrayHandle acts like a shared pointer, it is valid to
    // do the copy.
    return this->CastToTypeStorage<ValueType, StorageTag>();
  }

  /// Given a refernce to an ArrayHandle object, casts this array to the
  /// ArrayHandle's type and sets the given ArrayHandle to this array. Throws
  /// \c ErrorBadType if the cast does not work. Use \c
  /// ArrayHandleType to check if the cast can happen.
  ///
  /// Note that this is a shallow copy. The data are not copied and a change
  /// in the data in one array will be reflected in the other.
  ///
  template <typename ArrayHandleType>
  VTKM_CONT void CopyTo(ArrayHandleType& array) const
  {
    VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
    array = this->Cast<ArrayHandleType>();
  }

  /// Changes the types to try casting to when resolving this dynamic array,
  /// which is specified with a list tag like those in TypeListTag.h. Since C++
  /// does not allow you to actually change the template arguments, this method
  /// returns a new dynamic array object. This method is particularly useful to
  /// narrow down (or expand) the types when using an array of particular
  /// constraints.
  ///
  template <typename NewTypeList>
  VTKM_CONT DynamicArrayHandleBase<NewTypeList, StorageList> ResetTypeList(
    NewTypeList = NewTypeList()) const
  {
    VTKM_IS_LIST_TAG(NewTypeList);
    return DynamicArrayHandleBase<NewTypeList, StorageList>(*this);
  }

  /// Changes the array storage types to try casting to when resolving this
  /// dynamic array, which is specified with a list tag like those in
  /// StorageListTag.h. Since C++ does not allow you to actually change the
  /// template arguments, this method returns a new dynamic array object. This
  /// method is particularly useful to narrow down (or expand) the types when
  /// using an array of particular constraints.
  ///
  template <typename NewStorageList>
  VTKM_CONT DynamicArrayHandleBase<TypeList, NewStorageList> ResetStorageList(
    NewStorageList = NewStorageList()) const
  {
    VTKM_IS_LIST_TAG(NewStorageList);
    return DynamicArrayHandleBase<TypeList, NewStorageList>(*this);
  }

  /// Changes the value, and array storage types to try casting to when
  /// resolving this dynamic array. Since C++ does not allow you to actually
  /// change the template arguments, this method returns a new dynamic array
  /// object. This method is particularly useful when you have both custom
  /// types and traits.
  ///
  template <typename NewTypeList, typename NewStorageList>
  VTKM_CONT DynamicArrayHandleBase<NewTypeList, NewStorageList> ResetTypeAndStorageLists(
    NewTypeList = NewTypeList(),
    NewStorageList = NewStorageList()) const
  {
    VTKM_IS_LIST_TAG(NewTypeList);
    VTKM_IS_LIST_TAG(NewStorageList);
    return DynamicArrayHandleBase<NewTypeList, NewStorageList>(*this);
  }

  /// Attempts to cast the held array to a specific value type and storage,
  /// then call the given functor with the cast array. The types and storage
  /// tried in the cast are those in the lists defined by the TypeList and
  /// StorageList, respectively. The default \c DynamicArrayHandle sets these
  /// two lists to VTKM_DEFAULT_TYPE_LIST_TAG and VTK_DEFAULT_STORAGE_LIST_TAG,
  /// respectively.
  ///
  template <typename Functor>
  VTKM_CONT void CastAndCall(const Functor& f) const;

  /// \brief Create a new array of the same type as this array.
  ///
  /// This method creates a new array that is the same type as this one and
  /// returns a new dynamic array handle for it. This method is convenient when
  /// creating output arrays that should be the same type as some input array.
  ///
  VTKM_CONT
  DynamicArrayHandleBase<TypeList, StorageList> NewInstance() const
  {
    DynamicArrayHandleBase<TypeList, StorageList> newArray;
    newArray.ArrayContainer = this->ArrayContainer->NewInstance();
    return newArray;
  }

  /// \brief Get the number of components in each array value.
  ///
  /// This method will query the array type for the number of components in
  /// each value of the array. The number of components is determined by
  /// the \c VecTraits::NUM_COMPONENTS trait class.
  ///
  VTKM_CONT
  vtkm::IdComponent GetNumberOfComponents() const
  {
    return this->ArrayContainer->GetNumberOfComponents();
  }

  /// \brief Get the number of values in the array.
  ///
  VTKM_CONT
  vtkm::Id GetNumberOfValues() const { return this->ArrayContainer->GetNumberOfValues(); }

  VTKM_CONT
  void PrintSummary(std::ostream& out) const { this->ArrayContainer->PrintSummary(out); }

private:
  std::shared_ptr<vtkm::cont::detail::PolymorphicArrayHandleContainerBase> ArrayContainer;

  friend struct detail::DynamicArrayHandleCopyHelper;
};

using DynamicArrayHandle =
  vtkm::cont::DynamicArrayHandleBase<VTKM_DEFAULT_TYPE_LIST_TAG, VTKM_DEFAULT_STORAGE_LIST_TAG>;

namespace detail
{

template <typename Functor, typename Type>
struct DynamicArrayHandleTryStorage
{
  const DynamicArrayHandle* const Array;
  const Functor& Function;
  bool FoundCast;

  VTKM_CONT
  DynamicArrayHandleTryStorage(const DynamicArrayHandle& array, const Functor& f)
    : Array(&array)
    , Function(f)
    , FoundCast(false)
  {
  }

  template <typename Storage>
  VTKM_CONT void operator()(Storage)
  {
    this->DoCast(Storage(),
                 typename vtkm::cont::internal::IsValidArrayHandle<Type, Storage>::type());
  }

private:
  template <typename Storage>
  void DoCast(Storage, std::true_type)
  {
    if (!this->FoundCast && this->Array->template IsTypeAndStorage<Type, Storage>())
    {
      this->Function(this->Array->template CastToTypeStorage<Type, Storage>());
      this->FoundCast = true;
    }
  }

  template <typename Storage>
  void DoCast(Storage, std::false_type)
  {
    // This type of array handle cannot exist, so do nothing.
  }

  void operator=(const DynamicArrayHandleTryStorage<Functor, Type>&) = delete;
};

template <typename Functor, typename StorageList>
struct DynamicArrayHandleTryType
{
  const DynamicArrayHandle* const Array;
  const Functor& Function;
  bool FoundCast;

  VTKM_CONT
  DynamicArrayHandleTryType(const DynamicArrayHandle& array, const Functor& f)
    : Array(&array)
    , Function(f)
    , FoundCast(false)
  {
  }

  template <typename Type>
  VTKM_CONT void operator()(Type)
  {
    if (this->FoundCast)
    {
      return;
    }
    using TryStorageType = DynamicArrayHandleTryStorage<Functor, Type>;
    TryStorageType tryStorage = TryStorageType(*this->Array, this->Function);

    vtkm::ListForEach(tryStorage, StorageList());
    if (tryStorage.FoundCast)
    {
      this->FoundCast = true;
    }
  }

private:
  void operator=(const DynamicArrayHandleTryType<Functor, StorageList>&) = delete;
};

} // namespace detail

template <typename TypeList, typename StorageList>
template <typename Functor>
VTKM_CONT void DynamicArrayHandleBase<TypeList, StorageList>::CastAndCall(const Functor& f) const
{
  VTKM_IS_LIST_TAG(TypeList);
  VTKM_IS_LIST_TAG(StorageList);
  using TryTypeType = detail::DynamicArrayHandleTryType<Functor, StorageList>;

  // We cast this to a DynamicArrayHandle because at this point we are ignoring
  // the type/storage lists in it. There is no sense in adding more unnecessary
  // template cases.
  // The downside to this approach is that a copy is created, causing an
  // atomic increment, which affects both performance and library size.
  // For these reasons we have a specialization of this method to remove
  // the copy when the type/storage lists are the default
  DynamicArrayHandle t(*this);
  TryTypeType tryType = TryTypeType(t, f);

  vtkm::ListForEach(tryType, TypeList());
  if (!tryType.FoundCast)
  {
    std::ostringstream out;
    out << "Could not find appropriate cast for array in CastAndCall1.\n"
           "Array: ";
    this->PrintSummary(out);
    out << "TypeList: " << typeid(TypeList).name()
        << "\nStorageList: " << typeid(StorageList).name() << "\n";
    throw vtkm::cont::ErrorBadValue(out.str());
  }
}

template <>
template <typename Functor>
VTKM_CONT void
DynamicArrayHandleBase<VTKM_DEFAULT_TYPE_LIST_TAG, VTKM_DEFAULT_STORAGE_LIST_TAG>::CastAndCall(
  const Functor& f) const
{

  using TryTypeType = detail::DynamicArrayHandleTryType<Functor, VTKM_DEFAULT_STORAGE_LIST_TAG>;

  // We can remove the copy, as the current DynamicArrayHandle is already
  // the default one, and no reason to do an atomic increment and increase
  // library size, and reduce performance
  TryTypeType tryType = TryTypeType(*this, f);

  vtkm::ListForEach(tryType, VTKM_DEFAULT_TYPE_LIST_TAG());
  if (!tryType.FoundCast)
  {
    throw vtkm::cont::ErrorBadValue("Could not find appropriate cast for array in CastAndCall2.");
  }
}

namespace internal
{

template <typename TypeList, typename StorageList>
struct DynamicTransformTraits<vtkm::cont::DynamicArrayHandleBase<TypeList, StorageList>>
{
  using DynamicTag = vtkm::cont::internal::DynamicTransformTagCastAndCall;
};

} // namespace internal
}
} // namespace vtkm::cont

#endif //vtk_m_cont_DynamicArrayHandle_h
