#include <vtkm/cont/ArrayHandleGroupVec.h>
#include <vtkm/cont/CellSetSingleType.h>

#include <vtkm/exec/CellEdge.h>

#include <vtkm/worklet/DispatcherMapTopology.h>
#include <vtkm/worklet/ScatterCounting.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/filter/FilterDataSet.h>

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

namespace vtkm
{
namespace worklet
{

namespace
{

struct ExtractEdges
{
  ////
  //// BEGIN-EXAMPLE CreateMeshConstantShapeCount.cxx
  ////
  struct CountEdges : vtkm::worklet::WorkletMapPointToCell
  {
    using ControlSignature = void(CellSetIn cellSet, FieldOut<> numEdges);
    using ExecutionSignature = _2(CellShape, PointCount);
    using InputDomain = _1;

    template<typename CellShapeTag>
    VTKM_EXEC_CONT vtkm::IdComponent operator()(CellShapeTag shape,
                                                vtkm::IdComponent numPoints) const
    {
      return vtkm::exec::CellEdgeNumberOfEdges(numPoints, shape, *this);
    }
  };
  ////
  //// END-EXAMPLE CreateMeshConstantShapeCount.cxx
  ////

  ////
  //// BEGIN-EXAMPLE CreateMeshConstantShapeGenIds.cxx
  ////
  class EdgeIndices : public vtkm::worklet::WorkletMapPointToCell
  {
  public:
    using ControlSignature = void(CellSetIn cellSet, FieldOut<> edgeIds);
    using ExecutionSignature = void(CellShape, PointIndices, _2, VisitIndex);
    using InputDomain = _1;

    using ScatterType = vtkm::worklet::ScatterCounting;
    VTKM_CONT
    ScatterType GetScatter() const { return this->Scatter; }

    VTKM_CONT
    explicit EdgeIndices(const ScatterType& scatter)
      : Scatter(scatter)
    {
    }

    template<typename CountArrayType, typename Device>
    VTKM_CONT EdgeIndices(CountArrayType&& countArray, Device)
      : Scatter(std::forward<CountArrayType>(countArray), Device())
    {
    }

    template<typename CellShapeTag, typename PointIndexVecType>
    VTKM_EXEC void operator()(CellShapeTag shape,
                              const PointIndexVecType& pointIndices,
                              vtkm::Vec<vtkm::Id, 2>& edgeIndices,
                              vtkm::IdComponent visitIndex) const
    {
      vtkm::Vec<vtkm::IdComponent, 2> localEdgeIndices =
        vtkm::exec::CellEdgeLocalIndices(
          pointIndices.GetNumberOfComponents(), visitIndex, shape, *this);
      edgeIndices[0] = pointIndices[localEdgeIndices[0]];
      edgeIndices[1] = pointIndices[localEdgeIndices[1]];
    }

  private:
    ScatterType Scatter;
  };
  ////
  //// END-EXAMPLE CreateMeshConstantShapeGenIds.cxx
  ////

  template<typename CellSetType, typename Device>
  static VTKM_CONT vtkm::cont::CellSetSingleType<> Run(const CellSetType& inCellSet,
                                                       Device)
  {
    VTKM_IS_DYNAMIC_OR_STATIC_CELL_SET(CellSetType);

    ////
    //// BEGIN-EXAMPLE CreateMeshConstantShapeInvoke.cxx
    ////
    // Extract edges from a CellSet named inCellSet

    vtkm::cont::ArrayHandle<vtkm::IdComponent> edgeCounts;
    vtkm::worklet::DispatcherMapTopology<CountEdges, Device> countEdgeDispatcher;
    countEdgeDispatcher.Invoke(inCellSet, edgeCounts);

    vtkm::cont::ArrayHandle<vtkm::Id> connectivityArray;
    vtkm::worklet::DispatcherMapTopology<EdgeIndices, Device> edgeIndicesDispatcher(
      EdgeIndices(edgeCounts, Device()));
    edgeIndicesDispatcher.Invoke(
      inCellSet, vtkm::cont::make_ArrayHandleGroupVec<2>(connectivityArray));

    vtkm::cont::CellSetSingleType<> outCellSet(inCellSet.GetName());
    outCellSet.Fill(
      inCellSet.GetNumberOfPoints(), vtkm::CELL_SHAPE_LINE, 2, connectivityArray);
    ////
    //// END-EXAMPLE CreateMeshConstantShapeInvoke.cxx
    ////

    return outCellSet;
  }
};

} // anonymous namespace

} // namespace worklet
} // namespace vtkm

