//============================================================================
//  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/filter/clean_grid/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/VTKDataSetWriter.h>

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

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

// This example has 3d data read using SST. This example will also work using BP4 streaming,
// just comment out the lines setting SST engine (Fides will use BP4 if no other engine specified).
int main(int argc, char** argv)
{
  if (argc < 3)
  {
    std::cerr << "Usage ./xgc-sst <name of the json file> <path of data source folder> "
                 "<optional output method>\n";
    std::cerr << "Example 1: ./xgc-sst ~/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-sst ~/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

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

  std::string dataModelFile = std::string(argv[1]);
  fides::io::DataSetReader reader(dataModelFile);
  std::cout << "XGC Reader started\n";

  std::string data_dir = std::string(argv[2]);
  if (data_dir.back() != '/')
  {
    data_dir += '/';
  }
  // Data we want to read from file
  std::unordered_map<std::string, std::string> paths;
  paths["mesh"] = std::string(data_dir);
  paths["diag"] = std::string(data_dir);

  // Data to read using SST
  paths["3d"] = std::string(data_dir);
  fides::DataSourceParams params;
  params["engine_type"] = "SST";
  reader.SetDataSourceParameters("3d", params);

  size_t step = 0;
  bool metadataRead = false;
  while (true)
  {
    auto status = reader.PrepareNextStep(paths);
    if (status == fides::StepStatus::NotReady)
    {
      continue;
    }
    else if (status == fides::StepStatus::EndOfStream)
    {
      break;
    }
    // When using SST, a BeginStep call must happen before any attempt to read data.
    // For XGC, ReadMetaData requires reading the nphi variable in order to determine
    // the number of blocks, this means for XGC ReadMetaData must be called after
    // PrepareNextStep.
    fides::metadata::MetaData metaData;
    if (!metadataRead)
    {
      metaData = reader.ReadMetaData(paths);
      metadataRead = true;
    }
    fides::metadata::MetaData selections;

    auto output = reader.ReadDataSet(paths, selections);

    if (output.GetNumberOfPartitions() == 0)
    {
      continue;
    }

    vtkm::filter::clean_grid::CleanGrid clean;
    auto outputDataPDS = clean.Execute(output);
    if (outputVTK)
    {
      auto& outputData = outputDataPDS.GetPartition(0);
      std::string filename = "xgc-step-" + std::to_string(step) + ".vtk";
      std::cout << "writing output in vtk format to " << filename << std::endl;
      vtkm::io::VTKDataSetWriter writer(filename);
      writer.WriteDataSet(outputData);
    }
#ifdef USE_VTKM_RENDERING
    else
    {
      // 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);
        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-step-" + std::to_string(step) + ".png";
      std::cout << "Rendering image " << filename << std::endl;
      view.SaveAs(filename);
    }
    step++;
#endif
  }

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