//============================================================================
//  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 <vtkm/cont/ArrayHandleXGCCoordinates.h>
#include <vtkm/cont/CellSetExtrude.h>

#include <vtkm/io/VTKDataSetWriter.h>

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

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

static void CheckDataSet(const vtkm::cont::DataSet& ds, int addPlanes = 0)
{
  auto& cellSet = ds.GetCellSet();
  if (!cellSet.IsType<vtkm::cont::CellSetExtrude>())
  {
    throw std::runtime_error("CellSet for XGC should be type CellSetExtrude");
  }
  auto csXGC = cellSet.AsCellSet<vtkm::cont::CellSetExtrude>();
  if (!csXGC.GetIsPeriodic())
  {
    throw std::runtime_error("The json file specifies that this dataset should be periodic.");
  }

  vtkm::Id numExpectedPlanes = 4 * (1 + addPlanes);

  auto numPlanes = csXGC.GetNumberOfPlanes();
  if (numPlanes != numExpectedPlanes)
  {
    throw std::runtime_error("Wrong number of planes for XGC data.");
  }

  auto numPointsPerPlane = csXGC.GetNumberOfPointsPerPlane();
  if (numPointsPerPlane != 20694)
  {
    throw std::runtime_error("There should be 20,694 points per plane in this dataset.");
  }
  if (numPointsPerPlane * numPlanes != cellSet.GetNumberOfPoints())
  {
    throw std::runtime_error("The total number of points in the cell set is incorrect");
  }

  // CoordinateSystem checks
  const auto& coordSys = ds.GetCoordinateSystem(0);
  if (ds.GetNumberOfCoordinateSystems() < 1)
  {
    throw std::runtime_error("XGC should have a coordinate system");
  }
  if (numPointsPerPlane * numPlanes != coordSys.GetNumberOfPoints())
  {
    throw std::runtime_error("The total number of points in the coordinate system is incorrect");
  }
  const auto& coords =
    coordSys.GetData().AsArrayHandle<vtkm::cont::ArrayHandleXGCCoordinates<double>>();
  const auto coordsPortal = coords.ReadPortal();
  for (vtkm::Id i = 0; i < coordsPortal.GetNumberOfValues(); ++i)
  {
    auto v = coordsPortal.Get(i);

    if (v[2] < -0.63904 || v[2] > 2.33904)
    {
      throw std::runtime_error(
        "z-coordinate out of bounds specified by the rz field in xgc.mesh.bp");
    }
    for (vtkm::Id j = 0; j < 2; ++j)
    {
      if (!std::isfinite(v[j]))
      {
        throw std::runtime_error(
          "XGC coordinates should not be nan or inf. "
          "Check that make_ArrayHandleXGCCoordinates is being passed correct values.");
      }
    }
  }
}

template <typename T>
static void CheckField(const vtkm::cont::DataSet& ds,
                       const std::string& vname,
                       const vtkm::Range& minRange,
                       const vtkm::Range& maxRange)
{
  if (!ds.HasField(vname))
  {
    throw std::runtime_error("Field not found: " + vname);
  }

  vtkm::cont::ArrayHandle<vtkm::Range> rangeArray = ds.GetField(vname).GetRange();
  auto vRange = rangeArray.ReadPortal().Get(0);
  if (!minRange.Contains(vRange.Min) || !maxRange.Contains(vRange.Max))
  {
    throw std::runtime_error("Range for " + vname + " is incorrect");
  }
}

template <typename T>
static void CheckField(const vtkm::cont::DataSet& ds, const std::string& vname)
{
  vtkm::Range r(-1e20, 1e20);
  CheckField<T>(ds, vname, r, r);
}


