#include <vtkm/filter/FilterDataSetWithField.h>

#include <vtkm/worklet/Threshold.h>

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

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

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

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

//// PAUSE-EXAMPLE
namespace {

//// RESUME-EXAMPLE
class BlankCells : public vtkm::filter::FilterDataSetWithField<BlankCells>
{
public:
  template<typename ArrayHandleType, typename Policy, typename DeviceAdapter>
  VTKM_CONT
  vtkm::filter::Result
  DoExecute(const vtkm::cont::DataSet &inDataSet,
            const ArrayHandleType &inField,
            const vtkm::filter::FieldMetadata &fieldMetadata,
            vtkm::filter::PolicyBase<Policy>,
            DeviceAdapter);

  template <typename T, typename StorageType, typename Policy, typename Device>
  VTKM_CONT bool DoMapField(vtkm::filter::Result& result,
                            const vtkm::cont::ArrayHandle<T, StorageType>& input,
                            const vtkm::filter::FieldMetadata& fieldMeta,
                            const vtkm::filter::PolicyBase<Policy>& policy,
                            Device);

private:
  vtkm::worklet::Threshold Worklet;
};

//// PAUSE-EXAMPLE
} // anonymous namespace
//// RESUME-EXAMPLE

template<>
struct FilterTraits<vtkm::filter::BlankCells>
{
  struct InputFieldTypeList : vtkm::TypeListTagScalarAll {  };
};

}
} // namespace vtkm::filter
////
//// END-EXAMPLE BlankCellsFilterDeclaration.cxx
////

namespace vtkm {
namespace filter {

namespace {

////
//// BEGIN-EXAMPLE BlankCellsFilterDoExecute.cxx
////
template<typename ArrayHandleType, typename Policy, typename Device>
VTKM_CONT
vtkm::filter::Result
BlankCells::DoExecute(const vtkm::cont::DataSet &inData,
                      const ArrayHandleType &inField,
                      const vtkm::filter::FieldMetadata &fieldMetadata,
                      vtkm::filter::PolicyBase<Policy>,
                      Device)
{
  VTKM_IS_ARRAY_HANDLE(ArrayHandleType);
  VTKM_IS_DEVICE_ADAPTER_TAG(Device);

  if (!fieldMetadata.IsCellField())
  {
    throw vtkm::cont::ErrorBadValue("Blanking field must be a cell field.");
  }

  const vtkm::cont::DynamicCellSet& inCells =
      inData.GetCellSet(this->GetActiveCellSetIndex());

  vtkm::cont::DynamicCellSet outCells =
      this->Worklet.Run(vtkm::filter::ApplyPolicy(inCells, Policy()),
                        inField,
                        fieldMetadata.GetAssociation(),
                        vtkm::NotZeroInitialized(),
                        Device());

  vtkm::cont::DataSet outData;

  outData.AddCellSet(outCells);

  for (vtkm::IdComponent coordSystemIndex = 0;
       coordSystemIndex < inData.GetNumberOfCoordinateSystems();
       ++coordSystemIndex)
  {
    outData.AddCoordinateSystem(inData.GetCoordinateSystem(coordSystemIndex));
  }

  return vtkm::filter::Result(outData);
}
////
//// END-EXAMPLE BlankCellsFilterDoExecute.cxx
////

////
//// BEGIN-EXAMPLE BlankCellsFilterDoMapField.cxx
////
template <typename T, typename StorageType, typename Policy, typename Device>
inline VTKM_CONT
bool BlankCells::DoMapField(vtkm::filter::Result& result,
                            const vtkm::cont::ArrayHandle<T, StorageType>& input,
                            const vtkm::filter::FieldMetadata& fieldMeta,
                            const vtkm::filter::PolicyBase<Policy>&,
                            Device)
{
  vtkm::cont::Field output;

  if (fieldMeta.IsPointField())
  {
    output = fieldMeta.AsField(input); // pass through
  }
  else if (fieldMeta.IsCellField())
  {
    output = fieldMeta.AsField(this->Worklet.ProcessCellField(input, Device()));
  }
  else
  {
    return false;
  }

  result.GetDataSet().AddField(output);

  return true;
}
////
//// END-EXAMPLE BlankCellsFilterDoMapField.cxx
////

} // anonymous namespace

}
} // namespace vtkm::filter

VTKM_CONT
static void DoTest()
{
  std::cout << "Setting up data" << std::endl;
  vtkm::cont::testing::MakeTestDataSet makedata;
  vtkm::cont::DataSet inData = makedata.Make3DExplicitDataSetCowNose();

  vtkm::Id numInCells = inData.GetCellSet().GetNumberOfCells();

  using FieldType = vtkm::Float32;
  vtkm::cont::ArrayHandle<FieldType> inField;
  inField.Allocate(numInCells);
  SetPortal(inField.GetPortalControl());
  vtkm::cont::DataSetFieldAdd::AddCellField(inData, "field", inField);

  vtkm::cont::ArrayHandle<vtkm::IdComponent> maskArray;
  maskArray.Allocate(numInCells);
  auto maskPortal = maskArray.GetPortalControl();
  for (vtkm::Id cellIndex = 0; cellIndex < numInCells; ++cellIndex)
  {
    maskPortal.Set(cellIndex, static_cast<vtkm::IdComponent>(cellIndex%2));
  }
  vtkm::cont::DataSetFieldAdd::AddCellField(inData, "mask", maskArray);

  std::cout << "Run filter" << std::endl;
  vtkm::filter::BlankCells filter;

  vtkm::filter::Result result = filter.Execute(inData, "mask");
  VTKM_TEST_ASSERT(result.IsValid(), "Execute failed.");

  for (vtkm::IdComponent fieldIndex = 0;
       fieldIndex < inData.GetNumberOfFields();
       ++fieldIndex)
  {
    filter.MapFieldOntoOutput(result, inData.GetField(fieldIndex));
  }

  std::cout << "Checking output." << std::endl;
  vtkm::cont::DataSet outData = result.GetDataSet();
  vtkm::Id numOutCells = numInCells/2;
  VTKM_TEST_ASSERT(outData.GetCellSet().GetNumberOfCells() == numOutCells,
                   "Unexpected number of cells.");

  vtkm::cont::Field outCellField = outData.GetField("field");
  vtkm::cont::ArrayHandle<FieldType> outField;
  outCellField.GetData().CopyTo(outField);
  auto outFieldPortal = outField.GetPortalConstControl();
  for (vtkm::Id cellIndex = 0; cellIndex < numOutCells; ++cellIndex)
  {
    FieldType expectedValue = TestValue(2*cellIndex+1, FieldType());
    VTKM_TEST_ASSERT(outFieldPortal.Get(cellIndex) == expectedValue, "Bad value.");
  }
}

int FilterDataSetWithField(int, char*[])
{
  return vtkm::cont::testing::Testing::Run(DoTest);
}
