#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/cont/DataSet.h>
#include <vtkm/cont/DataSetFieldAdd.h>

#include <vtkm/exec/CellInterpolate.h>
#include <vtkm/exec/ParametricCoordinates.h>

////
//// BEGIN-EXAMPLE UseWorkletVisitCellsWithPoints.cxx
////
namespace vtkm
{
namespace worklet
{

struct CellCenter : public vtkm::worklet::WorkletVisitCellsWithPoints
{
public:
  using ControlSignature = void(CellSetIn cellSet,
                                FieldInPoint inputPointField,
                                FieldOut outputCellField);
  using ExecutionSignature = _3(_1, PointCount, _2);

  using InputDomain = _1;

  template<typename CellShape, typename InputPointFieldType>
  VTKM_EXEC typename InputPointFieldType::ComponentType operator()(
    CellShape shape,
    vtkm::IdComponent numPoints,
    const InputPointFieldType& inputPointField) const
  {
    vtkm::Vec3f parametricCenter =
      vtkm::exec::ParametricCoordinatesCenter(numPoints, shape, *this);
    return vtkm::exec::CellInterpolate(
      inputPointField, parametricCenter, shape, *this);
  }
};

} // namespace worklet
} // namespace vtkm
////
//// END-EXAMPLE UseWorkletVisitCellsWithPoints.cxx
////

#include <vtkm/filter/FilterField.h>

#include <vtkm/filter/CreateResult.h>

////
//// BEGIN-EXAMPLE UseFilterFieldWithCells.cxx
////
namespace vtkm
{
namespace filter
{

class CellCenters : public vtkm::filter::FilterField<CellCenters>
{
public:
  VTKM_CONT
  CellCenters();

  template<typename ArrayHandleType, typename Policy>
  VTKM_CONT vtkm::cont::DataSet DoExecute(
    const vtkm::cont::DataSet& inDataSet,
    const ArrayHandleType& inField,
    const vtkm::filter::FieldMetadata& FieldMetadata,
    vtkm::filter::PolicyBase<Policy>);
};

} // namespace filter
} // namespace vtkm
////
//// END-EXAMPLE UseFilterFieldWithCells.cxx
////

////
//// BEGIN-EXAMPLE FilterFieldWithCellsImpl.cxx
////
namespace vtkm
{
namespace filter
{

VTKM_CONT
CellCenters::CellCenters()
{
  this->SetOutputFieldName("");
}

template<typename ArrayHandleType, typename Policy>
VTKM_CONT cont::DataSet CellCenters::DoExecute(
  const vtkm::cont::DataSet& inDataSet,
  const ArrayHandleType& inField,
  const vtkm::filter::FieldMetadata& fieldMetadata,
  vtkm::filter::PolicyBase<Policy>)
{
  VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
  using ValueType = typename ArrayHandleType::ValueType;

  if (!fieldMetadata.IsPointField())
  {
    throw vtkm::cont::ErrorBadType("Cell Centers filter operates on point data.");
  }

  vtkm::cont::DynamicCellSet cellSet = inDataSet.GetCellSet();

  vtkm::cont::ArrayHandle<ValueType> outField;

  this->Invoke(vtkm::worklet::CellCenter{},
               vtkm::filter::ApplyPolicyCellSet(cellSet, Policy()),
               inField,
               outField);

  std::string outFieldName = this->GetOutputFieldName();
  if (outFieldName == "")
  {
    outFieldName = fieldMetadata.GetName() + "_center";
  }

  return vtkm::filter::CreateResultFieldCell(inDataSet, outField, outFieldName);
}

} // namespace filter
} // namespace vtkm
////
//// END-EXAMPLE FilterFieldWithCellsImpl.cxx
////

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

namespace
{

void CheckCellCenters(const vtkm::cont::DataSet& dataSet)
{
  std::cout << "Checking cell centers." << std::endl;
  vtkm::cont::CellSetStructured<3> cellSet;
  dataSet.GetCellSet().CopyTo(cellSet);

  vtkm::cont::ArrayHandle<vtkm::Vec3f> cellCentersArray;
  dataSet.GetCellField("cell_center").GetData().CopyTo(cellCentersArray);

  VTKM_TEST_ASSERT(cellSet.GetNumberOfCells() ==
                     cellCentersArray.GetNumberOfValues(),
                   "Cell centers array has wrong number of values.");

  vtkm::Id3 cellDimensions = cellSet.GetCellDimensions() - vtkm::Id3(1);

  vtkm::cont::ArrayHandle<vtkm::Vec3f>::PortalConstControl cellCentersPortal =
    cellCentersArray.GetPortalConstControl();

  vtkm::Id cellIndex = 0;
  for (vtkm::Id kIndex = 0; kIndex < cellDimensions[2]; kIndex++)
  {
    for (vtkm::Id jIndex = 0; jIndex < cellDimensions[1]; jIndex++)
    {
      for (vtkm::Id iIndex = 0; iIndex < cellDimensions[0]; iIndex++)
      {
        vtkm::Vec3f center = cellCentersPortal.Get(cellIndex);
        VTKM_TEST_ASSERT(test_equal(center[0], iIndex + 0.5), "Bad X coord.");
        VTKM_TEST_ASSERT(test_equal(center[1], jIndex + 0.5), "Bad Y coord.");
        VTKM_TEST_ASSERT(test_equal(center[2], kIndex + 0.5), "Bad Z coord.");
        cellIndex++;
      }
    }
  }
}

void Test()
{
  vtkm::cont::testing::MakeTestDataSet makeTestDataSet;

  std::cout << "Making test data set." << std::endl;
  vtkm::cont::DataSet dataSet = makeTestDataSet.Make3DUniformDataSet0();

  std::cout << "Finding cell centers with filter." << std::endl;
  vtkm::filter::CellCenters cellCentersFilter;
  cellCentersFilter.SetUseCoordinateSystemAsField(true);
  cellCentersFilter.SetOutputFieldName("cell_center");
  vtkm::cont::DataSet results = cellCentersFilter.Execute(dataSet);

  CheckCellCenters(results);
}

} // anonymous namespace

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