/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkDIYUtilities.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm 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 "vtkDIYUtilities.h"

#include "vtkAbstractArray.h"
#include "vtkArrayDispatch.h"
#include "vtkBoundingBox.h"
#include "vtkCellCenters.h"
#include "vtkCompositeDataIterator.h"
#include "vtkCompositeDataSet.h"
#include "vtkDataArray.h"
#include "vtkDataObjectTypes.h"
#include "vtkFieldData.h"
#include "vtkImageData.h"
#include "vtkImageDataToPointSet.h"
#include "vtkLogger.h"
#include "vtkMultiProcessController.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkRectilinearGrid.h"
#include "vtkRectilinearGridToPointSet.h"
#include "vtkUnstructuredGrid.h"
#include "vtkXMLDataObjectWriter.h"
#include "vtkXMLGenericDataObjectReader.h"

#if VTK_MODULE_ENABLE_VTK_ParallelMPI
#include "vtkMPI.h"
#include "vtkMPICommunicator.h"
#include "vtkMPIController.h"
#endif

#include <cassert>
#include <string>
#include <vector>

namespace
{
//==============================================================================
struct SaveArrayWorker
{
  SaveArrayWorker(diy::BinaryBuffer& bb)
    : BB(bb)
  {
  }

  template <class ArrayT>
  void operator()(ArrayT* array)
  {
    auto data = array->GetPointer(0);
    diy::save(this->BB, data, array->GetNumberOfValues());
  }

  diy::BinaryBuffer& BB;
};

//==============================================================================
struct LoadArrayWorker
{
  LoadArrayWorker(diy::BinaryBuffer& bb)
    : BB(bb)
  {
  }

  template <class ArrayT>
  void operator()(ArrayT* array)
  {
    int numberOfComponents;
    vtkIdType numberOfTuples;
    std::string name;
    diy::load(this->BB, numberOfComponents);
    diy::load(this->BB, numberOfTuples);
    diy::load(this->BB, name);

    array->SetNumberOfComponents(numberOfComponents);
    array->SetNumberOfTuples(numberOfTuples);
    array->SetName(name.c_str());

    auto data = array->GetPointer(0);
    diy::load(this->BB, data, numberOfComponents * numberOfTuples);
  }

  diy::BinaryBuffer& BB;
};
} // anonymous namespace

static unsigned int vtkDIYUtilitiesCleanupCounter = 0;
#if VTK_MODULE_ENABLE_VTK_ParallelMPI
static vtkMPIController* vtkDIYUtilitiesCleanupMPIController = nullptr;
#endif
vtkDIYUtilitiesCleanup::vtkDIYUtilitiesCleanup()
{
  ++vtkDIYUtilitiesCleanupCounter;
}

vtkDIYUtilitiesCleanup::~vtkDIYUtilitiesCleanup()
{
  if (--vtkDIYUtilitiesCleanupCounter == 0)
  {
#if VTK_MODULE_ENABLE_VTK_ParallelMPI
    if (vtkDIYUtilitiesCleanupMPIController)
    {
      vtkLogF(TRACE, "Cleaning up MPI controller created for DIY filters.");
      vtkDIYUtilitiesCleanupMPIController->Finalize();
      vtkDIYUtilitiesCleanupMPIController->Delete();
      vtkDIYUtilitiesCleanupMPIController = nullptr;
    }
#endif
  }
}

//------------------------------------------------------------------------------
vtkDIYUtilities::vtkDIYUtilities() = default;

//------------------------------------------------------------------------------
vtkDIYUtilities::~vtkDIYUtilities() = default;

//------------------------------------------------------------------------------
void vtkDIYUtilities::InitializeEnvironmentForDIY()
{
#if VTK_MODULE_ENABLE_VTK_ParallelMPI
  int mpiOk;
  MPI_Initialized(&mpiOk);
  if (!mpiOk)
  {
    vtkLogF(TRACE,
      "Initializing MPI for DIY filters since process did not do so in an MPI enabled build.");
    assert(vtkDIYUtilitiesCleanupMPIController == nullptr);
    vtkDIYUtilitiesCleanupMPIController = vtkMPIController::New();

    static int argc = 0;
    static char** argv = { nullptr };
    vtkDIYUtilitiesCleanupMPIController->Initialize(&argc, &argv);
  }
#endif
}

