#include <vtkm/worklet/WorkletPointNeighborhood.h>

#include <vtkm/exec/BoundaryState.h>
#include <vtkm/exec/FieldNeighborhood.h>

#include <vtkm/cont/DataSet.h>

namespace vtkm
{
namespace worklet
{

////
//// BEGIN-EXAMPLE UseWorkletPointNeighborhood.cxx
////
class ApplyBoxKernel : public vtkm::worklet::WorkletPointNeighborhood
{
private:
  vtkm::IdComponent NumberOfLayers;

public:
  using ControlSignature = void(CellSetIn cellSet,
                                FieldInNeighborhood inputField,
                                FieldOut outputField);
  using ExecutionSignature = _3(_2, Boundary);

  using InputDomain = _1;

  ApplyBoxKernel(vtkm::IdComponent kernelSize)
  {
    VTKM_ASSERT(kernelSize >= 3);
    VTKM_ASSERT((kernelSize % 2) == 1);

    this->NumberOfLayers = (kernelSize - 1) / 2;
  }

  template<typename InputFieldPortalType>
  VTKM_EXEC typename InputFieldPortalType::ValueType operator()(
    const vtkm::exec::FieldNeighborhood<InputFieldPortalType>& inputField,
    const vtkm::exec::BoundaryState& boundary) const
  {
    using T = typename InputFieldPortalType::ValueType;

    ////
    //// BEGIN-EXAMPLE GetNeighborhoodBoundary.cxx
    ////
    auto minIndices = boundary.MinNeighborIndices(this->NumberOfLayers);
    auto maxIndices = boundary.MaxNeighborIndices(this->NumberOfLayers);

    T sum = 0;
    vtkm::IdComponent size = 0;
    for (vtkm::IdComponent k = minIndices[2]; k <= maxIndices[2]; ++k)
    {
      for (vtkm::IdComponent j = minIndices[1]; j <= maxIndices[1]; ++j)
      {
        for (vtkm::IdComponent i = minIndices[0]; i <= maxIndices[0]; ++i)
        {
          ////
          //// BEGIN-EXAMPLE GetNeighborhoodFieldValue.cxx
          ////
          sum = sum + inputField.Get(i, j, k);
          ////
          //// END-EXAMPLE GetNeighborhoodFieldValue.cxx
          ////
          ++size;
        }
      }
    }
    ////
    //// END-EXAMPLE GetNeighborhoodBoundary.cxx
    ////

    return static_cast<T>(sum / size);
  }
};
////
//// END-EXAMPLE UseWorkletPointNeighborhood.cxx
////

} // namespace worklet
} // namespace vtkm

#include <vtkm/filter/FilterField.h>

#include <vtkm/filter/CreateResult.h>

namespace vtkm
{
namespace filter
{

class BoxFilter : public vtkm::filter::FilterField<BoxFilter>
{
public:
  VTKM_CONT
  BoxFilter(vtkm::IdComponent kernelSize);

  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>);

private:
  vtkm::IdComponent KernelSize;
};

} // namespace filter
} // namespace vtkm

namespace vtkm
{
namespace filter
{

VTKM_CONT
BoxFilter::BoxFilter(vtkm::IdComponent kernelSize)
  : KernelSize(kernelSize)
{
  this->SetOutputFieldName("");
}

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

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

  ////
  //// BEGIN-EXAMPLE GetCellSetStructured.cxx
  ////
  auto cellSet = vtkm::filter::ApplyPolicyCellSetStructured(
    inDataSet.GetCellSet(), Policy{}, *this);
  ////
  //// END-EXAMPLE GetCellSetStructured.cxx
  ////

  using ValueType = typename ArrayHandleType::ValueType;

  vtkm::cont::ArrayHandle<ValueType> outField;
  this->Invoke(
    vtkm::worklet::ApplyBoxKernel(this->KernelSize), cellSet, inField, outField);

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

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

} // namespace filter
} // namespace vtkm

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

namespace
{

void CheckBoxFilter(const vtkm::cont::DataSet& dataSet)
{
  std::cout << "Check box filter." << std::endl;

  static const vtkm::Id NUM_POINTS = 18;
  static const vtkm::Float32 expected[NUM_POINTS] = {
    60.1875f, 65.2f,    70.2125f, 60.1875f, 65.2f,    70.2125f,
    90.2667f, 95.2778f, 100.292f, 90.2667f, 95.2778f, 100.292f,
    120.337f, 125.35f,  130.363f, 120.337f, 125.35f,  130.363f
  };

  vtkm::cont::ArrayHandle<vtkm::Float32> outputArray;
  dataSet.GetPointField("pointvar_average").GetData().CopyTo(outputArray);

  vtkm::cont::printSummary_ArrayHandle(outputArray, std::cout, true);

  VTKM_TEST_ASSERT(outputArray.GetNumberOfValues() == NUM_POINTS);

  auto portal = outputArray.ReadPortal();
  for (vtkm::Id index = 0; index < portal.GetNumberOfValues(); ++index)
  {
    vtkm::Float32 computed = portal.Get(index);
    VTKM_TEST_ASSERT(
      test_equal(expected[index], computed), "Unexpected value at index ", index);
  }
}

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

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

  std::cout << "Running box filter." << std::endl;
  vtkm::filter::BoxFilter boxFilter(3);
  boxFilter.SetActiveField("pointvar");
  boxFilter.SetOutputFieldName("pointvar_average");
  vtkm::cont::DataSet results = boxFilter.Execute(dataSet);

  CheckBoxFilter(results);
}

} // anonymous namespace

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