//============================================================================
//  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/CellSetExtrude.h>
#include <vtkm/cont/ColorTable.h>
#include <vtkm/cont/Initialize.h>
#include <vtkm/filter/CleanGrid.h>
#ifdef USE_VTKM_RENDERING
#include <vtkm/rendering/Camera.h>
#include <vtkm/rendering/CanvasRayTracer.h>
#include <vtkm/rendering/MapperPoint.h>
#include <vtkm/rendering/Scene.h>
#include <vtkm/rendering/View3D.h>
#endif
#include <vtkm/io/writer/VTKDataSetWriter.h>

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

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

int main(int argc, char** argv)
{
  if (argc < 3)
  {
    // output method is 1 for writing VTK files (default), 0 to use VTK-m rendering
    std::cerr << "Usage ./xgc <name of the json file> <path of data source folder> "
                 "<optional output method>\n";
    std::cerr << "Example 1: ./xgc ~/fides/examples/xgc/xgc-reader.json ~/fides/tests/data 0\n"
                 "\tThis outputs the data in .png format using VTK-m rendering.\n";
    std::cerr << "Example 2: ./xgc ~/fides/examples/xgc/xgc-reader.json ~/fides/tests/data 1\n"
                 "\tThis outputs the data in .vtk format.\n";
    return 1;
  }

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

  int retVal = 0;
  bool outputVTK = true;
  if (argc == 4)
  {
    outputVTK = std::atoi(argv[3]) == 1;
  }

  std::string dataModelFile = std::string(argv[1]);
  std::cout << "data-model: " << dataModelFile << "\n";

  fides::io::DataSetReader reader(dataModelFile);

  std::unordered_map<std::string, std::string> paths;

  // This would be so much easier with C++17 filesystem!
  std::string data_dir = std::string(argv[2]);
  if (data_dir.back() != '/')
  {
    data_dir += '/';
  }
  paths["mesh"] = data_dir;
  paths["3d"] = data_dir;
  paths["diag"] = data_dir;

  auto metaData = reader.ReadMetaData(paths);
  auto& nBlocks = metaData.Get<fides::metadata::Size>(fides::keys::NUMBER_OF_BLOCKS());
  std::cout << "num blocks " << nBlocks.NumberOfItems << std::endl;

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

  vtkm::cont::PartitionedDataSet output = reader.ReadDataSet(paths, selections);

  if (output.GetNumberOfPartitions() == 0)
  {
#ifdef FIDES_USE_MPI
    MPI_Finalize();
#endif
    return 0;
  }

  vtkm::filter::clean_grid::CleanGrid clean;
  auto outputDataPDS = clean.Execute(output);
  if (outputVTK)
  {
    auto& outputData = outputDataPDS.GetPartition(0);
    std::cout << "writing output in vtk format...\n";
    vtkm::io::VTKDataSetWriter writer("xgc-output.vtk");
    writer.WriteDataSet(outputData);
  }
#ifdef USE_VTKM_RENDERING
  else
  {
    std::cout << "rendering image using vtk-m...\n";
    // Create a mapper, canvas and view that will be used to render the scene
    vtkm::rendering::Scene scene;
    vtkm::rendering::MapperPoint mapper;
    vtkm::rendering::CanvasRayTracer canvas(1024, 1024);
    vtkm::rendering::Color bg(0.2f, 0.2f, 0.2f, 1.0f);
    vtkm::cont::ColorTable colorTable(vtkm::cont::ColorTable::Preset::CoolToWarm);

    // Render an image of the output
    std::cout << "Rendering image\n";
    for (vtkm::Id i = 0; i < outputDataPDS.GetNumberOfPartitions(); ++i)
    {
      auto& outputData = outputDataPDS.GetPartition(i);
      std::cout << "Adding partition " << i << std::endl;
      scene.AddActor(vtkm::rendering::Actor(outputData.GetCellSet(),
                                            outputData.GetCoordinateSystem(),
                                            outputData.GetField("dpot"),
                                            colorTable));
    }
    vtkm::rendering::View3D view(scene, mapper, canvas);
    auto& camera = view.GetCamera();
    camera.SetViewUp(vtkm::Vec3f_64{ 0.5, -0.2, 0.8 });
    camera.SetPosition(vtkm::Vec3f_32{ -9, 5, 7 });
    vtkm::Float32 zoom = 0.5;
    camera.Zoom(zoom);
    view.Paint();
    std::string filename = "xgc.png";
    view.SaveAs(filename);
  }
#endif
#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif
  return retVal;
}
