#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>
  VTKM_CONT vtkm::cont::DataSet DoExecute(
    const vtkm::cont::DataSet& inDataSet,
    const ArrayHandleType& inField,
    const vtkm::filter::FieldMetadata& fieldMetadata,
    vtkm::filter::PolicyBase<Policy>);

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

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

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

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

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

namespace vtkm
{
namespace filter
{

namespace
{

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

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

  vtkm::cont::DataSet outData;

  outData.AddCellSet(outCells);

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

  return outData;
}
////
//// END-EXAMPLE BlankCellsFilterDoExecute.cxx
////

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

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

  result.AddField(output);

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

} // anonymous namespace

} // namespace filter
} // namespace vtkm

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;
  filter.SetActiveField("mask", vtkm::cont::Field::Association::CELL_SET);

  // NOTE 2018-03-21: I expect this to fail in the short term. Right now no fields
  // are copied from input to output. The default should be changed to copy them
  // all. (Also, I'm thinking it would be nice to have a mode to select all except
  // a particular field or list of fields.)
  vtkm::cont::DataSet outData = filter.Execute(inData);

  std::cout << "Checking output." << std::endl;
  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);
}
