#include <vtkm/Assert.h>

#include <vtkm/cont/ArrayHandle.h>
#include <vtkm/cont/DataSet.h>
#include <vtkm/cont/DataSetBuilderUniform.h>
#include <vtkm/cont/DeviceAdapter.h>

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

#include <vtkm/worklet/WorkletMapField.h>

#include <vtkm/VectorAnalysis.h>

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

#include <vector>

namespace
{

////
//// BEGIN-EXAMPLE FictitiousFieldStorage.cxx
////
struct FooAttributes
{
  float Pressure;
  float Temperature;
  float Velocity[3];
  // And so on...
};

class FooFields
{
public:
  FooAttributes* GetAttributesArray();

  std::size_t GetSize() const;

  void Resize(std::size_t numberOfElements);
  //// PAUSE-EXAMPLE

  FooFields()
    : AttributesArray(new std::vector<FooAttributes>)
  {
  }

  FooFields(const FooFields&) = default;
  FooFields& operator=(const FooFields&) = default;

private:
  std::shared_ptr<std::vector<FooAttributes>> AttributesArray;
  //// RESUME-EXAMPLE
};
////
//// END-EXAMPLE FictitiousFieldStorage.cxx
////

inline FooAttributes* FooFields::GetAttributesArray()
{
  return this->AttributesArray->data();
}

std::size_t FooFields::GetSize() const
{
  return this->AttributesArray->size();
}

void FooFields::Resize(std::size_t numberOfElements)
{
  this->AttributesArray->resize(numberOfElements);
}

} // namespace

////
//// BEGIN-EXAMPLE ArrayPortalAdapter.cxx
////
namespace vtkm
{
namespace internal
{

// Note: FooAttributesPointer expected to be either FooAttributes* or
// const FooAttributes*
template<typename FooAttributesPointer>
class ArrayPortalFooPressure
{
  FooAttributesPointer AttributesArray = nullptr;
  vtkm::Id NumberOfValues = 0;

public:
  using ValueType = float;

  ArrayPortalFooPressure() = default;

  VTKM_CONT ArrayPortalFooPressure(FooAttributesPointer array,
                                   vtkm::Id numberOfValues)
    : AttributesArray(array)
    , NumberOfValues(numberOfValues)
  {
  }

  VTKM_EXEC_CONT vtkm::Id GetNumberOfValues() const { return this->NumberOfValues; }

  VTKM_EXEC_CONT ValueType Get(vtkm::Id index) const
  {
    VTKM_ASSERT(index >= 0);
    VTKM_ASSERT(index < this->GetNumberOfValues());
    return this->AttributesArray[index].Pressure;
  }

  // This template is a trick to not define Set if FooAttributesPointer
  // is const. That saves us from having to create separate implementations
  // of this portal for the read and write versions.
  template<typename T = FooAttributesPointer,
           typename = typename std::enable_if<
             !std::is_const<std::remove_pointer<T>>::value>::type>
  VTKM_EXEC_CONT void Set(vtkm::Id index, ValueType value) const
  {
    VTKM_ASSERT(index >= 0);
    VTKM_ASSERT(index < this->GetNumberOfValues());
    this->AttributesArray[index].Pressure = value;
  }
};

}
} // namespace vtkm::internal
////
//// END-EXAMPLE ArrayPortalAdapter.cxx
////

////
//// BEGIN-EXAMPLE StoragePrototype.cxx
////
namespace vtkm
{
namespace cont
{
namespace internal
{

template<typename T, class StorageTag>
class Storage;

}
}
} // namespace vtkm::cont::internal
////
//// END-EXAMPLE StoragePrototype.cxx
////

////
//// BEGIN-EXAMPLE StorageAdapter.cxx
////
// Includes or definition for ArrayPortalFooPressure

namespace vtkm
{
namespace cont
{
namespace internal
{

struct StorageTagFooPressure
{
};

template<>
class Storage<float, StorageTagFooPressure>
{
public:
  using ReadPortalType =
    vtkm::internal::ArrayPortalFooPressure<const FooAttributes*>;
  using WritePortalType = vtkm::internal::ArrayPortalFooPressure<FooAttributes*>;

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

  VTKM_CONT static vtkm::Id GetNumberOfValues(
    const vtkm::cont::internal::Buffer* buffers)
  {
    return static_cast<vtkm::Id>(
      buffers->GetNumberOfBytes() /
      static_cast<vtkm::BufferSizeType>(sizeof(FooAttributes)));
  }

  VTKM_CONT static void ResizeBuffers(vtkm::Id numValues,
                                      vtkm::cont::internal::Buffer* buffers,
                                      vtkm::CopyFlag preserve,
                                      vtkm::cont::Token& token)
  {
    buffers[0].SetNumberOfBytes(
      vtkm::internal::NumberOfValuesToNumberOfBytes<FooAttributes>(numValues),
      preserve,
      token);
  }

  VTKM_CONT static ReadPortalType CreateReadPortal(
    const vtkm::cont::internal::Buffer* buffers,
    vtkm::cont::DeviceAdapterId device,
    vtkm::cont::Token& token)
  {
    return ReadPortalType(reinterpret_cast<const FooAttributes*>(
                            buffers[0].ReadPointerDevice(device, token)),
                          GetNumberOfValues(buffers));
  }

