//============================================================================
//  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 <fides/DataSetWriter.h>

#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/ArrayHandleCartesianProduct.h>
#include <vtkm/cont/ArrayHandleUniformPointCoordinates.h>
#include <vtkm/cont/ArrayPortalToIterators.h>
#include <vtkm/cont/CellSetStructured.h>
#include <vtkm/cont/DataSetBuilderExplicit.h>
#include <vtkm/cont/DataSetBuilderRectilinear.h>
#include <vtkm/cont/DataSetBuilderUniform.h>

#include <numeric>
#include <random>
#include <string>
#include <vector>

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

//vector of cell shape info. Pair is: <cellType, numVertices>
static const std::vector<std::pair<vtkm::UInt8, vtkm::IdComponent>> CellInfo = {
  std::make_pair(vtkm::CELL_SHAPE_VERTEX, 1),    std::make_pair(vtkm::CELL_SHAPE_LINE, 2),
  std::make_pair(vtkm::CELL_SHAPE_TRIANGLE, 3),  std::make_pair(vtkm::CELL_SHAPE_QUAD, 4),
  std::make_pair(vtkm::CELL_SHAPE_TETRA, 4),     std::make_pair(vtkm::CELL_SHAPE_HEXAHEDRON, 8),
  std::make_pair(vtkm::CELL_SHAPE_WEDGE, 6),     std::make_pair(vtkm::CELL_SHAPE_PYRAMID, 5),
  std::make_pair(vtkm::CELL_SHAPE_POLY_LINE, 2), std::make_pair(vtkm::CELL_SHAPE_POLYGON, 3)
};

class RandomNumbers
{
public:
  RandomNumbers()
    : Gen(0)
    , Dist(0, 1)
  {
  }

  vtkm::FloatDefault RandomFloat(vtkm::FloatDefault v0, vtkm::FloatDefault v1)
  {
    auto r = this->Dist(this->Gen);
    return v0 + (v1 - v0) * r;
  }

  template <typename T>
  T RandomVal(T v0, T v1)
  {
    vtkm::FloatDefault t = this->Dist(this->Gen);
    vtkm::FloatDefault dV = static_cast<vtkm::FloatDefault>(v1 - v0);

    return v0 + static_cast<T>(t * dV);
  }

  template <typename T>
  std::vector<T> RandomVec(T v0, T v1, std::size_t n)
  {
    std::vector<T> vec(n);
    do
    {
      for (auto& val : vec)
      {
        val = this->RandomVal(v0, v1);
      }
    } while (std::accumulate(vec.begin(), vec.end(), 0) == 0);
    //We don't want a vector full of 0, so if we get that, try again.

    return vec;
  }

private:
  std::default_random_engine Gen;
  std::uniform_real_distribution<vtkm::FloatDefault> Dist;
};

vtkm::cont::ArrayHandle<vtkm::FloatDefault> MakeArray(std::size_t numValues,
                                                      RandomNumbers& randomNumbers)
{
  std::vector<vtkm::FloatDefault> arr;

  arr.resize(numValues);
  for (auto& v : arr)
  {
    v = randomNumbers.RandomFloat(-100, 100);
  }

  return vtkm::cont::make_ArrayHandle(arr, vtkm::CopyFlag::On);
}

std::vector<vtkm::Vec3f> MakeVecVector(std::size_t numValues,
                                       vtkm::FloatDefault minVal,
                                       vtkm::FloatDefault maxVal,
                                       RandomNumbers& randomNumbers)
{
  std::vector<vtkm::Vec3f> arr;
  arr.resize(numValues);
  for (auto& v : arr)
  {
    v[0] = randomNumbers.RandomFloat(minVal, maxVal);
    v[1] = randomNumbers.RandomFloat(minVal, maxVal);
    v[2] = randomNumbers.RandomFloat(minVal, maxVal);
  }
  return arr;
}

vtkm::cont::ArrayHandle<vtkm::Vec3f> MakeVecArray(std::size_t numValues,
                                                  RandomNumbers& randomNumbers)
{
  auto arr = MakeVecVector(numValues, -100, 100, randomNumbers);
  return vtkm::cont::make_ArrayHandle(arr, vtkm::CopyFlag::On);
}

fides::metadata::Vector<std::size_t> CreateBlockSelection(int rank,
                                                          const std::vector<int>& numBlocks)
{
  int b0 = 0;
  for (int i = 0; i < rank; i++)
  {
    b0 += numBlocks[i];
  }
  int b1 = b0 + numBlocks[rank];

  fides::metadata::Vector<std::size_t> blockSelection;
  for (int bid = b0; bid < b1; bid++)
  {
    blockSelection.Data.push_back(static_cast<std::size_t>(bid));
  }

  return blockSelection;
}

vtkm::cont::PartitionedDataSet DataSetsByRank(const std::vector<vtkm::cont::DataSet>& dataSets,
                                              int rank,
                                              const std::vector<int>& blocksPerRank)
{
  int offset = 0;
  for (std::size_t i = 0; i < static_cast<size_t>(rank); i++)
  {
    offset += blocksPerRank[i];
  }

  vtkm::cont::PartitionedDataSet pds;
  for (int i = 0; i < blocksPerRank[rank]; i++)
  {
    pds.AppendPartition(dataSets[offset + i]);
  }
  return pds;
}

