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

#include <vtkm/exec/CellEdge.h>

#include <vtkm/Hash.h>

#include <vtkm/worklet/AverageByKey.h>
#include <vtkm/worklet/Keys.h>
#include <vtkm/worklet/ScatterCounting.h>
#include <vtkm/worklet/WorkletMapTopology.h>
#include <vtkm/worklet/WorkletReduceByKey.h>

#include <vtkm/filter/FilterDataSet.h>

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

//#define CHECK_COLLISIONS

namespace vtkm
{
namespace worklet
{

namespace
{

////
//// BEGIN-EXAMPLE GenerateMeshHashCount.cxx
////
struct CountEdgesWorklet : vtkm::worklet::WorkletVisitCellsWithPoints
{
  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 cellShape,
    vtkm::IdComponent numPointsInCell) const
  {
    vtkm::IdComponent numEdges;
    vtkm::ErrorCode status =
      vtkm::exec::CellEdgeNumberOfEdges(numPointsInCell, cellShape, numEdges);
    if (status != vtkm::ErrorCode::Success)
    {
      // There is an error in the cell. As good as it would be to return an
      // error, we probably don't want to invalidate the entire run if there
      // is just one malformed cell. Instead, ignore the cell.
      return 0;
    }
    return numEdges;
  }
};
////
//// END-EXAMPLE GenerateMeshHashCount.cxx
////

////
//// BEGIN-EXAMPLE GenerateMeshHashGenHashes.cxx
////
class EdgeHashesWorklet : public vtkm::worklet::WorkletVisitCellsWithPoints
{
public:
  using ControlSignature = void(CellSetIn cellSet, FieldOut hashValues);
  using ExecutionSignature = _2(CellShape cellShape,
                                PointIndices globalPointIndices,
                                VisitIndex localEdgeIndex);
  using InputDomain = _1;

  using ScatterType = vtkm::worklet::ScatterCounting;

  template<typename CellShapeTag, typename PointIndexVecType>
  VTKM_EXEC vtkm::HashType operator()(
    CellShapeTag cellShape,
    const PointIndexVecType& globalPointIndicesForCell,
    vtkm::IdComponent localEdgeIndex) const
  {
    vtkm::IdComponent numPointsInCell =
      globalPointIndicesForCell.GetNumberOfComponents();
    vtkm::Id2 canonicalId;
    vtkm::ErrorCode status =
      vtkm::exec::CellEdgeCanonicalId(numPointsInCell,
                                      localEdgeIndex,
                                      cellShape,
                                      globalPointIndicesForCell,
                                      canonicalId);
    if (status != vtkm::ErrorCode::Success)
    {
      this->RaiseError(vtkm::ErrorString(status));
      return vtkm::HashType(-1);
    }
    //// PAUSE-EXAMPLE
#ifndef CHECK_COLLISIONS
    //// RESUME-EXAMPLE
    return vtkm::Hash(canonicalId);
    //// PAUSE-EXAMPLE
#else  // ! CHECK_COLLISIONS
    // Intentionally use a bad hash value to cause collisions to check to make
    // sure that collision resolution works.
    return vtkm::HashType(canonicalId[0]);
#endif // !CHECK_COLLISIONS
    //// RESUME-EXAMPLE
  }
};
////
//// END-EXAMPLE GenerateMeshHashGenHashes.cxx
////

////
//// BEGIN-EXAMPLE GenerateMeshHashResolveCollisions.cxx
////
class EdgeHashCollisionsWorklet : public vtkm::worklet::WorkletReduceByKey
{
public:
  using ControlSignature = void(KeysIn keys,
                                WholeCellSetIn<> inputCells,
                                ValuesIn originCells,
                                ValuesIn originEdges,
                                ValuesOut localEdgeIndices,
                                ReducedValuesOut numEdges);
  using ExecutionSignature = _6(_2 inputCells,
                                _3 originCells,
                                _4 originEdges,
                                _5 localEdgeIndices);
  using InputDomain = _1;

