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

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

#include <vtkm/Version.h>
#include <vtkm/cont/Algorithm.h>
#include <vtkm/cont/ArrayRangeCompute.h>
#include <vtkm/worklet/DispatcherMapTopology.h>
#include <vtkm/worklet/ScatterPermutation.h>
#include <vtkm/worklet/WorkletMapTopology.h>

static const int num_steps = 10;
static const int size = 5;

void analysis(fides::io::DataSetReader& reader, int min, int max)
{
  reader.PrepareNextStep();
  auto metaData = reader.ReadMetaData();
  using FieldInfoType = fides::metadata::Vector<fides::metadata::FieldInformation>;
  auto& fields = metaData.Get<FieldInfoType>(fides::keys::FIELDS());
  if (fields.Data.size() != 1)
  {
    std::stringstream ss;
    ss << "Error: expected 1 field, got " << fields.Data.size();
    throw std::runtime_error(ss.str());
  }

  fides::metadata::MetaData selections;
  fides::metadata::Vector<size_t> blockSelection;
  blockSelection.Data.push_back(0);
  selections.Set(fides::keys::BLOCK_SELECTION(), blockSelection);

  FieldInfoType fieldSelection;
  fieldSelection.Data.push_back(
    fides::metadata::FieldInformation("testvar", vtkm::cont::Field::Association::Points));
  selections.Set(fides::keys::FIELDS(), fieldSelection);

  vtkm::cont::PartitionedDataSet output = reader.ReadDataSet(selections);
  if (output.GetNumberOfPartitions() != 1)
  {
    std::stringstream ss;
    ss << "Error: expected 1 block, got " << output.GetNumberOfPartitions();
    throw std::runtime_error(ss.str());
  }

  vtkm::cont::DataSet ds = output.GetPartition(0);
  vtkm::IdComponent nFields = ds.GetNumberOfFields();

  // Two fields: One for the output; another one for the coordinates.
  if (nFields != 2)
  {
    std::stringstream ss;
    ss << "Error: expected 1 output array, got " << nFields;
    throw std::runtime_error(ss.str());
  }

  if (!ds.HasField("testvar", vtkm::cont::Field::Association::Points))
  {
    throw std::runtime_error("Error: expected testvar array, but not found");
  }

  const auto& arrHandle =
    ds.GetField("testvar").GetData().AsArrayHandle<vtkm::cont::ArrayHandle<int>>();
  vtkm::cont::ArrayHandle<vtkm::Range> rangeArray = vtkm::cont::ArrayRangeCompute(arrHandle);
  auto rangePortal = rangeArray.ReadPortal();
  if (rangePortal.Get(0).Min != min)
  {
    std::stringstream ss;
    ss << "Error: expected min = " << min << ", but got " << rangePortal.Get(0).Min;
    throw std::runtime_error(ss.str());
  }

  if (rangePortal.Get(0).Max != max)
  {
    std::stringstream ss;
    ss << "Error: expected max = " << max << ", but got " << rangePortal.Get(0).Max;
    throw std::runtime_error(ss.str());
  }
}

int main(int argc, char** argv)
{
#ifdef FIDES_USE_MPI
  MPI_Init(&argc, &argv);
  adios2::ADIOS adios(MPI_COMM_WORLD);
#else
  adios2::ADIOS adios;
#endif

  const std::string writer_id = "output.bp";
  const std::string source_name = "the_source";

  // First set up ADIOS and Fides
  adios2::IO io = adios.DeclareIO("DataWriter");
  //params["verbose"] = "5"; // optional; sets verbose in ADIOS
  fides::io::DataSetReader reader(argv[1]);
  fides::DataSourceParams params;
  params["engine_type"] = "Inline";
  reader.SetDataSourceParameters(source_name, std::move(params));
  reader.SetDataSourceIO(source_name, &io);

  adios2::Variable<int> testVar =
    io.DefineVariable<int>("testvar", { size, size, size }, { 0, 0, 0 }, { size, size, size });

  // in the application if you call io.Open before setting up Fides,
  // you'll have to call io.SetEngine("Inline") as well.
  // If Fides is set up before the io.Open() call, Fides handles it.
  adios2::Engine writer = io.Open(writer_id, adios2::Mode::Write);

  std::vector<int> data(size * size * size, 0);
  int min = INT_MAX, max = 0;
  for (int ii = 0; ii < num_steps; ++ii)
  {
    for (size_t idx = 0; idx < data.size(); ++idx)
    {
      data[idx] += idx;
      if (data[idx] < min)
        min = data[idx];
      if (data[idx] > max)
        max = data[idx];
    }

    writer.BeginStep();
    writer.Put<int>(testVar, data.data());
    writer.EndStep();

    analysis(reader, min, max);
  }

  writer.Close();
#ifdef FIDES_USE_MPI
  MPI_Finalize();
#endif
  return 0;
}