void AddFields(vtkm::cont::DataSet& ds, RandomNumbers& randomNumbers)
{
  std::size_t numPoints = static_cast<std::size_t>(ds.GetNumberOfPoints());
  std::size_t numCells = static_cast<std::size_t>(ds.GetNumberOfCells());

  // Add scalar fields.
  ds.AddPointField("varPoint", MakeArray(numPoints, randomNumbers));
  ds.AddCellField("varCell", MakeArray(numCells, randomNumbers));

  // Add vector fields.
  ds.AddPointField("vecPoint", MakeVecArray(numPoints, randomNumbers));
  ds.AddCellField("vecCell", MakeVecArray(numCells, randomNumbers));
}

std::vector<vtkm::cont::PartitionedDataSet> MakeUniformData(const std::vector<int>& numBlocks,
                                                            int rank,
                                                            RandomNumbers& randomNumbers)
{
  int totalNumBlocks = std::accumulate(numBlocks.begin(), numBlocks.end(), 0);

  std::vector<vtkm::cont::DataSet> dataSets;
  for (int i = 0; i < totalNumBlocks; i++)
  {
    vtkm::Vec3f spacing, origin;
    vtkm::Id3 dim;

    for (int j = 0; j < 3; j++)
    {
      spacing[j] = randomNumbers.RandomFloat(0.1, 10.);
      origin[j] = randomNumbers.RandomFloat(-100, 100);
      dim[j] = randomNumbers.RandomVal<int>(2, 10);
    }

    auto ds = vtkm::cont::DataSetBuilderUniform::Create(dim, origin, spacing);
    AddFields(ds, randomNumbers);
    dataSets.push_back(ds);
  }

  std::vector<vtkm::cont::PartitionedDataSet> ret;
  ret.push_back(DataSetsByRank(dataSets, rank, numBlocks));
  return ret;
}

std::vector<vtkm::cont::PartitionedDataSet> MakeRectilinearData(const std::vector<int>& numBlocks,
                                                                int rank,
                                                                RandomNumbers& randomNumbers)
{
  int totalNumBlocks = std::accumulate(numBlocks.begin(), numBlocks.end(), 0);
  std::vector<vtkm::cont::DataSet> dataSets;
  for (int i = 0; i < totalNumBlocks; i++)
  {
    std::vector<vtkm::FloatDefault> coords[3];

    for (int j = 0; j < 3; j++)
    {
      coords[j].resize(randomNumbers.RandomVal<int>(2, 20));
      auto v0 = randomNumbers.RandomFloat(-100, 100);
      for (auto& v : coords[j])
      {
        v = v0 + randomNumbers.RandomFloat(.1, 10);
      }
    }

    auto ds = vtkm::cont::DataSetBuilderRectilinear::Create(coords[0], coords[1], coords[2]);
    AddFields(ds, randomNumbers);
    dataSets.push_back(ds);
  }

  std::vector<vtkm::cont::PartitionedDataSet> ret;
  ret.push_back(DataSetsByRank(dataSets, rank, numBlocks));
  return ret;
}

vtkm::cont::ArrayHandle<vtkm::Id> MakeConnectivity(size_t numCoords,
                                                   size_t numPtsPerCell,
                                                   RandomNumbers& randomNumbers)
{
  size_t numCells = randomNumbers.RandomVal<size_t>(numCoords, 3 * numCoords);
  size_t totalNum = numCells * numPtsPerCell;
  std::vector<vtkm::Id> connectivity(totalNum, 0);

  for (auto& c : connectivity)
  {
    c = randomNumbers.RandomVal<vtkm::Id>(0, numCoords - 1);
  }

  return vtkm::cont::make_ArrayHandle(connectivity, vtkm::CopyFlag::On);
}

std::vector<vtkm::cont::PartitionedDataSet> MakeUnstructuredSingleData(
  const std::vector<int>& numBlocks,
  int rank,
  RandomNumbers& randomNumbers)
{
  std::vector<vtkm::cont::PartitionedDataSet> ret;
  int totalNumBlocks = std::accumulate(numBlocks.begin(), numBlocks.end(), 0);

  constexpr int numShapes = 8;

  for (int shapeIdx = 0; shapeIdx < numShapes; shapeIdx++)
  {
    std::vector<vtkm::cont::DataSet> dataSets;
    for (int i = 0; i < totalNumBlocks; i++)
    {
      size_t numCoords = randomNumbers.RandomVal<size_t>(10, 100);
      auto coords = MakeVecArray(numCoords, randomNumbers);

      vtkm::cont::DataSet ds;
      auto shape = CellInfo[shapeIdx].first;
      auto numVerts = CellInfo[shapeIdx].second;
      //Variable vert cells MUST have the same number of verts for all datasets.
      if (shape == vtkm::CELL_SHAPE_POLY_LINE || shape == vtkm::CELL_SHAPE_POLYGON)
      {
        numVerts = numVerts * 3;
      }

      auto connArray = MakeConnectivity(numCoords, numVerts, randomNumbers);
      vtkm::CellShapeTagGeneric shapeTag(shape);

      ds = vtkm::cont::DataSetBuilderExplicit::Create(coords, shapeTag, numVerts, connArray);
      AddFields(ds, randomNumbers);
      dataSets.push_back(ds);
    }

    ret.push_back(DataSetsByRank(dataSets, rank, numBlocks));
  }
  return ret;
}