  template<typename CellSetType,
           typename OriginCellsType,
           typename OriginEdgesType,
           typename localEdgeIndicesType>
  VTKM_EXEC vtkm::IdComponent operator()(
    const CellSetType& cellSet,
    const OriginCellsType& originCells,
    const OriginEdgesType& originEdges,
    localEdgeIndicesType& localEdgeIndices) const
  {
    vtkm::IdComponent numEdgesInHash = localEdgeIndices.GetNumberOfComponents();

    // Sanity checks.
    VTKM_ASSERT(originCells.GetNumberOfComponents() == numEdgesInHash);
    VTKM_ASSERT(originEdges.GetNumberOfComponents() == numEdgesInHash);

    // Clear out localEdgeIndices
    for (vtkm::IdComponent index = 0; index < numEdgesInHash; ++index)
    {
      localEdgeIndices[index] = -1;
    }

    // Count how many unique edges there are and create an id for each;
    vtkm::IdComponent numUniqueEdges = 0;
    for (vtkm::IdComponent firstEdgeIndex = 0; firstEdgeIndex < numEdgesInHash;
         ++firstEdgeIndex)
    {
      if (localEdgeIndices[firstEdgeIndex] == -1)
      {
        vtkm::IdComponent edgeId = numUniqueEdges;
        localEdgeIndices[firstEdgeIndex] = edgeId;
        // Find all matching edges.
        vtkm::Id firstCellIndex = originCells[firstEdgeIndex];
        vtkm::Id2 canonicalEdgeId;
        vtkm::exec::CellEdgeCanonicalId(cellSet.GetNumberOfIndices(firstCellIndex),
                                        originEdges[firstEdgeIndex],
                                        cellSet.GetCellShape(firstCellIndex),
                                        cellSet.GetIndices(firstCellIndex),
                                        canonicalEdgeId);
        for (vtkm::IdComponent laterEdgeIndex = firstEdgeIndex + 1;
             laterEdgeIndex < numEdgesInHash;
             ++laterEdgeIndex)
        {
          vtkm::Id laterCellIndex = originCells[laterEdgeIndex];
          vtkm::Id2 otherCanonicalEdgeId;
          vtkm::exec::CellEdgeCanonicalId(cellSet.GetNumberOfIndices(laterCellIndex),
                                          originEdges[laterEdgeIndex],
                                          cellSet.GetCellShape(laterCellIndex),
                                          cellSet.GetIndices(laterCellIndex),
                                          otherCanonicalEdgeId);
          if (canonicalEdgeId == otherCanonicalEdgeId)
          {
            localEdgeIndices[laterEdgeIndex] = edgeId;
          }
        }
        ++numUniqueEdges;
      }
    }

    return numUniqueEdges;
  }
};
////
//// END-EXAMPLE GenerateMeshHashResolveCollisions.cxx
////

////
//// BEGIN-EXAMPLE GenerateMeshHashGenIndices.cxx
////
class EdgeIndicesWorklet : public vtkm::worklet::WorkletReduceByKey
{
public:
  using ControlSignature = void(KeysIn keys,
                                WholeCellSetIn<> inputCells,
                                ValuesIn originCells,
                                ValuesIn originEdges,
                                ValuesIn localEdgeIndices,
                                ReducedValuesOut connectivityOut);
  using ExecutionSignature = void(_2 inputCells,
                                  _3 originCell,
                                  _4 originEdge,
                                  _5 localEdgeIndices,
                                  VisitIndex localEdgeIndex,
                                  _6 connectivityOut);
  using InputDomain = _1;

  using ScatterType = vtkm::worklet::ScatterCounting;

