#include <vtkm/filter/FilterDataSetWithField.h>

#include <vtkm/worklet/Threshold.h>

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

#include <vtkm/cont/DataSet.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:
  using SupportedTypes = vtkm::TypeListScalarAll;

  template<typename T, typename StorageType, typename Policy>
  VTKM_CONT vtkm::cont::DataSet DoExecute(
    const vtkm::cont::DataSet& inDataSet,
    const vtkm::cont::ArrayHandle<T, StorageType>& 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);
  //// PAUSE-EXAMPLE
  template<typename T, typename StorageType, typename Policy>
  VTKM_CONT void DemoApplyPolicyFieldOfType(
    const vtkm::cont::DataSet& inDataSet,
    const vtkm::cont::ArrayHandle<T, StorageType>& inField,
    const vtkm::filter::FieldMetadata& fieldMetadata,
    vtkm::filter::PolicyBase<Policy>);
  //// RESUME-EXAMPLE

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

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

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

namespace vtkm
{
namespace filter
{

namespace
{

////
//// BEGIN-EXAMPLE BlankCellsFilterDoExecute.cxx
////
template<typename T, typename StorageType, typename Policy>
VTKM_CONT vtkm::cont::DataSet BlankCells::DoExecute(
  const vtkm::cont::DataSet& inData,
  const vtkm::cont::ArrayHandle<T, StorageType>& inField,
  const vtkm::filter::FieldMetadata& fieldMetadata,
  vtkm::filter::PolicyBase<Policy>)
{
  //// PAUSE-EXAMPLE
  DemoApplyPolicyFieldOfType(inData, inField, fieldMetadata, Policy{});
  //// RESUME-EXAMPLE
  if (!fieldMetadata.IsCellField())
  {
    throw vtkm::cont::ErrorBadValue("Blanking field must be a cell field.");
  }

  ////
  //// BEGIN-EXAMPLE GetCellSet.cxx
  ////
  auto inCells =
    vtkm::filter::ApplyPolicyCellSet(inData.GetCellSet(), Policy{}, *this);
  ////
  //// END-EXAMPLE GetCellSet.cxx
  ////

  vtkm::cont::DynamicCellSet outCells = this->Worklet.Run(
    inCells, inField, fieldMetadata.GetAssociation(), vtkm::NotZeroInitialized());
  vtkm::cont::DataSet outData;

  outData.SetCellSet(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
////

template<typename T, typename StorageType, typename Policy>
VTKM_CONT void BlankCells::DemoApplyPolicyFieldOfType(
  const vtkm::cont::DataSet& inData,
  const vtkm::cont::ArrayHandle<T, StorageType>&,
  const vtkm::filter::FieldMetadata&,
  vtkm::filter::PolicyBase<Policy>)
{
  std::string secondaryFieldName = this->GetActiveFieldName();
  ////
  //// BEGIN-EXAMPLE GetSecondaryField.cxx
  ////
  vtkm::cont::Field secondaryField = inData.GetField(secondaryFieldName);
  //// LABEL secondaryFieldArrayHandle
  auto secondaryFieldArrayHandle =
    vtkm::filter::ApplyPolicyFieldOfType<T>(secondaryField, Policy{}, *this);
  ////
  //// END-EXAMPLE GetSecondaryField.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.WritePortal());
  inData.AddCellField("field", inField);

  vtkm::cont::ArrayHandle<vtkm::IdComponent> maskArray;
  maskArray.Allocate(numInCells);
  auto maskPortal = maskArray.WritePortal();
  for (vtkm::Id cellIndex = 0; cellIndex < numInCells; ++cellIndex)
  {
    maskPortal.Set(cellIndex, static_cast<vtkm::IdComponent>(cellIndex % 2));
  }
  inData.AddCellField("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.ReadPortal();
  for (vtkm::Id cellIndex = 0; cellIndex < numOutCells; ++cellIndex)
  {
    FieldType expectedValue = TestValue(2 * cellIndex + 1, FieldType());
    VTKM_TEST_ASSERT(test_equal(outFieldPortal.Get(cellIndex), expectedValue));
  }
}

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