std::vector<vtkm::cont::PartitionedDataSet> MakeExplicitData(const std::vector<int>& numBlocks,
                                                             int rank,
                                                             RandomNumbers& randomNumbers)
{
  std::vector<vtkm::cont::PartitionedDataSet> ret;

  int totalNumBlocks = std::accumulate(numBlocks.begin(), numBlocks.end(), 0);

  std::vector<vtkm::cont::DataSet> dataSets;
  for (int i = 0; i < totalNumBlocks; i++)
  {
    size_t numCoords = randomNumbers.RandomVal<size_t>(10, 100);
    numCoords = 10;
    auto coords = MakeVecVector(numCoords, -10, 10, randomNumbers);

    int numCells = randomNumbers.RandomVal<int>(1, 100);

    vtkm::cont::DataSetBuilderExplicit dsb;
    std::vector<vtkm::Id> connectivity;
    std::vector<vtkm::IdComponent> numIndices;
    std::vector<vtkm::UInt8> shapes;
    for (int j = 0; j < numCells; j++)
    {
      std::size_t cellIdx = randomNumbers.RandomVal<size_t>(0, CellInfo.size());

      auto shape = CellInfo[cellIdx].first;
      auto numVerts = CellInfo[cellIdx].second;
      if (shape == vtkm::CELL_SHAPE_POLY_LINE || shape == vtkm::CELL_SHAPE_POLYGON)
      {
        numVerts = randomNumbers.RandomVal<vtkm::IdComponent>(numVerts, 3 * numVerts);
      }

      numIndices.push_back(numVerts);
      shapes.push_back(shape);
      for (vtkm::Id k = 0; k < numVerts; k++)
      {
        connectivity.push_back(
          randomNumbers.RandomVal<vtkm::Id>(0, static_cast<vtkm::Id>(numCoords)));
      }
    }
    auto ds = vtkm::cont::DataSetBuilderExplicit::Create(coords, shapes, numIndices, connectivity);

    AddFields(ds, randomNumbers);
    dataSets.push_back(ds);
  }

  ret.push_back(DataSetsByRank(dataSets, rank, numBlocks));
  return ret;
}

vtkm::cont::PartitionedDataSet ReadData(const std::vector<int>& numBlocks,
                                        int rank,
                                        const std::string& fname)
{
  auto blockSelection = CreateBlockSelection(rank, numBlocks);
  fides::io::DataSetReader reader(fname, fides::io::DataSetReader::DataModelInput::BPFile);
  std::unordered_map<std::string, std::string> paths;
  paths["source"] = fname;
  auto metaData = reader.ReadMetaData(paths);
  metaData.Set(fides::keys::BLOCK_SELECTION(), blockSelection);
  auto output = reader.ReadDataSet(paths, metaData);

  return output;
}

int ValidateScalarArray(const vtkm::cont::ArrayHandle<vtkm::FloatDefault>& arr1,
                        const vtkm::cont::ArrayHandle<vtkm::FloatDefault>& arr2)
{
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> basicArr1(arr1);
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> basicArr2(arr2);

  vtkm::Id numVals = arr1.GetNumberOfValues();
  if (numVals != arr2.GetNumberOfValues())
  {
    std::cerr << "Wrong number of values in array" << std::endl;
    return 1;
  }

  const vtkm::FloatDefault* buff1 = basicArr1.GetReadPointer();
  const vtkm::FloatDefault* buff2 = basicArr2.GetReadPointer();
  const auto eps = std::numeric_limits<vtkm::FloatDefault>::epsilon();

  for (vtkm::Id i = 0; i < numVals; i++)
  {
    if (vtkm::Abs(buff1[i] - buff2[i]) > eps)
    {
      std::cerr << "Wrong array value at index " << i << ": " << buff1[i] << " " << buff2[i]
                << std::endl;
      return 1;
    }
  }

  return 0;
}

int ValidateVectorArray(const vtkm::cont::ArrayHandle<vtkm::Vec3f>& arr1,
                        const vtkm::cont::ArrayHandle<vtkm::Vec3f>& arr2)
{
  vtkm::cont::ArrayHandleBasic<vtkm::Vec3f> basicArr1(arr1);
  vtkm::cont::ArrayHandleBasic<vtkm::Vec3f> basicArr2(arr2);

  vtkm::Id numVals = arr1.GetNumberOfValues();
  if (numVals != arr2.GetNumberOfValues())
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " Wrong number of values in vector array"
              << std::endl;
    return 1;
  }

  const vtkm::Vec3f* buff1 = basicArr1.GetReadPointer();
  const vtkm::Vec3f* buff2 = basicArr2.GetReadPointer();
  const auto eps = std::numeric_limits<vtkm::FloatDefault>::epsilon();

  for (vtkm::Id i = 0; i < numVals; i++)
  {
    if (vtkm::Magnitude(buff1[i] - buff2[i]) > eps)
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Wrong array value at index " << i << ": "
                << buff1[i] << " " << buff2[i] << std::endl;
      return 1;
    }
  }

  return 0;
}

