#include <vtkm/worklet/DispatcherMapField.h>
#include <vtkm/worklet/WorkletMapField.h>

#include <vtkm/cont/ArrayCopy.h>
#include <vtkm/cont/DataSetBuilderUniform.h>
#include <vtkm/cont/PointLocatorUniformGrid.h>

#include <vtkm/Math.h>

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

namespace
{

constexpr vtkm::Id DimensionSize = 50;
const vtkm::Id3 DimensionSizes = vtkm::Id3(DimensionSize);

////
//// BEGIN-EXAMPLE UsePointLocator.cxx
////
/// Worklet that generates for each input coordinate a unit vector that points
/// to the closest point in a locator.
struct PointToClosest : public vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn, ExecObject, WholeArrayIn, FieldOut);
  using ExecutionSignature = void(_1, _2, _3, _4);

  template<typename Point,
           typename PointLocatorExecObject,
           typename CoordinateSystemPortal,
           typename OutType>
  VTKM_EXEC void operator()(const Point& queryPoint,
                            const PointLocatorExecObject& pointLocator,
                            const CoordinateSystemPortal& coordinateSystem,
                            OutType& out) const
  {
    // Use the point locator to find the point in the locator closest to the point
    // given.
    vtkm::Id pointId;
    vtkm::FloatDefault distanceSquared;
    pointLocator->FindNearestNeighbor(queryPoint, pointId, distanceSquared);

    // Use this information to find the nearest point and create a unit vector
    // pointing to it.
    if (pointId >= 0)
    {
      // Get nearest point coordinate.
      auto point = coordinateSystem.Get(pointId);

      // Get the vector pointing to this point
      out = point - queryPoint;

      // Convert to unit vector (if possible)
      if (distanceSquared > vtkm::Epsilon<vtkm::FloatDefault>())
      {
        out = vtkm::RSqrt(distanceSquared) * out;
      }
    }
    else
    {
      this->RaiseError("Locator could not find closest point.");
    }
  }

  VTKM_CONT static vtkm::cont::ArrayHandle<vtkm::Vec3f> Run(
    vtkm::cont::PointLocator* pointLocator,
    const vtkm::cont::ArrayHandle<vtkm::Vec3f>& queryPoints)
  {
    pointLocator->Update();

    vtkm::cont::ArrayHandle<vtkm::Vec3f> pointDirections;

    vtkm::worklet::DispatcherMapField<PointToClosest> dispatcher;
    dispatcher.Invoke(
      queryPoints, pointLocator, pointLocator->GetCoordinates(), pointDirections);

    return pointDirections;
  }
}; // struct QueryCells
////
//// END-EXAMPLE UsePointLocator.cxx
////

template<typename PointLocatorType>
void TestPointLocator(PointLocatorType& pointLocator)
{
  using ValueType = vtkm::Vec3f;
  using ArrayType = vtkm::cont::ArrayHandle<ValueType>;

  ArrayType points;
  vtkm::cont::ArrayCopy(vtkm::cont::ArrayHandleUniformPointCoordinates(
                          DimensionSizes - vtkm::Id3(1), ValueType(0.75f)),
                        points);

  ArrayType pointers = PointToClosest::Run(&pointLocator, points);

  auto expected = vtkm::cont::make_ArrayHandleConstant(vtkm::Vec3f(0.57735f),
                                                       points.GetNumberOfValues());

  std::cout << "Expected: ";
  vtkm::cont::printSummary_ArrayHandle(expected, std::cout);

  std::cout << "Calculated: ";
  vtkm::cont::printSummary_ArrayHandle(pointers, std::cout);

  VTKM_TEST_ASSERT(test_equal_portals(expected.GetPortalConstControl(),
                                      pointers.GetPortalConstControl()));
}

void TestPointLocatorUniformGrid()
{
  std::cout << "*** UniformGrid" << std::endl;
  vtkm::cont::DataSet data =
    vtkm::cont::DataSetBuilderUniform::Create(DimensionSizes);

  // Note: when more point locators are created, we might want to switch the
  // example to a different (perhaps more general) one.
  ////
  //// BEGIN-EXAMPLE ConstructPointLocator.cxx
  ////
  vtkm::cont::PointLocatorUniformGrid locator;
  locator.SetCoordinates(data.GetCoordinateSystem());
  locator.Update();
  ////
  //// END-EXAMPLE ConstructPointLocator.cxx
  ////

  TestPointLocator(locator);
}

void Run()
{
  TestPointLocatorUniformGrid();
}

} // anonymous namespace

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