namespace vtkm
{
namespace filter
{

//// PAUSE-EXAMPLE
namespace
{

//// RESUME-EXAMPLE
class ExtractEdges : public vtkm::filter::FilterDataSet<ExtractEdges>
{
public:
  template<typename Policy, typename Device>
  VTKM_CONT vtkm::filter::Result DoExecute(const vtkm::cont::DataSet& inData,
                                           vtkm::filter::PolicyBase<Policy> policy,
                                           Device);
};

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

#if 0
//// RESUME-EXAMPLE
#include <vtkm/filter/ExtractEdges.hxx>
//// PAUSE-EXAMPLE
#endif
//// RESUME-EXAMPLE
} // namespace filter
} // namespace vtkm

namespace vtkm
{
namespace filter
{

//// PAUSE-EXAMPLE
namespace
{

//// RESUME-EXAMPLE
template<typename Policy, typename Device>
inline VTKM_CONT vtkm::filter::Result ExtractEdges::DoExecute(
  const vtkm::cont::DataSet& inData,
  vtkm::filter::PolicyBase<Policy> policy,
  Device)
{
  VTKM_IS_DEVICE_ADAPTER_TAG(Device);

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

  vtkm::cont::CellSetSingleType<> outCells = vtkm::worklet::ExtractEdges::Run(
    vtkm::filter::ApplyPolicy(inCells, policy), 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);
}

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

//// RESUME-EXAMPLE
} // namespace filter
} // namespace vtkm

void CheckOutput(const vtkm::cont::CellSetSingleType<>& cellSet)
{
  std::cout << "Num cells: " << cellSet.GetNumberOfCells() << std::endl;
  VTKM_TEST_ASSERT(cellSet.GetNumberOfCells() == 12 + 8 + 6 + 9,
                   "Wrong # of cells.");

  auto connectivity = cellSet.GetConnectivityArray(vtkm::TopologyElementTagPoint(),
                                                   vtkm::TopologyElementTagCell());
  std::cout << "Connectivity:" << std::endl;
  vtkm::cont::printSummary_ArrayHandle(connectivity, std::cout, true);

  auto connectivityPortal = connectivity.GetPortalConstControl();
  VTKM_TEST_ASSERT(connectivityPortal.Get(0) == 0, "Bad edge index");
  VTKM_TEST_ASSERT(connectivityPortal.Get(1) == 1, "Bad edge index");
  VTKM_TEST_ASSERT(connectivityPortal.Get(2) == 1, "Bad edge index");
  VTKM_TEST_ASSERT(connectivityPortal.Get(3) == 5, "Bad edge index");
  VTKM_TEST_ASSERT(connectivityPortal.Get(68) == 9, "Bad edge index");
  VTKM_TEST_ASSERT(connectivityPortal.Get(69) == 10, "Bad edge index");
}

namespace
{

void TryWorklet()
{
  std::cout << std::endl << "Trying calling worklet." << std::endl;
  vtkm::cont::DataSet dataSet =
    vtkm::cont::testing::MakeTestDataSet().Make3DExplicitDataSet5();
  vtkm::cont::CellSetExplicit<> inCellSet;
  dataSet.GetCellSet().CopyTo(inCellSet);

  vtkm::cont::CellSetSingleType<> outCellSet =
    vtkm::worklet::ExtractEdges::Run(inCellSet, VTKM_DEFAULT_DEVICE_ADAPTER_TAG());
  CheckOutput(outCellSet);
}

void TryFilter()
{
  std::cout << std::endl << "Trying calling filter." << std::endl;
  vtkm::cont::DataSet dataSet =
    vtkm::cont::testing::MakeTestDataSet().Make3DExplicitDataSet5();

  vtkm::filter::ExtractEdges filter;

  vtkm::filter::Result result = filter.Execute(dataSet);
  VTKM_TEST_ASSERT(result.IsDataSetValid(), "Execute failed.");

  vtkm::cont::CellSetSingleType<> outCellSet;
  result.GetDataSet().GetCellSet().CopyTo(outCellSet);
  CheckOutput(outCellSet);
}

void DoTest()
{
  TryWorklet();
  TryFilter();
}

} // anonymous namespace

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