int ValidateScalar(const vtkm::cont::Field& field1, const vtkm::cont::Field& field2)
{
  auto arr1 = field1.GetData().AsArrayHandle<vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault>>();
  auto arr2 = field2.GetData().AsArrayHandle<vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault>>();
  vtkm::Id numVals = arr1.GetNumberOfValues();
  if (numVals != arr2.GetNumberOfValues())
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " Wrong number of values in field "
              << field1.GetName() << std::endl;
    return 1;
  }
  const vtkm::FloatDefault* buff1 = arr1.GetReadPointer();
  const vtkm::FloatDefault* buff2 = arr2.GetReadPointer();
  const auto eps = std::numeric_limits<vtkm::FloatDefault>::epsilon();

  for (vtkm::Id k = 0; k < numVals; k++)
  {
    if (vtkm::Abs(buff1[k] - buff2[k]) > eps)
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Wrong field value: " << field1.GetName()
                << " at " << k << ": " << buff1[k] << " " << buff2[k] << std::endl;
      return 1;
    }
  }
  return 0;
}

int ValidateVector(const vtkm::cont::Field& field1, const vtkm::cont::Field& field2)
{
  auto arr1 = field1.GetData().AsArrayHandle<vtkm::cont::ArrayHandleBasic<vtkm::Vec3f>>();
  auto arr2 = field2.GetData().AsArrayHandle<vtkm::cont::ArrayHandleBasic<vtkm::Vec3f>>();

  vtkm::Id numVals = arr1.GetNumberOfValues();
  if (numVals != arr2.GetNumberOfValues())
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " Wrong number of values in field "
              << field1.GetName() << std::endl;
    return 1;
  }
  const vtkm::Vec3f* buff1 = arr1.GetReadPointer();
  const vtkm::Vec3f* buff2 = arr2.GetReadPointer();
  const auto eps = std::numeric_limits<vtkm::FloatDefault>::epsilon();

  for (vtkm::Id k = 0; k < numVals; k++)
  {
    if (vtkm::Magnitude(buff1[k] - buff2[k]) > eps)
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Wrong field value: " << field1.GetName()
                << " at " << k << ": " << buff1[k] << " " << buff2[k] << std::endl;
      return 1;
    }
  }
  return 0;
}

int ValidateFields(const vtkm::cont::DataSet& ds1,
                   const vtkm::cont::DataSet& ds2,
                   const std::vector<std::string>& writeVars)
{
  bool writeAll = (writeVars.size() == 1 && writeVars[0] == "ALL");

  vtkm::IdComponent numFields = ds1.GetNumberOfFields();
  //ds2 should always have <= than ds1.
  if (numFields < ds2.GetNumberOfFields())
  {
    std::cerr << "Wrong number of fields in data set" << std::endl;
    return 1;
  }

  for (vtkm::IdComponent j = 0; j < numFields; j++)
  {
    auto field1 = ds1.GetField(j);
    auto fieldName = field1.GetName();

    // We do not validate CoordinatesSystems here
    if (ds1.HasCoordinateSystem(fieldName))
    {
      continue;
    }

    bool ds2HasField = ds2.HasField(fieldName);
    bool writeField =
      (writeAll || (std::find(writeVars.begin(), writeVars.end(), fieldName) != writeVars.end()));

    // ds2HasField and writeField are always the same.
    // otherwise it's an error.
    if (ds2HasField != writeField)
    {
      std::cerr << "Missing field in data set: " << fieldName << std::endl;
      return 1;
    }

    //We didn't write out this field, so nothing to validate.
    if (!ds2HasField)
    {
      continue;
    }

    auto field2 = ds2.GetField(fieldName);

    if (field1.GetName() != field2.GetName())
    {
      std::cerr << "Field names do not match for data set: " << field1.GetName() << " "
                << field2.GetName() << std::endl;
      return 1;
    }
    if (field1.GetAssociation() != field2.GetAssociation())
    {
      std::cerr << "Field assocation do not match for data set: " << field1.GetName() << " "
                << field2.GetName() << std::endl;
      return 1;
    }

    if (field1.GetData().IsType<vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault>>() &&
        field2.GetData().IsType<vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault>>())
    {
      if (ValidateScalar(field1, field2) != 0)
      {
        return 1;
      }
    }
    else if (field1.GetData().IsType<vtkm::cont::ArrayHandleBasic<vtkm::Vec3f>>() &&
             field2.GetData().IsType<vtkm::cont::ArrayHandleBasic<vtkm::Vec3f>>())
    {
      if (ValidateVector(field1, field2) != 0)
      {
        return 1;
      }
    }
    else
    {
      std::cerr << "Field data type do not match: " << field1.GetName() << " " << field2.GetName()
                << std::endl;
      return 1;
    }
  }

  return 0;
}