int main(int argc, char** argv)
{
  if (argc < 3)
  {
    std::cerr
      << "./xgc <name of the json file> <path of data source folder> <optional write data>\n";
    return 0;
  }

  bool writeVTK = false;
  if (argc >= 4)
    writeVTK = std::atoi(argv[3]) == 1;

#ifdef FIDES_USE_MPI
  MPI_Init(&argc, &argv);
#endif

  std::vector<bool> boolVals = { true, false };
  for (const auto& isPeriodic : boolVals)
  {
    for (int addPlanes = 0; addPlanes < 3; addPlanes++)
    {
      std::string dataModelFile = std::string(argv[1]);
      fides::io::DataSetReader reader(dataModelFile);

      std::unordered_map<std::string, std::string> paths;
      paths["mesh"] = std::string(argv[2]) + "/";
      paths["3d"] = std::string(argv[2]) + "/";
      paths["diag"] = std::string(argv[2]) + "/";
      paths["units"] = std::string(argv[2]) + "/";

      auto metaData = reader.ReadMetaData(paths);

      fides::metadata::MetaData selections;
      fides::metadata::Bool periodic(isPeriodic);
      selections.Set(fides::keys::fusion::FUSION_PERIODIC_CELLSET(), periodic);
      if (addPlanes > 0)
      {
        fides::metadata::Size numAddPlanes(addPlanes);
        selections.Set(fides::keys::fusion::PLANE_INSERTION(), numAddPlanes);
      }

      fides::metadata::Index idx(2);
      selections.Set(fides::keys::STEP_SELECTION(), idx);

      fides::metadata::Bool addR(true), addPhi(true), addPsi(true);
      selections.Set(fides::keys::fusion::ADD_R_FIELD(), addR);
      selections.Set(fides::keys::fusion::ADD_PHI_FIELD(), addPhi);
      selections.Set(fides::keys::fusion::ADD_PSI_FIELD(), addPsi);

      vtkm::cont::PartitionedDataSet output = reader.ReadDataSet(paths, selections);
      if (output.GetNumberOfPartitions() == 0)
      {
        throw std::runtime_error("No output from XGC reader");
      }

      const auto& ds = output.GetPartition(0);

      CheckDataSet(ds, addPlanes);

      CheckField<vtkm::Float64>(
        ds, "pot0", vtkm::Range(1.67863, 1.67865), vtkm::Range(90.7623, 90.7625));
      CheckField<vtkm::Float64>(
        ds, "dpot", vtkm::Range(-15.1787, -15.1785), vtkm::Range(15.5977, 15.5979));

      CheckField<vtkm::Float64>(ds, "R", vtkm::Range(1.060, 1.061), vtkm::Range(2.339, 2.34));
      CheckField<vtkm::Float64>(
        ds, "Psi", vtkm::Range(-.001, 0.001), vtkm::Range(0.02463, 0.02464));

      vtkm::Float64 N = static_cast<vtkm::Float64>((4 * (1 + addPlanes)));
      vtkm::Float64 phiMax = vtkm::TwoPi() * (N - 1) / N;
      CheckField<vtkm::Float64>(
        ds, "Phi", vtkm::Range(-.001, 0.001), vtkm::Range(phiMax - 0.01, phiMax + 0.01));

      if (writeVTK)
      {
        std::string fname = "xgc-output.vtk";
        if (addPlanes > 0)
          fname = "xgc-output-" + std::to_string(addPlanes) + ".vtk";
        std::cout << "writing output in vtk format to " << fname << std::endl;
        vtkm::io::VTKDataSetWriter writer(fname);
        writer.WriteDataSet(ds);
      }
    }
  }

  //Do a time varying test.
  std::string dataModelFile = std::string(argv[1]);
  fides::io::DataSetReader reader(dataModelFile);

  std::unordered_map<std::string, std::string> paths;
  paths["mesh"] = std::string(argv[2]) + "/";
  paths["3d"] = std::string(argv[2]) + "/";
  paths["diag"] = std::string(argv[2]) + "/";
  paths["units"] = std::string(argv[2]) + "/";

  auto metaData = reader.ReadMetaData(paths);

  bool val = true;
  fides::metadata::Bool addR(val), addPhi(val), addPsi(val);
  metaData.Set(fides::keys::fusion::ADD_R_FIELD(), addR);
  metaData.Set(fides::keys::fusion::ADD_PHI_FIELD(), addPhi);
  metaData.Set(fides::keys::fusion::ADD_PSI_FIELD(), addPsi);

  int stepCount = 0;
  while (true)
  {
    std::cout << "Step= " << stepCount << std::endl;
    fides::StepStatus status = reader.PrepareNextStep(paths);
    if (status == fides::StepStatus::EndOfStream)
    {
      break;
    }

    auto output = reader.ReadDataSet(paths, metaData);
    if (output.GetNumberOfPartitions() != 1)
    {
      std::cerr << "Error: expected 1 output block, got " << output.GetNumberOfPartitions()
                << std::endl;
    }
    const auto& ds = output.GetPartition(0);
    CheckDataSet(ds);

    CheckField<vtkm::Float64>(ds, "pot0");
    CheckField<vtkm::Float64>(ds, "dpot");

    CheckField<vtkm::Float64>(ds, "R", vtkm::Range(1.060, 1.061), vtkm::Range(2.339, 2.34));
    CheckField<vtkm::Float64>(ds, "Psi", vtkm::Range(-.001, 0.001), vtkm::Range(0.02463, 0.02464));

    vtkm::Float64 phiMax = vtkm::TwoPi() * 0.75;
    CheckField<vtkm::Float64>(
      ds, "Phi", vtkm::Range(-.001, 0.001), vtkm::Range(phiMax - 0.01, phiMax + 0.01));

    stepCount++;
  }


#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif
  return 0;
}