  template<typename CellSetType,
           typename OriginCellsType,
           typename OriginEdgesType,
           typename LocalEdgeIndicesType>
  VTKM_EXEC void operator()(const CellSetType& cellSet,
                            const OriginCellsType& originCells,
                            const OriginEdgesType& originEdges,
                            const LocalEdgeIndicesType& localEdgeIndices,
                            vtkm::IdComponent localEdgeIndex,
                            vtkm::Id2& connectivityOut) const
  {
    // Find the first edge that matches the index given and return it.
    for (vtkm::IdComponent edgeIndex = 0;; ++edgeIndex)
    {
      if (localEdgeIndices[edgeIndex] == localEdgeIndex)
      {
        vtkm::Id cellIndex = originCells[edgeIndex];
        vtkm::IdComponent numPointsInCell = cellSet.GetNumberOfIndices(cellIndex);
        vtkm::IdComponent edgeInCellIndex = originEdges[edgeIndex];
        auto cellShape = cellSet.GetCellShape(cellIndex);

        vtkm::IdComponent pointInCellIndex0;
        vtkm::exec::CellEdgeLocalIndex(
          numPointsInCell, 0, edgeInCellIndex, cellShape, pointInCellIndex0);
        vtkm::IdComponent pointInCellIndex1;
        vtkm::exec::CellEdgeLocalIndex(
          numPointsInCell, 1, edgeInCellIndex, cellShape, pointInCellIndex1);

        auto globalPointIndicesForCell = cellSet.GetIndices(cellIndex);
        connectivityOut[0] = globalPointIndicesForCell[pointInCellIndex0];
        connectivityOut[1] = globalPointIndicesForCell[pointInCellIndex1];

        break;
      }
    }
  }
};
////
//// END-EXAMPLE GenerateMeshHashGenIndices.cxx
////

////
//// BEGIN-EXAMPLE GenerateMeshHashAverageField.cxx
////
class AverageCellEdgesFieldWorklet : public vtkm::worklet::WorkletReduceByKey
{
public:
  using ControlSignature = void(KeysIn keys,
                                ValuesIn inFieldValues,
                                ValuesIn localEdgeIndices,
                                ReducedValuesOut averagedField);
  using ExecutionSignature = _4(_2 inFieldValues,
                                _3 localEdgeIndices,
                                VisitIndex localEdgeIndex);
  using InputDomain = _1;

  using ScatterType = vtkm::worklet::ScatterCounting;