int ValidateUniformDataSet(const vtkm::cont::DataSet& ds1, const vtkm::cont::DataSet& ds2)
{
  using UniformCoordType = vtkm::cont::ArrayHandleUniformPointCoordinates;
  using UniformCellType = vtkm::cont::CellSetStructured<3>;

  if (!ds2.GetCoordinateSystem().GetData().IsType<UniformCoordType>())
  {
    std::cerr << "Wrong dataset type." << std::endl;
    return 1;
  }

  auto ucoords1 = ds1.GetCoordinateSystem().GetData().AsArrayHandle<UniformCoordType>();
  auto origin1 = ucoords1.ReadPortal().GetOrigin();
  auto spacing1 = ucoords1.ReadPortal().GetSpacing();
  auto cellSet1 = ds1.GetCellSet().AsCellSet<UniformCellType>();
  auto dim1 = cellSet1.GetPointDimensions();

  auto ucoords2 = ds2.GetCoordinateSystem().GetData().AsArrayHandle<UniformCoordType>();
  auto origin2 = ucoords2.ReadPortal().GetOrigin();
  auto spacing2 = ucoords2.ReadPortal().GetSpacing();
  auto cellSet2 = ds2.GetCellSet().AsCellSet<UniformCellType>();
  auto dim2 = cellSet2.GetPointDimensions();

  const auto eps = std::numeric_limits<vtkm::FloatDefault>::epsilon();
  if (dim1 != dim2)
  {
    std::cerr << "Uniform dataset dimensions do not match: " << dim1 << " " << dim2 << std::endl;
    return 1;
  }
  if (vtkm::MagnitudeSquared(origin1 - origin2) > eps)
  {
    std::cerr << "Uniform dataset origins do not match: " << origin1 << " " << origin2 << std::endl;
    return 1;
  }
  if (vtkm::MagnitudeSquared(spacing1 - spacing2) > eps)
  {
    std::cerr << "Uniform dataset spacings do not match: " << spacing1 << " " << spacing2
              << std::endl;
    return 1;
  }

  return 0;
}

int ValidateRectilinearDataSet(const vtkm::cont::DataSet& ds1, const vtkm::cont::DataSet& ds2)
{
  using RectilinearCoordType =
    vtkm::cont::ArrayHandleCartesianProduct<vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
                                            vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
                                            vtkm::cont::ArrayHandle<vtkm::FloatDefault>>;
  using StructuredCellType = vtkm::cont::CellSetStructured<3>;

  if (!ds2.GetCoordinateSystem().GetData().IsType<RectilinearCoordType>())
  {
    std::cerr << "Wrong dataset type." << std::endl;
    return 1;
  }

  const auto& coords1 = ds1.GetCoordinateSystem().GetData().AsArrayHandle<RectilinearCoordType>();
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords1x(coords1.GetFirstArray());
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords1y(coords1.GetSecondArray());
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords1z(coords1.GetThirdArray());
  auto cellSet1 = ds1.GetCellSet().AsCellSet<StructuredCellType>();
  auto dim1 = cellSet1.GetPointDimensions();

  const auto& coords2 = ds2.GetCoordinateSystem().GetData().AsArrayHandle<RectilinearCoordType>();
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords2x(coords2.GetFirstArray());
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords2y(coords2.GetSecondArray());
  vtkm::cont::ArrayHandleBasic<vtkm::FloatDefault> coords2z(coords2.GetThirdArray());
  auto cellSet2 = ds2.GetCellSet().AsCellSet<StructuredCellType>();
  auto dim2 = cellSet2.GetPointDimensions();

  if (dim1 != dim2)
  {
    std::cerr << "Rectilinear cellset dimensions do not match: " << dim1 << " " << dim2
              << std::endl;
    return 1;
  }

  if (ValidateScalarArray(coords1x, coords2x) != 0 ||
      ValidateScalarArray(coords1y, coords2y) != 0 || ValidateScalarArray(coords1z, coords2z) != 0)
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " Rectilinear coordinates are wrong" << std::endl;
    return 1;
  }

  return 0;
}

