//============================================================================
//  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>
#ifdef FIDES_USE_MPI
#include <mpi.h>
#endif
#include <string>
#include <unordered_map>

#include <vtkm/cont/ArrayRangeCompute.h>
#include <vtkm/io/VTKDataSetWriter.h>

int main(int argc, char** argv)
{
  if (argc < 3)
  {
    std::cerr << "Usage ./test-gtc <name of the json file> <path of data source folder>"
                 "<optional write data>\n";
    std::cerr << "Example: ./test-gtc ~/fides/tests/test-gtc.json ~/fides/tests/data 1\n"
                 "\tThis outputs the data in .vtk format.\n";
    return 0;
  }

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

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

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

      std::cout << "Reading: " << argv[1] << std::endl;
      std::cout << "ap= " << addPlanes << " per= " << isPeriodic << std::endl;
      auto metaData = reader.ReadMetaData();

      fides::metadata::Bool periodic(isPeriodic);
      metaData.Set(fides::keys::fusion::FUSION_PERIODIC_CELLSET(), periodic);

      if (addPlanes > 0)
      {
        fides::metadata::Size numAddPlanes(addPlanes);
        metaData.Set(fides::keys::fusion::PLANE_INSERTION(), numAddPlanes);
      }

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

      int stepCount = 0;
      while (true)
      {
        std::cout << "Step= " << stepCount << std::endl;

        if (stepCount > 2)
        {
          std::cerr << "Should only be 2 steps in GTC file." << std::endl;
          retVal = 1;
          break;
        }

        fides::StepStatus status = reader.PrepareNextStep();
        if (status == fides::StepStatus::EndOfStream)
        {
          break;
        }

        auto output = reader.ReadDataSet(metaData);
        if (output.GetNumberOfPartitions() != 1)
        {
          std::cerr << "Error: expected 1 output block, got " << output.GetNumberOfPartitions()
                    << std::endl;
          retVal = 1;
        }
        vtkm::cont::DataSet ds = output.GetPartition(0);
        vtkm::cont::UnknownCellSet cellSet = ds.GetCellSet();

        if (!cellSet.IsType<vtkm::cont::CellSetSingleType<>>())
        {
          std::cerr << "Wrong cell type. Expected CellSetSingleType." << std::endl;
          retVal = 1;
        }

        vtkm::cont::CellSetSingleType<> css = cellSet.AsCellSet<vtkm::cont::CellSetSingleType<>>();
        static constexpr vtkm::Id numTrisPerPlane = 3868;
        static constexpr vtkm::Id numPlanes = 8;
        vtkm::Id totNumPlanes = numPlanes * (1 + addPlanes);
        //If *not* periodic, we don't add an extra set of cells.
        if (!isPeriodic)
          totNumPlanes--;
        vtkm::Id numExpectedCells = numTrisPerPlane * totNumPlanes;

        if (css.GetNumberOfCells() != numExpectedCells)
        {
          std::cerr << "expected= " << numExpectedCells << std::endl;
          std::cerr << "Wrong number of cells: " << css.GetNumberOfCells() << std::endl;
          std::cerr << numExpectedCells << std::endl;
          retVal = 1;
        }
        if (css.GetCellShapeAsId() != vtkm::CellShapeTagWedge::Id)
        {
          std::cerr << "Wrong cell type." << std::endl;
          retVal = 1;
        }

        vtkm::Id nPts = ds.GetNumberOfPoints();
        vtkm::Id numExpectedPoints = 17160 * (1 + addPlanes);
        if (nPts != numExpectedPoints)
        {
          std::cerr << "Wrong number of points: " << nPts << " expected: " << numExpectedPoints
                    << std::endl;
          retVal = 1;
        }

        //Make sure we got what we asked for.
        std::vector<std::string> reqFields = { "phi", "b", "R", "Phi" };
        for (const auto& f : reqFields)
        {
          if (!ds.HasPointField(f))
          {
            std::cerr << "Missing field: " << f << std::endl;
            retVal = 1;
          }
        }

        //Check that all fields are the right size.
        for (vtkm::Id i = 0; i < ds.GetNumberOfFields(); i++)
        {
          const auto& field = ds.GetField(i);
          if (field.IsPointField())
          {
            if (field.GetNumberOfValues() != numExpectedPoints)
            {
              std::cerr << "Wrong number of values in " << field.GetName() << std::endl;
              retVal = 1;
            }
          }
          else if (field.IsCellField())
          {
            if (field.GetNumberOfValues() != numExpectedCells)
            {
              std::cerr << "Wrong number of values in " << field.GetName() << std::endl;
              retVal = 1;
            }
          }

          if (field.GetName() == "phi")
          {
            const auto& phiHandle = field.GetData().AsArrayHandle<vtkm::cont::ArrayHandle<float>>();
            auto rangeArray = vtkm::cont::ArrayRangeCompute(phiHandle);
            auto rangePortal = rangeArray.ReadPortal();

            if (!(rangePortal.Get(0).Min > -2.0e-11 || rangePortal.Get(0).Max < 1.59e-11))
            {
              std::cerr << "Wrong range for phi variable." << std::endl;
              retVal = 1;
            }
          }
          else if (field.GetName() == "b")
          {
            const auto& bHandle = field.GetData().AsArrayHandle<vtkm::cont::ArrayHandle<float>>();
            auto rangeArray = vtkm::cont::ArrayRangeCompute(bHandle);
            auto rangePortal = rangeArray.ReadPortal();

            if (!(rangePortal.Get(0).Min > 0.762 || rangePortal.Get(0).Max < 1.485))
            {
              std::cerr << "Wrong range for b variable." << std::endl;
              retVal = 1;
            }
          }
        }

        if (writeVTK)
        {
          std::string fname = "gtc-output.vtk";
          if (addPlanes > 0)
            fname = "gtc-output-" + std::to_string(addPlanes) + ".step" +
              std::to_string(stepCount) + ".vtk";

          std::cout << "writing output in vtk format to " << fname << std::endl;
          vtkm::io::VTKDataSetWriter writer(fname);
          writer.WriteDataSet(ds);
        }
        stepCount++;
      }

      if (stepCount != 2)
      {
        std::cerr << "Should be 2 steps in GTC file." << std::endl;
        retVal = 1;
        break;
      }
    }
  }

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