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

#include <vtkm/cont/ArrayCopy.h>
#include <vtkm/cont/CellLocatorBoundingIntervalHierarchy.h>
#include <vtkm/cont/CellLocatorGeneral.h>
#include <vtkm/cont/CellLocatorRectilinearGrid.h>
#include <vtkm/cont/CellLocatorUniformBins.h>
#include <vtkm/cont/CellLocatorUniformGrid.h>
#include <vtkm/cont/DataSetBuilderRectilinear.h>
#include <vtkm/cont/DataSetBuilderUniform.h>

#include <vtkm/VecFromPortalPermute.h>

#include <vtkm/exec/CellInterpolate.h>

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

namespace
{

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

////
//// BEGIN-EXAMPLE UseCellLocator.cxx
////
struct QueryCells : public vtkm::worklet::WorkletMapField
{
  using ControlSignature =
    void(FieldIn, ExecObject, WholeCellSetIn<Cell, Point>, WholeArrayIn, FieldOut);
  using ExecutionSignature = void(_1, _2, _3, _4, _5);

  template<typename Point,
           typename CellLocatorExecObject,
           typename CellSet,
           typename FieldPortal,
           typename OutType>
  VTKM_EXEC void operator()(const Point& point,
                            const CellLocatorExecObject& cellLocator,
                            const CellSet& cellSet,
                            const FieldPortal& field,
                            OutType& out) const
  {
    // Use the cell locator to find the cell containing the point and the parametric
    // coordinates within that cell.
    vtkm::Id cellId;
    vtkm::Vec3f parametric;
    cellLocator->FindCell(point, cellId, parametric, *this);

    // Use this information to interpolate the point field to the given location.
    if (cellId >= 0)
    {
      // Get shape information about the cell containing the point coordinate
      auto cellShape = cellSet.GetCellShape(cellId);
      auto indices = cellSet.GetIndices(cellId);

      // Make a Vec-like containing the field data at the cell's points
      auto fieldValues = vtkm::make_VecFromPortalPermute(&indices, &field);

      // Do the interpolation
      out = vtkm::exec::CellInterpolate(fieldValues, parametric, cellShape, *this);
    }
    else
    {
      this->RaiseError("Given point outside of the cell set.");
    }
  }

  template<typename FieldType>
  VTKM_CONT static vtkm::cont::ArrayHandle<FieldType> Run(
    vtkm::cont::CellLocator* cellLocator,
    const vtkm::cont::ArrayHandle<vtkm::Vec3f>& queryPoints,
    const vtkm::cont::ArrayHandle<FieldType>& field)
  {
    cellLocator->Update();

    vtkm::cont::ArrayHandle<FieldType> interpolatedField;

    vtkm::worklet::DispatcherMapField<QueryCells> dispatcher;
    dispatcher.Invoke(
      queryPoints, cellLocator, cellLocator->GetCellSet(), field, interpolatedField);

    return interpolatedField;
  }
}; // struct QueryCells
////
//// END-EXAMPLE UseCellLocator.cxx
////

template<typename CellLocatorType>
void TestCellLocator(CellLocatorType& cellLocator)
{
  using ValueType = vtkm::Vec3f;
  using ArrayType = vtkm::cont::ArrayHandle<ValueType>;

  ArrayType inField;
  vtkm::cont::ArrayCopy(vtkm::cont::ArrayHandleUniformPointCoordinates(
                          DimensionSizes, ValueType(0.0f), ValueType(2.0f)),
                        inField);

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

  ArrayType interpolated = QueryCells::Run(&cellLocator, points, inField);

  vtkm::cont::ArrayHandleUniformPointCoordinates expected(
    DimensionSizes - vtkm::Id3(1), ValueType(1.0f), ValueType(2.0f));

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

  std::cout << "Interpolated: ";
  vtkm::cont::printSummary_ArrayHandle(interpolated, std::cout);

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

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

  vtkm::cont::CellLocatorUniformGrid locator;
  locator.SetCellSet(data.GetCellSet());
  locator.SetCoordinates(data.GetCoordinateSystem());

  TestCellLocator(locator);
}

void TestCellLocatorRectilinearGrid()
{
  std::cout << "*** RectilinearGrid" << std::endl;
  vtkm::cont::ArrayHandle<vtkm::FloatDefault> axis;
  vtkm::cont::ArrayCopy(
    vtkm::cont::ArrayHandleCounting<vtkm::FloatDefault>(0, 1, DimensionSize), axis);
  vtkm::cont::DataSet data =
    vtkm::cont::DataSetBuilderRectilinear::Create(axis, axis, axis);

  vtkm::cont::CellLocatorRectilinearGrid locator;
  locator.SetCellSet(data.GetCellSet());
  locator.SetCoordinates(data.GetCoordinateSystem());

  TestCellLocator(locator);
}

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

  vtkm::cont::CellLocatorUniformBins locator;
  locator.SetDensityL1(30.0f);
  locator.SetDensityL2(1.0f);
  locator.SetCellSet(data.GetCellSet());
  locator.SetCoordinates(data.GetCoordinateSystem());

  TestCellLocator(locator);
}

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

  vtkm::cont::CellLocatorBoundingIntervalHierarchy locator;
  locator.SetCellSet(data.GetCellSet());
  locator.SetCoordinates(data.GetCoordinateSystem());
  locator.SetNumberOfSplittingPlanes(5);
  locator.SetMaxLeafSize(10);

  TestCellLocator(locator);
}

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

  ////
  //// BEGIN-EXAMPLE ConstructCellLocator.cxx
  ////
  vtkm::cont::CellLocatorGeneral locator;
  locator.SetCellSet(data.GetCellSet());
  locator.SetCoordinates(data.GetCoordinateSystem());
  locator.Update();
  ////
  //// END-EXAMPLE ConstructCellLocator.cxx
  ////

  TestCellLocator(locator);
}

void Run()
{
  TestCellLocatorUniformGrid();
  TestCellLocatorRectilinearGrid();
  TestCellLocatorUniformBins();
  TestCellLocatorBIH();
  TestCellLocatorGeneral();
}

} // anonymous namespace

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