int ValidateUnstructuredSingleTypeDataSet(const vtkm::cont::DataSet& ds1,
                                          const vtkm::cont::DataSet& ds2)
{
  using UnstructuredSingleType = vtkm::cont::CellSetSingleType<>;
  using UnstructuredCoordsType = vtkm::cont::ArrayHandle<vtkm::Vec3f>;

  if (!ds1.GetCoordinateSystem().GetData().IsType<UnstructuredCoordsType>() ||
      !ds2.GetCoordinateSystem().GetData().IsType<UnstructuredCoordsType>() ||
      !ds1.GetCellSet().IsType<UnstructuredSingleType>() ||
      !ds2.GetCellSet().IsType<UnstructuredSingleType>())
  {
    std::cerr << "Wrong dataset type." << std::endl;
    return 1;
  }

  const auto& coords1 = ds1.GetCoordinateSystem().GetData().AsArrayHandle<UnstructuredCoordsType>();
  const auto& coords2 = ds2.GetCoordinateSystem().GetData().AsArrayHandle<UnstructuredCoordsType>();

  if (ValidateVectorArray(coords1, coords2) != 0)
  {
    std::cerr << "Unstructured coordinates do not match" << std::endl;
    return 1;
  }

  const auto& cellSet1 = ds1.GetCellSet().AsCellSet<UnstructuredSingleType>();
  const auto& cellSet2 = ds2.GetCellSet().AsCellSet<UnstructuredSingleType>();
  vtkm::Id numCells1 = cellSet1.GetNumberOfCells();
  if (numCells1 != cellSet2.GetNumberOfCells())
  {
    std::cerr << "UnstructuredSingleType has wrong number of cells" << std::endl;
    return 1;
  }

  if (cellSet1.GetCellShapeAsId() != cellSet2.GetCellShapeAsId())
  {
    std::cerr << "UnstructuredSingleType cell types do not match" << std::endl;
    return 1;
  }

  for (vtkm::Id i = 0; i < numCells1; i++)
  {
    vtkm::cont::ArrayHandle<vtkm::Id> ids1, ids2;

    vtkm::Id nids1 = cellSet1.GetNumberOfPointsInCell(i);
    if (nids1 != cellSet2.GetNumberOfPointsInCell(i))
    {
      std::cerr << "Wrong number of points in cell at cell index " << i << std::endl;
      return 1;
    }

    cellSet1.GetIndices(i, ids1);
    cellSet2.GetIndices(i, ids2);

    auto ids1ReadPortal = ids1.ReadPortal();
    auto ids2ReadPortal = ids2.ReadPortal();
    for (vtkm::Id j = 0; j < nids1; j++)
    {
      if (ids1ReadPortal.Get(j) != ids2ReadPortal.Get(j))
      {
        std::cerr << "Wrong connectivity index at cell index " << i << ", pointIndex " << j
                  << std::endl;
        return 1;
      }
    }
  }

  return 0;
}

int ValidateExplicitDataSet(const vtkm::cont::DataSet& ds1, const vtkm::cont::DataSet& ds2)
{
  using Explicit = vtkm::cont::CellSetExplicit<>;
  using UnstructuredCoordsType = vtkm::cont::ArrayHandle<vtkm::Vec3f>;

  const auto& coords1 = ds1.GetCoordinateSystem().GetData().AsArrayHandle<UnstructuredCoordsType>();
  const auto& coords2 = ds2.GetCoordinateSystem().GetData().AsArrayHandle<UnstructuredCoordsType>();

  int retVal = 0;
  if (ValidateVectorArray(coords1, coords2) != 0)
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " Unstructured coordinates do not match\n";
    retVal = 1;
  }

  const vtkm::cont::CellSetExplicit<>& cellSet1 = ds1.GetCellSet().AsCellSet<Explicit>();
  const auto& cellSet2 = ds2.GetCellSet().AsCellSet<Explicit>();
  vtkm::Id numCells1 = cellSet1.GetNumberOfCells();
  vtkm::Id numCells2 = cellSet2.GetNumberOfCells();

  if (numCells1 != numCells2)
  {
    std::cerr << __FILE__ << ":" << __LINE__ << " The recovered number of cells is " << numCells2
              << " but the expected number of cells is " << numCells1 << "\n";
    retVal = 1;
  }

  for (vtkm::Id i = 0; i < numCells1; ++i)
  {
    vtkm::Id nids1 = cellSet1.GetNumberOfPointsInCell(i);
    if (nids1 != cellSet2.GetNumberOfPointsInCell(i))
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Wrong number of points in cell at cell index "
                << i << "\n";
      retVal = 1;
    }

    auto shape1 = cellSet1.GetCellShape(i);
    auto shape2 = cellSet2.GetCellShape(i);
    if (shape1 != shape2)
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Wrong shape at index " << i << "\n";
      retVal = 1;
    }
  }

  auto conn1 =
    cellSet1.GetConnectivityArray(vtkm::TopologyElementTagCell(), vtkm::TopologyElementTagPoint());
  ;
  auto conn2 =
    cellSet2.GetConnectivityArray(vtkm::TopologyElementTagCell(), vtkm::TopologyElementTagPoint());
  ;
  auto rp1 = conn1.ReadPortal();
  auto rp2 = conn2.ReadPortal();
  if (rp1.GetNumberOfValues() != rp2.GetNumberOfValues())
  {
    std::cerr << __FILE__ << ":" << __LINE__
              << " Wrong number of values in the recovered connectivity array.\n";
    retVal = 1;
  }
  for (vtkm::Id i = 0; i < rp1.GetNumberOfValues(); ++i)
  {
    if (rp1.Get(i) != rp2.Get(i))
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Connectivity wrong at index " << i << "\n";
      retVal = 1;
    }
  }
  return retVal;
}