  template<typename InFieldValuesType, typename LocalEdgeIndicesType>
  VTKM_EXEC typename InFieldValuesType::ComponentType operator()(
    const InFieldValuesType& inFieldValues,
    const LocalEdgeIndicesType& localEdgeIndices,
    vtkm::IdComponent localEdgeIndex) const
  {
    using FieldType = typename InFieldValuesType::ComponentType;

    FieldType averageField = FieldType(0);
    vtkm::IdComponent numValues = 0;
    for (vtkm::IdComponent reduceIndex = 0;
         reduceIndex < inFieldValues.GetNumberOfComponents();
         ++reduceIndex)
    {
      if (localEdgeIndices[reduceIndex] == localEdgeIndex)
      {
        FieldType fieldValue = inFieldValues[reduceIndex];
        averageField = averageField + fieldValue;
        ++numValues;
      }
    }
    VTKM_ASSERT(numValues > 0);
    return static_cast<FieldType>(averageField / numValues);
  }
};
////
//// END-EXAMPLE GenerateMeshHashAverageField.cxx
////

} // 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>
  VTKM_CONT vtkm::cont::DataSet DoExecute(const vtkm::cont::DataSet& inData,
                                          vtkm::filter::PolicyBase<Policy> 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::ScatterCounting::OutputToInputMapType OutputToInputCellMap;
  vtkm::worklet::Keys<vtkm::HashType> CellToEdgeKeys;
  vtkm::cont::ArrayHandle<vtkm::IdComponent> LocalEdgeIndices;
  std::shared_ptr<vtkm::worklet::ScatterCounting> HashCollisionScatter;
};

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

namespace vtkm
{
namespace filter
{

//// PAUSE-EXAMPLE
namespace
{

//// RESUME-EXAMPLE
////
//// BEGIN-EXAMPLE GenerateMeshHashInvoke.cxx
////
template<typename Policy>
inline VTKM_CONT vtkm::cont::DataSet ExtractEdges::DoExecute(
  const vtkm::cont::DataSet& inData,
  vtkm::filter::PolicyBase<Policy> policy)
{
  auto inCellSet =
    vtkm::filter::ApplyPolicyCellSet(inData.GetCellSet(), policy, *this);

  // First, count the edges in each cell.
  vtkm::cont::ArrayHandle<vtkm::IdComponent> edgeCounts;
  this->Invoke(vtkm::worklet::CountEdgesWorklet{}, inCellSet, edgeCounts);

  // Second, using these counts build a scatter that repeats a cell's visit
  // for each edge in the cell.
  vtkm::worklet::ScatterCounting scatter(edgeCounts);
  this->OutputToInputCellMap =
    scatter.GetOutputToInputMap(inCellSet.GetNumberOfCells());
  vtkm::worklet::ScatterCounting::VisitArrayType outputToInputEdgeMap =
    scatter.GetVisitArray(inCellSet.GetNumberOfCells());

  // Third, for each edge, extract a hash.
  vtkm::cont::ArrayHandle<vtkm::HashType> hashValues;
  this->Invoke(vtkm::worklet::EdgeHashesWorklet{}, scatter, inCellSet, hashValues);

  // Fourth, use a Keys object to combine all like hashes.
  this->CellToEdgeKeys = vtkm::worklet::Keys<vtkm::HashType>(hashValues);

  // Fifth, use a reduce-by-key to collect like hash values, resolve collisions,
  // and count the number of unique edges associated with each hash.
  vtkm::cont::ArrayHandle<vtkm::IdComponent> numUniqueEdgesInEachHash;
  this->Invoke(vtkm::worklet::EdgeHashCollisionsWorklet{},
               this->CellToEdgeKeys,
               inCellSet,
               this->OutputToInputCellMap,
               outputToInputEdgeMap,
               this->LocalEdgeIndices,
               numUniqueEdgesInEachHash);

  // Sixth, use a reduce-by-key to extract indices for each unique edge.
  // (Note that HashCollisionScatter is an std::shared_ptr because we cannot copy
  // ScatterCounting objects.)
  this->HashCollisionScatter.reset(
    new vtkm::worklet::ScatterCounting(numUniqueEdgesInEachHash));

  vtkm::cont::ArrayHandle<vtkm::Id> connectivityArray;
  //// LABEL InvokeEdgeIndices
  this->Invoke(vtkm::worklet::EdgeIndicesWorklet{},
               *this->HashCollisionScatter,
               this->CellToEdgeKeys,
               inCellSet,
               this->OutputToInputCellMap,
               outputToInputEdgeMap,
               this->LocalEdgeIndices,
               vtkm::cont::make_ArrayHandleGroupVec<2>(connectivityArray));

  // Seventh, use the created connectivity array to build a cell set.
  vtkm::cont::CellSetSingleType<> outCellSet;
  outCellSet.Fill(
    inCellSet.GetNumberOfPoints(), vtkm::CELL_SHAPE_LINE, 2, connectivityArray);

  vtkm::cont::DataSet outData;

  outData.SetCellSet(outCellSet);

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

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

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

  if (fieldMeta.IsPointField())
  {
    outputField = fieldMeta.AsField(inputArray); // pass through
  }
  else if (fieldMeta.IsCellField())
  {
    vtkm::cont::ArrayHandle<T> averageField;
    this->Invoke(vtkm::worklet::AverageCellEdgesFieldWorklet{},
                 *this->HashCollisionScatter,
                 this->CellToEdgeKeys,
                 vtkm::cont::make_ArrayHandlePermutation(this->OutputToInputCellMap,
                                                         inputArray),
                 this->LocalEdgeIndices,
                 averageField);
    outputField = fieldMeta.AsField(averageField);
  }
  else
  {
    return false;
  }

  result.AddField(outputField);

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

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

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

namespace
{

template<typename ConnectivityPortalType>
vtkm::Id FindEdge(const ConnectivityPortalType& connectivity, const vtkm::Id2& edge)
{
  vtkm::Id edgeIndex = 0;
  bool foundEdge = false;

  for (vtkm::Id connectivityIndex = 0;
       connectivityIndex < connectivity.GetNumberOfValues() - 1;
       connectivityIndex += 2)
  {
    vtkm::Id p0 = connectivity.Get(connectivityIndex + 0);
    vtkm::Id p1 = connectivity.Get(connectivityIndex + 1);
    if (((edge[0] == p0) && (edge[1] == p1)) || ((edge[0] == p1) && (edge[1] == p0)))
    {
      foundEdge = true;
      break;
    }
    edgeIndex++;
  }

  VTKM_TEST_ASSERT(foundEdge, "Did not find expected edge.");

  for (vtkm::Id connectivityIndex = 2 * (edgeIndex + 1);
       connectivityIndex < connectivity.GetNumberOfValues() - 1;
       connectivityIndex += 2)
  {
    vtkm::Id p0 = connectivity.Get(connectivityIndex + 0);
    vtkm::Id p1 = connectivity.Get(connectivityIndex + 1);
    if (((edge[0] == p0) && (edge[1] == p1)) || ((edge[0] == p1) && (edge[1] == p0)))
    {
      VTKM_TEST_FAIL("Edge duplicated.");
    }
  }

  return edgeIndex;
}

template<typename ConnectivityPortalType, typename FieldPortalType>
void CheckEdge(const ConnectivityPortalType& connectivity,
               const FieldPortalType& field,
               const vtkm::Id2& edge,
               const vtkm::Float32 expectedFieldValue)
{
  std::cout << "  Checking for edge " << edge << " with field value "
            << expectedFieldValue << std::endl;

  vtkm::Id edgeIndex = FindEdge(connectivity, edge);

  vtkm::Float32 fieldValue = field.Get(edgeIndex);
  VTKM_TEST_ASSERT(test_equal(expectedFieldValue, fieldValue), "Bad field value.");
}

void CheckOutput(const vtkm::cont::CellSetSingleType<>& cellSet,
                 const vtkm::cont::ArrayHandle<vtkm::Float32>& cellField)
{
  std::cout << "Num cells: " << cellSet.GetNumberOfCells() << std::endl;
  VTKM_TEST_ASSERT(cellSet.GetNumberOfCells() == 22, "Wrong # of cells.");

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

  std::cout << "Cell field:" << std::endl;
  vtkm::cont::printSummary_ArrayHandle(cellField, std::cout, true);

  auto connectivityPortal = connectivity.ReadPortal();
  auto fieldPortal = cellField.ReadPortal();
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(0, 1), 100.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(0, 3), 100.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(0, 4), 100.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(7, 4), 115.3f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(7, 6), 115.3f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(7, 3), 100.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(5, 1), 105.05f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(5, 4), 115.3f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(5, 6), 115.2f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(2, 1), 105.05f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(2, 3), 100.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(2, 6), 105.05f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(8, 1), 110.0f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(8, 2), 110.0f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(8, 5), 115.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(8, 6), 115.1f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(10, 5), 125.35f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(10, 6), 125.35f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(10, 8), 120.2f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(9, 4), 130.5f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(9, 7), 130.5f);
  CheckEdge(connectivityPortal, fieldPortal, vtkm::Id2(9, 10), 130.5f);
}

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

  vtkm::filter::ExtractEdges filter;

  vtkm::cont::DataSet outDataSet = filter.Execute(inDataSet);

  vtkm::cont::CellSetSingleType<> outCellSet;
  outDataSet.GetCellSet().AsCellSet(outCellSet);

  vtkm::cont::Field outCellField = outDataSet.GetField("cellvar");
  vtkm::cont::ArrayHandle<vtkm::Float32> outCellData;
  outCellField.GetData().AsArrayHandle(outCellData);

  CheckOutput(outCellSet, outCellData);
}

void DoTest()
{
  TryFilter();
}

} // anonymous namespace

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