  VTKM_CONT static WritePortalType CreateWritePortal(
    vtkm::cont::internal::Buffer* buffers,
    vtkm::cont::DeviceAdapterId device,
    vtkm::cont::Token& token)
  {
    return WritePortalType(
      reinterpret_cast<FooAttributes*>(buffers[0].WritePointerDevice(device, token)),
      GetNumberOfValues(buffers));
  }
};

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

namespace
{

////
//// BEGIN-EXAMPLE ArrayHandleAdapterMemManage.cxx
////
VTKM_CONT void FooFieldsDeleter(void* container)
{
  FooFields* fields = reinterpret_cast<FooFields*>(container);
  delete fields;
}

VTKM_CONT void FooFieldsReallocater(void*& memory,
                                    void*& container,
                                    vtkm::BufferSizeType oldSize,
                                    vtkm::BufferSizeType newSize)
{
  FooFields* fields = reinterpret_cast<FooFields*>(container);
  VTKM_ASSERT(static_cast<std::size_t>(oldSize) == fields->GetSize());
  fields->Resize(static_cast<std::size_t>(newSize) / sizeof(FooAttributes));
  memory = fields->GetAttributesArray();
}
////
//// END-EXAMPLE ArrayHandleAdapterMemManage.cxx
////

////
//// BEGIN-EXAMPLE ArrayHandleAdapter.cxx
////
class ArrayHandleFooPressure
  : public vtkm::cont::ArrayHandle<float,
                                   vtkm::cont::internal::StorageTagFooPressure>
{
private:
  using StorageType =
    vtkm::cont::internal::Storage<float,
                                  vtkm::cont::internal::StorageTagFooPressure>;

  VTKM_CONT vtkm::cont::internal::Buffer MakeBufferFromFields(
    const FooFields& fields)
  {
    FooFields* fieldsCopy = new FooFields(fields);
    //// LABEL make-buffer-start
    return vtkm::cont::internal::MakeBuffer(
      vtkm::cont::DeviceAdapterTagUndefined{},
      fieldsCopy->GetAttributesArray(),
      fieldsCopy,
      static_cast<vtkm::BufferSizeType>(fieldsCopy->GetSize() *
                                        sizeof(FooAttributes)),
      FooFieldsDeleter,
      //// LABEL make-buffer-end
      FooFieldsReallocater);
  }

public:
  VTKM_ARRAY_HANDLE_SUBCLASS_NT(
    ArrayHandleFooPressure,
    (vtkm::cont::ArrayHandle<float, vtkm::cont::internal::StorageTagFooPressure>));

  //// LABEL wrapper-constructor
  VTKM_CONT ArrayHandleFooPressure(const FooFields& fields)
    : Superclass(vtkm::cont::internal::CreateBuffers(MakeBufferFromFields(fields)))
  {
  }
};
////
//// END-EXAMPLE ArrayHandleAdapter.cxx
////

////
//// BEGIN-EXAMPLE UsingArrayHandleAdapter.cxx
////
struct AirPressureWorklet : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn elevation, FieldOut airPressure);

  VTKM_EXEC void operator()(vtkm::Vec3f position, float& airPressure) const
  {
    // Use linear interpolation to estimate atmospheric pressure based on
    // elevation in meters (0 = sea level). Atmospheric pressure is 101325 Pa
    // at sea level and drops about 12 Pa per meter.
    airPressure =
      vtkm::Lerp(101325.0f, 77325.0f, static_cast<float>(position[2] / 2000.0f));
  }
};

VTKM_CONT
void GetElevationAirPressure(vtkm::cont::DataSet grid, const FooFields& fields)
{
  // Make an array handle that points to the pressure values in the fields.
  ArrayHandleFooPressure pressureHandle(fields);

  vtkm::cont::Invoker invoke;
  invoke(AirPressureWorklet{},
         grid.GetCoordinateSystem().GetDataAsMultiplexer(),
         pressureHandle);

  // Make sure the values are flushed back to the control environment.
  pressureHandle.SyncControlArray();

  // Now the pressure field is in the fields container.
}
////
//// END-EXAMPLE UsingArrayHandleAdapter.cxx
////

void Test()
{
  vtkm::cont::DataSet grid =
    vtkm::cont::DataSetBuilderUniform::Create(vtkm::Id3(2, 2, 50));

  FooFields fields;
  fields.Resize(4 * 50);
  GetElevationAirPressure(grid, fields);

  vtkm::Float32 value = 101325.0f;
  for (std::size_t heightIndex = 0; heightIndex < 50; heightIndex++)
  {
    for (std::size_t slabIndex = 0; slabIndex < 4; slabIndex++)
    {
      VTKM_TEST_ASSERT(
        test_equal(fields.GetAttributesArray()[4 * heightIndex + slabIndex].Pressure,
                   value),
        "Bad value.");
    }
    value -= 12.0f;
  }
}

} // anonymous namespace

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