int ValidateDataSets(const vtkm::cont::PartitionedDataSet& pds1,
                     const vtkm::cont::PartitionedDataSet& pds2,
                     const std::vector<std::string>& writeVars)
{
  using UniformCoordType = vtkm::cont::ArrayHandleUniformPointCoordinates;
  using RectilinearCoordType =
    vtkm::cont::ArrayHandleCartesianProduct<vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
                                            vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
                                            vtkm::cont::ArrayHandle<vtkm::FloatDefault>>;
  using UnstructuredSingleType = vtkm::cont::CellSetSingleType<>;
  using Explicit = vtkm::cont::CellSetExplicit<>;

  if (pds1.GetNumberOfPartitions() != pds2.GetNumberOfPartitions())
  {
    std::cerr << __FILE__ << ":" << __LINE__
              << " Number of datasets does not match: " << pds2.GetNumberOfPartitions()
              << " should be " << pds1.GetNumberOfPartitions() << std::endl;
    return 1;
  }

  // Validate the field data for each data set.
  vtkm::Id numDataSets = pds1.GetNumberOfPartitions();
  for (vtkm::Id i = 0; i < numDataSets; i++)
  {
    const auto& ds1 = pds1.GetPartition(i);
    const auto& ds2 = pds2.GetPartition(i);

    if (ValidateFields(ds1, ds2, writeVars) != 0)
    {
      std::cerr << __FILE__ << ":" << __LINE__ << " Fields are invalid.\n";
      return 1;
    }

    if (ds1.GetCoordinateSystem().GetData().IsType<UniformCoordType>())
    {
      if (ValidateUniformDataSet(ds1, ds2) != 0)
      {
        std::cerr << __FILE__ << ":" << __LINE__ << " Uniform dataset is invalid.\n";
        return 1;
      }
    }
    else if (ds1.GetCoordinateSystem().GetData().IsType<RectilinearCoordType>())
    {
      if (ValidateRectilinearDataSet(ds1, ds2) != 0)
      {
        std::cerr << __FILE__ << ":" << __LINE__ << " Rectilinear dataset is invalid.\n";
        return 1;
      }
    }
    else if (ds1.GetCellSet().IsType<UnstructuredSingleType>())
    {
      if (ValidateUnstructuredSingleTypeDataSet(ds1, ds2))
      {
        std::cerr << __FILE__ << ":" << __LINE__ << " Unstructured dataset is invalid.\n";
        return 1;
      }
    }
    else if (ds1.GetCellSet().IsType<Explicit>())
    {
      if (ValidateExplicitDataSet(ds1, ds2))
      {
        std::cerr << __FILE__ << ":" << __LINE__ << " Explicit dataset is invalid.\n";
        return 1;
      }
    }
    else
    {
      std::string err =
        std::string(__FILE__) + ":" + std::to_string(__LINE__) + " Unrecognized CellSet type.\n";
      throw std::domain_error(err);
    }
  }

  return 0;
}

std::vector<vtkm::cont::PartitionedDataSet> GenerateTestDataSets(const std::string& dataSetType,
                                                                 const std::vector<int>& numBlocks,
                                                                 int rank,
                                                                 RandomNumbers& randomNumbers)
{
  std::vector<vtkm::cont::PartitionedDataSet> vecPDS;

  if (dataSetType == "uniform")
  {
    vecPDS = MakeUniformData(numBlocks, rank, randomNumbers);
  }
  else if (dataSetType == "rectilinear")
  {
    vecPDS = MakeRectilinearData(numBlocks, rank, randomNumbers);
  }
  else if (dataSetType == "unstructured_single")
  {
    vecPDS = MakeUnstructuredSingleData(numBlocks, rank, randomNumbers);
  }
  else if (dataSetType == "explicit")
  {
    vecPDS = MakeExplicitData(numBlocks, rank, randomNumbers);
  }
  else
  {
    throw std::domain_error("Unrecognized dataset type " + dataSetType);
  }
  return vecPDS;
}

int RunTests(const std::vector<std::vector<int>>& numBlocksArray,
             const std::string& dataSetType,
             int rank,
             int numRanks,
             RandomNumbers& randomNumbers)
{
  int retVal = 0;
  const std::string fname = "test.bp";
  std::stringstream ss;
  ss << "rm -rf ./" << fname;

  for (auto& numBlocks : numBlocksArray)
  {
    if (static_cast<int>(numBlocks.size()) != numRanks)
      throw std::domain_error("Wrong number of blocks in numBlocksArray");

    auto vecPDS1 = GenerateTestDataSets(dataSetType, numBlocks, rank, randomNumbers);

    std::vector<std::vector<std::string>> writeFields;
    writeFields = { { "ALL" }, //write all fields
                    {},
                    { "varPoint" },
                    { "varPoint", "vecPoint" },
                    { "varPoint", "varCell", "vecPoint", "vecCell" } };

    for (const auto& wf : writeFields)
    {
      for (const auto& pds1 : vecPDS1)
      {
        fides::io::DataSetWriter writer(fname);
        if (!(wf.size() == 1 && wf[0] == "ALL"))
          writer.SetWriteFields(wf);
        writer.Write(pds1, "BPFile");
        auto pds2 = ReadData(numBlocks, rank, fname);
        retVal = ValidateDataSets(pds1, pds2, wf);

#ifdef FIDES_USE_MPI
        MPI_Barrier(MPI_COMM_WORLD);
#endif
        if (rank == 0)
        {
          int result = system(ss.str().c_str());
          if (result != 0)
          {
            std::cout << "rm command failed!\n";
          }
        }

        if (retVal != 0)
        {
          return retVal;
        }
      }
    }
  }

  return retVal;
}