//------------------------------------------------------------------------------
diy::mpi::communicator vtkDIYUtilities::GetCommunicator(vtkMultiProcessController* controller)
{
  vtkDIYUtilities::InitializeEnvironmentForDIY();

#if VTK_MODULE_ENABLE_VTK_ParallelMPI
  vtkMPICommunicator* vtkcomm =
    vtkMPICommunicator::SafeDownCast(controller ? controller->GetCommunicator() : nullptr);
  return vtkcomm ? diy::mpi::communicator(*vtkcomm->GetMPIComm()->GetHandle())
                 : diy::mpi::communicator(MPI_COMM_SELF);
#else
  (void)controller;
  return diy::mpi::communicator();
#endif
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::AllReduce(diy::mpi::communicator& comm, vtkBoundingBox& bbox)
{
  if (comm.size() > 1)
  {
    std::vector<double> local_minpoint(3), local_maxpoint(3);
    bbox.GetMinPoint(&local_minpoint[0]);
    bbox.GetMaxPoint(&local_maxpoint[0]);

    std::vector<double> global_minpoint(3), global_maxpoint(3);
    diy::mpi::all_reduce(comm, local_minpoint, global_minpoint, diy::mpi::minimum<float>());
    diy::mpi::all_reduce(comm, local_maxpoint, global_maxpoint, diy::mpi::maximum<float>());

    bbox.SetMinPoint(&global_minpoint[0]);
    bbox.SetMaxPoint(&global_maxpoint[0]);
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Save(diy::BinaryBuffer& bb, vtkDataArray* array)
{
  if (!array)
  {
    diy::save(bb, static_cast<int>(VTK_VOID));
  }
  else
  {
    diy::save(bb, array->GetDataType());
    diy::save(bb, array->GetNumberOfComponents());
    diy::save(bb, array->GetNumberOfTuples());
    if (array->GetName())
    {
      diy::save(bb, std::string(array->GetName()));
    }
    else
    {
      diy::save(bb, std::string(""));
    }

    SaveArrayWorker worker(bb);
    vtkArrayDispatch::Dispatch::Execute(array, worker);
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Save(diy::BinaryBuffer& bb, vtkFieldData* fd)
{
  if (!fd)
  {
    diy::save(bb, 0);
  }
  else
  {
    diy::save(bb, fd->GetNumberOfArrays());
    for (int id = 0; id < fd->GetNumberOfArrays(); ++id)
    {
      vtkDIYUtilities::Save(bb, vtkArrayDownCast<vtkDataArray>(fd->GetAbstractArray(id)));
    }
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Save(diy::BinaryBuffer& bb, vtkDataSet* p)
{
  if (p)
  {
    diy::save(bb, p->GetDataObjectType());
    auto writer = vtkXMLDataObjectWriter::NewWriter(p->GetDataObjectType());
    if (writer)
    {
      writer->WriteToOutputStringOn();
      writer->SetCompressorTypeToLZ4();
      writer->SetEncodeAppendedData(false);
      writer->SetInputDataObject(p);
      writer->Write();
      diy::save(bb, writer->GetOutputString());
      writer->Delete();
    }
    else
    {
      vtkLogF(
        ERROR, "Cannot serialize `%s` yet. Aborting for debugging purposes.", p->GetClassName());
      abort();
    }
  }
  else
  {
    diy::save(bb, static_cast<int>(-1)); // can't be VTK_VOID since VTK_VOID == VTK_POLY_DATA.
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Load(diy::BinaryBuffer& bb, vtkDataArray*& array)
{
  int type;
  diy::load(bb, type);
  if (type == VTK_VOID)
  {
    array = nullptr;
  }
  else
  {
    array = vtkDataArray::SafeDownCast(vtkAbstractArray::CreateArray(type));
    LoadArrayWorker worker(bb);
    vtkArrayDispatch::Dispatch::Execute(array, worker);
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Load(diy::BinaryBuffer& bb, vtkFieldData*& fd)
{
  int numberOfArrays;
  diy::load(bb, numberOfArrays);
  if (!numberOfArrays)
  {
    fd = nullptr;
  }
  else
  {
    vtkNew<vtkFieldData> tmp;
    for (int id = 0; id < numberOfArrays; ++id)
    {
      vtkDataArray* array = nullptr;
      vtkDIYUtilities::Load(bb, array);
      if (array)
      {
        tmp->AddArray(array);
        array->FastDelete();
      }
    }
    tmp->Register(nullptr);
    fd = tmp.GetPointer();
  }
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Load(diy::BinaryBuffer& bb, vtkDataSet*& p)
{
  p = nullptr;
  int type;
  diy::load(bb, type);
  if (type == -1)
  {
    p = nullptr;
  }
  else
  {
    std::string data;
    diy::load(bb, data);

    vtkSmartPointer<vtkDataSet> ds;
    if (auto reader = vtkXMLGenericDataObjectReader::CreateReader(type, /*parallel*/ false))
    {
      reader->ReadFromInputStringOn();
      reader->SetInputString(data);
      reader->Update();
      ds = vtkDataSet::SafeDownCast(reader->GetOutputDataObject(0));
    }
    else
    {
      vtkLogF(ERROR, "Currrently type '%d' (%s) is not supported.", type,
        vtkDataObjectTypes::GetClassNameFromTypeId(type));
      // aborting for debugging purposes.
      abort();
    }

    ds->Register(nullptr);
    p = ds.GetPointer();
  }
}

//------------------------------------------------------------------------------
diy::ContinuousBounds vtkDIYUtilities::Convert(const vtkBoundingBox& bbox)
{
  if (bbox.IsValid())
  {
    diy::ContinuousBounds bds(3);
    bds.min[0] = static_cast<float>(bbox.GetMinPoint()[0]);
    bds.min[1] = static_cast<float>(bbox.GetMinPoint()[1]);
    bds.min[2] = static_cast<float>(bbox.GetMinPoint()[2]);
    bds.max[0] = static_cast<float>(bbox.GetMaxPoint()[0]);
    bds.max[1] = static_cast<float>(bbox.GetMaxPoint()[1]);
    bds.max[2] = static_cast<float>(bbox.GetMaxPoint()[2]);
    return bds;
  }
  return diy::ContinuousBounds(3);
}

//------------------------------------------------------------------------------
vtkBoundingBox vtkDIYUtilities::Convert(const diy::ContinuousBounds& bds)
{
  double bounds[6];
  bounds[0] = bds.min[0];
  bounds[1] = bds.max[0];
  bounds[2] = bds.min[1];
  bounds[3] = bds.max[1];
  bounds[4] = bds.min[2];
  bounds[5] = bds.max[2];
  vtkBoundingBox bbox;
  bbox.SetBounds(bounds);
  return bbox;
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::Broadcast(
  diy::mpi::communicator& comm, std::vector<vtkBoundingBox>& boxes, int source)
{
  // FIXME: can't this be made more elegant?
  std::vector<double> raw_bounds;
  if (comm.rank() == source)
  {
    raw_bounds.resize(6 * boxes.size());
    for (size_t cc = 0; cc < boxes.size(); ++cc)
    {
      boxes[cc].GetBounds(&raw_bounds[6 * cc]);
    }
  }
  diy::mpi::broadcast(comm, raw_bounds, source);
  if (comm.rank() != source)
  {
    boxes.resize(raw_bounds.size() / 6);
    for (size_t cc = 0; cc < boxes.size(); ++cc)
    {
      boxes[cc].SetBounds(&raw_bounds[6 * cc]);
    }
  }
}

//------------------------------------------------------------------------------
std::vector<vtkSmartPointer<vtkPoints>> vtkDIYUtilities::ExtractPoints(
  const std::vector<vtkDataSet*>& datasets, bool use_cell_centers)
{
  vtkNew<vtkCellCenters> cellCenterFilter;
  cellCenterFilter->SetVertexCells(false);
  cellCenterFilter->SetCopyArrays(false);

  vtkNew<vtkRectilinearGridToPointSet> convertorRG;
  vtkNew<vtkImageDataToPointSet> convertorID;

  std::vector<vtkSmartPointer<vtkPoints>> all_points;
  for (auto ds : datasets)
  {
    if (use_cell_centers)
    {
      cellCenterFilter->SetInputDataObject(ds);
      cellCenterFilter->Update();
      ds = cellCenterFilter->GetOutput();
    }
    if (auto ps = vtkPointSet::SafeDownCast(ds))
    {
      all_points.emplace_back(ps->GetPoints());
    }
    else if (auto rg = vtkRectilinearGrid::SafeDownCast(ds))
    {
      convertorRG->SetInputDataObject(rg);
      convertorRG->Update();
      all_points.emplace_back(convertorRG->GetOutput()->GetPoints());
    }
    else if (auto id = vtkImageData::SafeDownCast(ds))
    {
      convertorID->SetInputDataObject(id);
      convertorID->Update();
      all_points.emplace_back(convertorID->GetOutput()->GetPoints());
    }
    else
    {
      // need a placeholder for dataset.
      all_points.emplace_back(nullptr);
    }
  }
  return all_points;
}

//------------------------------------------------------------------------------
vtkBoundingBox vtkDIYUtilities::GetLocalBounds(vtkDataObject* dobj)
{
  double bds[6];
  vtkMath::UninitializeBounds(bds);
  if (auto ds = vtkDataSet::SafeDownCast(dobj))
  {
    ds->GetBounds(bds);
  }
  else if (auto cd = vtkCompositeDataSet::SafeDownCast(dobj))
  {
    cd->GetBounds(bds);
  }
  return vtkBoundingBox(bds);
}

//------------------------------------------------------------------------------
void vtkDIYUtilities::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}
