//============================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//
//============================================================================

#include <fides/DataSetReader.h>

#include <string>
#include <unordered_map>
#include <vector>

#include <vtkm/cont/Algorithm.h>
#include <vtkm/worklet/DispatcherMapTopology.h>
#include <vtkm/worklet/ScatterPermutation.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#ifdef FIDES_USE_MPI
#include <mpi.h>
#endif

class CheckTopology : public vtkm::worklet::WorkletVisitCellsWithPoints
{
public:
  using ControlSignature = void(CellSetIn cellset, FieldInCell trueConn, FieldOutCell outCells);
  using ExecutionSignature = void(PointIndices, _2, _3);
  using InputDomain = _1;
  using ScatterType = vtkm::worklet::ScatterPermutation<>;

  template <typename ConnType>
  VTKM_EXEC void operator()(const ConnType& inputConn,
                            const vtkm::Id3& trueConn,
                            bool& isTrue) const
  {
    isTrue = true;
    for (int i = 0; i < 3; i++)
    {
      if (inputConn[i] != trueConn[i])
      {
        isTrue = false;
        break;
      }
    }
  }
};

int main(int argc, char** argv)
{
#ifdef FIDES_USE_MPI
  MPI_Init(&argc, &argv);
#endif

  if (argc != 3)
  {
    std::cerr << "Usage: " << argv[0] << "/path/to.json /path/to/dataroot\n";
    return 1;
  }

  int retVal = 0;
  fides::io::DataSetReader reader(argv[1]);
  std::unordered_map<std::string, std::string> paths;
  paths["source"] = std::string(argv[2]) + "/tris-explicit.bp";
  auto metaData = reader.ReadMetaData(paths);
  auto& nBlocks = metaData.Get<fides::metadata::Size>(fides::keys::NUMBER_OF_BLOCKS());
  if (nBlocks.NumberOfItems != 1)
  {
    std::cerr << "Error: expected 1 blocks, got " << nBlocks.NumberOfItems << std::endl;
    retVal = 1;
  }
  fides::metadata::MetaData selections;
  vtkm::cont::PartitionedDataSet output = reader.ReadDataSet(paths, selections);
  if (output.GetNumberOfPartitions() != 1)
  {
    std::cerr << "Error: expected 1 output blocks, got " << output.GetNumberOfPartitions()
              << std::endl;
    retVal = 1;
  }
  vtkm::cont::DataSet ds = output.GetPartition(0);

  std::vector<vtkm::Id3> groundTruth = { { 2, 8, 0 }, { 8, 14, 0 } };
  groundTruth.resize(96);
  vtkm::cont::ArrayHandle<bool> result;
  auto cellSet = ds.GetCellSet().AsCellSet<vtkm::cont::CellSetExplicit<>>();
  std::vector<vtkm::Id> cellsToVisit;
  cellsToVisit.push_back(0);
  cellsToVisit.push_back(1);
  vtkm::worklet::ScatterPermutation<> scatter(
    vtkm::cont::make_ArrayHandle(cellsToVisit, vtkm::CopyFlag::Off));
  vtkm::worklet::DispatcherMapTopology<CheckTopology> dispatcher(scatter);
  dispatcher.Invoke(
    cellSet, vtkm::cont::make_ArrayHandle(groundTruth, vtkm::CopyFlag::Off), result);
  bool ret = vtkm::cont::Algorithm::Reduce(result, true, vtkm::Product{});
  if (!ret)
  {
    std::cerr << "Error: The topology does not match the ground truth." << std::endl;
    retVal = 1;
  }

#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif

  return retVal;
}