int RunAppendTests(const std::vector<std::vector<int>>& numBlocksArray,
                   const std::string& dataSetType,
                   int rank,
                   int numRanks,
                   RandomNumbers& randomNumbers)
{
  int retVal = 0;
  const std::string fname = "test.bp";
  std::stringstream ss;
  ss << "rm -rf ./" << fname;

  for (auto& numBlocks : numBlocksArray)
  {
    if (static_cast<int>(numBlocks.size()) != numRanks)
      throw std::domain_error("Wrong number of blocks in numBlocksArray");

    std::vector<std::vector<vtkm::cont::PartitionedDataSet>> allPDS1;

    static const std::size_t nTime = 5;
    for (std::size_t i = 0; i < nTime; i++)
    {
      auto pds = GenerateTestDataSets(dataSetType, numBlocks, rank, randomNumbers);
      allPDS1.push_back(pds);
    }
    std::size_t numPDS = allPDS1[0].size();

    std::vector<std::vector<std::string>> writeFields;
    writeFields = { { "ALL" }, //write all fields
                    {},
                    { "varPoint" },
                    { "varPoint", "vecPoint" },
                    { "varPoint", "varCell", "vecPoint", "vecCell" } };

    for (const auto& wf : writeFields)
    {
      for (std::size_t i = 0; i < numPDS; i++)
      {
        //Need to remove file BEFORE create the writer. Otherwise the engine close when
        //the reader destructs will fail.
#ifdef FIDES_USE_MPI
        MPI_Barrier(MPI_COMM_WORLD);
#endif
        if (rank == 0)
        {
          int result = system(ss.str().c_str());
          if (result != 0)
          {
            std::cout << "rm command failed!\n";
          }
        }

        //Write each step.
        fides::io::DataSetAppendWriter writer(fname);
        if (!(wf.size() == 1 && wf[0] == "ALL"))
          writer.SetWriteFields(wf);

        for (std::size_t t = 0; t < nTime; t++)
        {
          auto pds = allPDS1[t][i];
          writer.Write(pds, "BPFile");
        }
        writer.Close();

        //Now, read them back in.
        auto blockSelection = CreateBlockSelection(rank, numBlocks);

        fides::io::DataSetReader reader(fname, fides::io::DataSetReader::DataModelInput::BPFile);
        std::unordered_map<std::string, std::string> paths;
        paths["source"] = fname;
        auto metaData = reader.ReadMetaData(paths);
        metaData.Set(fides::keys::BLOCK_SELECTION(), blockSelection);

        for (std::size_t t = 0; t < nTime; t++)
        {
          fides::StepStatus status = reader.PrepareNextStep(paths);
          if (status == fides::StepStatus::EndOfStream)
          {
            retVal = 1;
            break;
          }
          auto pds2 = reader.ReadDataSet(paths, metaData);
          retVal = ValidateDataSets(allPDS1[t][i], pds2, wf);
          if (retVal)
            break;
        }
      }
    }
  }
  return retVal;
}

int main(int argc, char** argv)
{
  int rank = 0, numRanks = 1;
#ifdef FIDES_USE_MPI
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &numRanks);
#endif

  RandomNumbers randomNumbers;
  std::vector<std::string> dataSetTypes = {
    "uniform", "rectilinear", "unstructured_single", "explicit"
  };

  //Some convenient code for debugging.
#if 0
  {
    //std::vector<std::vector<int>> numBlocks = {{1,1}};
    //std::vector<std::vector<int>> numBlocks = {{1}};
    std::vector<std::vector<int>> numBlocks;
    std::vector<int> x;
    for (int i = 0; i < numRanks; i++)
      x.push_back(1);
    numBlocks.push_back(x);

    RunAppendTests(numBlocks, "explicit", rank, numRanks, randomNumbers);
    //RunTests(numBlocks, "uniform", rank, numRanks, randomNumbers);
  }
#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif
  return 0;
#endif

  std::vector<std::vector<int>> numBlocks;
  //First, add some a constant number of blocks per rank, ranging from 1 to 6.
  for (int i = 0; i < 5; i++)
  {
    std::vector<int> nblocks(numRanks, i + 1);
    numBlocks.push_back(nblocks);
  }

  //Add a random number of blocks per rank.
  for (int i = 0; i < 5; i++)
  {
    auto vec = randomNumbers.RandomVec(0, 5, numRanks);
    numBlocks.push_back(vec);
  }

  int retVal = 0;
  for (auto& dsType : dataSetTypes)
  {
    int status = RunTests(numBlocks, dsType, rank, numRanks, randomNumbers);
    if (status)
    {
      retVal = 1;
      // If the test fails here, you can't continue in MPI mode for some reason:
#ifdef FIDES_USE_MPI
      return MPI_Abort(MPI_COMM_WORLD, 1);
#endif
    }
  }

  //test the append writer
  for (auto& dsType : dataSetTypes)
  {
    int status = RunAppendTests(numBlocks, dsType, rank, numRanks, randomNumbers);
    if (status)
    {
      retVal = 1;
      // If the test fails here, you can't continue in MPI mode for some reason:
#ifdef FIDES_USE_MPI
      return MPI_Abort(MPI_COMM_WORLD, 1);
#endif
    }
  }

#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif

  return retVal;
}
