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

  Program:   Visualization Toolkit
  Module:    vtkXdmf3HeavyDataHandler.cxx
  Language:  C++

  Copyright (c) 1993-2002 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 "vtkXdmf3HeavyDataHandler.h"

#include "vtkCompositeDataSet.h"
#include "vtkDataObject.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkMutableDirectedGraph.h"
#include "vtkRectilinearGrid.h"
#include "vtkStructuredGrid.h"
#include "vtkUniformGrid.h"
#include "vtkUnstructuredGrid.h"
#include "vtkXdmf3ArrayKeeper.h"
#include "vtkXdmf3ArraySelection.h"
#include "vtkXdmf3DataSet.h"

// clang-format off
#include "vtk_xdmf3.h"
#include VTKXDMF3_HEADER(XdmfCurvilinearGrid.hpp)
#include VTKXDMF3_HEADER(XdmfDomain.hpp)
#include VTKXDMF3_HEADER(XdmfGraph.hpp)
#include VTKXDMF3_HEADER(XdmfGrid.hpp)
#include VTKXDMF3_HEADER(XdmfGridCollection.hpp)
#include VTKXDMF3_HEADER(XdmfGridCollectionType.hpp)
#include VTKXDMF3_HEADER(core/XdmfItem.hpp)
#include VTKXDMF3_HEADER(XdmfRectilinearGrid.hpp)
#include VTKXDMF3_HEADER(XdmfRegularGrid.hpp)
#include VTKXDMF3_HEADER(XdmfSet.hpp)
#include VTKXDMF3_HEADER(XdmfUnstructuredGrid.hpp)
// clang-format on

#include <cassert>

//------------------------------------------------------------------------------
shared_ptr<vtkXdmf3HeavyDataHandler> vtkXdmf3HeavyDataHandler::New(vtkXdmf3ArraySelection* fs,
  vtkXdmf3ArraySelection* cs, vtkXdmf3ArraySelection* ps, vtkXdmf3ArraySelection* gc,
  vtkXdmf3ArraySelection* sc, unsigned int processor, unsigned int nprocessors, bool dt, double t,
  vtkXdmf3ArrayKeeper* keeper, bool asTime)
{
  shared_ptr<vtkXdmf3HeavyDataHandler> p(new vtkXdmf3HeavyDataHandler());
  p->FieldArrays = fs;
  p->CellArrays = cs;
  p->PointArrays = ps;
  p->GridsCache = gc;
  p->SetsCache = sc;
  p->Rank = processor;
  p->NumProcs = nprocessors;
  p->doTime = dt;
  p->time = t;
  p->Keeper = keeper;
  p->AsTime = asTime;
  return p;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::Populate(
  shared_ptr<XdmfDomain> item, vtkDataObject* toFill)
{
  assert(toFill);

  shared_ptr<XdmfDomain> group = shared_dynamic_cast<XdmfDomain>(item);

  shared_ptr<XdmfGridCollection> asGC = shared_dynamic_cast<XdmfGridCollection>(item);
  bool isDomain = !asGC;
  bool isTemporal = false;
  if (asGC)
  {
    if (asGC->getType() == XdmfGridCollectionType::Temporal())
    {
      isTemporal = true;
    }
  }

  // ignore groups that are not in timestep we were asked for
  // but be sure to return everything within them
  bool lastTime = this->doTime;
  if (this->doTime && !(isDomain || isTemporal))
  {
    if (asGC->getTime())
    {
      if (asGC->getTime()->getValue() != this->time)
      {
        // don't return MB that doesn't match the requested time
        return nullptr;
      }
      // inside a match, make sure we get everything underneath
      this->doTime = false;
    }
  }

  vtkMultiBlockDataSet* topB = vtkMultiBlockDataSet::SafeDownCast(toFill);
  vtkDataObject* result;
  unsigned int cnt = 0;
  unsigned int nGridCollections = group->getNumberGridCollections();

  for (unsigned int i = 0; i < nGridCollections; i++)
  {
    if (!this->AsTime)
    {
      if (isDomain && !this->ShouldRead(i, nGridCollections))
      {
        topB->SetBlock(cnt++, nullptr);
        continue;
      }
      vtkMultiBlockDataSet* child = vtkMultiBlockDataSet::New();
      // result = this->Populate(group->getGridCollection(i), child);
      shared_ptr<XdmfDomain> tempDomain =
        shared_dynamic_cast<XdmfDomain>(group->getGridCollection(i));
      result = this->Populate(tempDomain, child);
      topB->SetBlock(cnt++, result);
      child->Delete();
    }
    else
    {
      vtkMultiBlockDataSet* child = vtkMultiBlockDataSet::New();
      // result = this->Populate(group->getGridCollection(i), child);
      shared_ptr<XdmfDomain> tempDomain =
        shared_dynamic_cast<XdmfDomain>(group->getGridCollection(i));
      result = this->Populate(tempDomain, child);
      if (result)
      {
        topB->SetBlock(cnt++, result);
      }
      child->Delete();
    }
  }
  unsigned int nUnstructuredGrids = group->getNumberUnstructuredGrids();
  for (unsigned int i = 0; i < nUnstructuredGrids; i++)
  {
    if (this->AsTime && !isTemporal && !this->ShouldRead(i, nUnstructuredGrids))
    {
      topB->SetBlock(cnt++, nullptr);
      continue;
    }
    shared_ptr<XdmfUnstructuredGrid> cGrid = group->getUnstructuredGrid(i);
    unsigned int nSets = cGrid->getNumberSets();
    vtkDataObject* child;
    if (nSets > 0)
    {
      child = vtkMultiBlockDataSet::New();
    }
    else
    {
      child = vtkUnstructuredGrid::New();
    }
    result = this->Populate(group->getUnstructuredGrid(i), child);
    if (result)
    {
      topB->SetBlock(cnt, result);
      topB->GetMetaData(cnt)->Set(vtkCompositeDataSet::NAME(), cGrid->getName().c_str());
      cnt++;
    }
    child->Delete();
  }
  unsigned int nRectilinearGrids = group->getNumberRectilinearGrids();
  for (unsigned int i = 0; i < nRectilinearGrids; i++)
  {
    if (this->AsTime && !isTemporal && !this->ShouldRead(i, nRectilinearGrids))
    {
      topB->SetBlock(cnt++, nullptr);
      continue;
    }
    shared_ptr<XdmfRectilinearGrid> cGrid = group->getRectilinearGrid(i);
    unsigned int nSets = cGrid->getNumberSets();
    vtkDataObject* child;
    if (nSets > 0)
    {
      child = vtkMultiBlockDataSet::New();
    }
    else
    {
      child = vtkRectilinearGrid::New();
    }
    result = this->Populate(cGrid, child);
    if (result)
    {
      topB->SetBlock(cnt, result);
      topB->GetMetaData(cnt)->Set(vtkCompositeDataSet::NAME(), cGrid->getName().c_str());
      cnt++;
    }
    child->Delete();
  }
  unsigned int nCurvilinearGrids = group->getNumberCurvilinearGrids();
  for (unsigned int i = 0; i < nCurvilinearGrids; i++)
  {
    if (this->AsTime && !isTemporal && !this->ShouldRead(i, nCurvilinearGrids))
    {
      topB->SetBlock(cnt++, nullptr);
      continue;
    }
    shared_ptr<XdmfCurvilinearGrid> cGrid = group->getCurvilinearGrid(i);
    unsigned int nSets = cGrid->getNumberSets();
    vtkDataObject* child;
    if (nSets > 0)
    {
      child = vtkMultiBlockDataSet::New();
    }
    else
    {
      child = vtkStructuredGrid::New();
    }
    result = this->Populate(cGrid, child);
    if (result)
    {
      topB->SetBlock(cnt, result);
      topB->GetMetaData(cnt)->Set(vtkCompositeDataSet::NAME(), cGrid->getName().c_str());
      cnt++;
    }
    child->Delete();
  }
  unsigned int nRegularGrids = group->getNumberRegularGrids();
  for (unsigned int i = 0; i < nRegularGrids; i++)
  {
    if (this->AsTime && !isTemporal && !this->ShouldRead(i, nRegularGrids))
    {
      topB->SetBlock(cnt++, nullptr);
      continue;
    }
    shared_ptr<XdmfRegularGrid> cGrid = group->getRegularGrid(i);
    unsigned int nSets = cGrid->getNumberSets();
    vtkDataObject* child;
    if (nSets > 0)
    {
      child = vtkMultiBlockDataSet::New();
    }
    else
    {
      child = vtkUniformGrid::New();
    }
    result = this->Populate(cGrid, child);
    if (result)
    {
      topB->SetBlock(cnt, result);
      topB->GetMetaData(cnt)->Set(vtkCompositeDataSet::NAME(), cGrid->getName().c_str());
      cnt++;
    }
    child->Delete();
  }
  unsigned int nGraphs = group->getNumberGraphs();
  for (unsigned int i = 0; i < nGraphs; i++)
  {
    if (this->AsTime && !isTemporal && !this->ShouldRead(i, nGraphs))
    {
      topB->SetBlock(cnt++, nullptr);
      continue;
    }
    vtkMutableDirectedGraph* child = vtkMutableDirectedGraph::New();
    result = this->Populate(group->getGraph(i), child);
    if (result)
    {
      topB->SetBlock(cnt, result);
      topB->GetMetaData(cnt)->Set(
        vtkCompositeDataSet::NAME(), group->getGraph(i)->getName().c_str());
      cnt++;
    }
    child->Delete();
  }

  if (lastTime)
  {
    // restore time search now that we've done the group contents
    this->doTime = true;
  }

  if (isTemporal && topB->GetNumberOfBlocks() == 1)
  {
    // temporal collection is just a place holder for its content
    return topB->GetBlock(0);
  }

  return topB;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::Populate(shared_ptr<XdmfGrid> item, vtkDataObject* toFill)
{
  assert(toFill);

  shared_ptr<XdmfUnstructuredGrid> unsGrid = shared_dynamic_cast<XdmfUnstructuredGrid>(item);
  if (unsGrid)
  {
    unsigned int nSets = unsGrid->getNumberSets();
    if (nSets > 0)
    {
      vtkMultiBlockDataSet* mbds = vtkMultiBlockDataSet::SafeDownCast(toFill);
      vtkUnstructuredGrid* child = vtkUnstructuredGrid::New();
      mbds->SetBlock(0, this->MakeUnsGrid(unsGrid, child, this->Keeper));
      mbds->GetMetaData((unsigned int)0)
        ->Set(vtkCompositeDataSet::NAME(), unsGrid->getName().c_str());
      for (unsigned int i = 0; i < nSets; i++)
      {
        vtkUnstructuredGrid* sub = vtkUnstructuredGrid::New();
        mbds->SetBlock(i + 1, this->ExtractSet(i, unsGrid, child, sub, this->Keeper));
        mbds->GetMetaData(i + 1)->Set(
          vtkCompositeDataSet::NAME(), unsGrid->getSet(i)->getName().c_str());
        sub->Delete();
      }
      child->Delete();
      return mbds;
    }
    return this->MakeUnsGrid(unsGrid, vtkUnstructuredGrid::SafeDownCast(toFill), this->Keeper);
  }

  shared_ptr<XdmfRectilinearGrid> recGrid = shared_dynamic_cast<XdmfRectilinearGrid>(item);
  if (recGrid)
  {
    unsigned int nSets = recGrid->getNumberSets();
    if (nSets > 0)
    {
      vtkMultiBlockDataSet* mbds = vtkMultiBlockDataSet::SafeDownCast(toFill);
      vtkRectilinearGrid* child = vtkRectilinearGrid::New();
      mbds->SetBlock(0, this->MakeRecGrid(recGrid, child, this->Keeper));
      mbds->GetMetaData((unsigned int)0)
        ->Set(vtkCompositeDataSet::NAME(), recGrid->getName().c_str());
      for (unsigned int i = 0; i < nSets; i++)
      {
        vtkUnstructuredGrid* sub = vtkUnstructuredGrid::New();
        mbds->SetBlock(i + 1, this->ExtractSet(i, recGrid, child, sub, this->Keeper));
        mbds->GetMetaData(i + 1)->Set(
          vtkCompositeDataSet::NAME(), recGrid->getSet(i)->getName().c_str());
        sub->Delete();
      }
      child->Delete();
      return mbds;
    }
    return this->MakeRecGrid(recGrid, vtkRectilinearGrid::SafeDownCast(toFill), this->Keeper);
  }

  shared_ptr<XdmfCurvilinearGrid> crvGrid = shared_dynamic_cast<XdmfCurvilinearGrid>(item);
  if (crvGrid)
  {
    unsigned int nSets = crvGrid->getNumberSets();
    if (nSets > 0)
    {
      vtkMultiBlockDataSet* mbds = vtkMultiBlockDataSet::SafeDownCast(toFill);
      vtkStructuredGrid* child = vtkStructuredGrid::New();
      mbds->SetBlock(0, this->MakeCrvGrid(crvGrid, child, this->Keeper));
      mbds->GetMetaData((unsigned int)0)
        ->Set(vtkCompositeDataSet::NAME(), crvGrid->getName().c_str());
      for (unsigned int i = 0; i < nSets; i++)
      {
        vtkUnstructuredGrid* sub = vtkUnstructuredGrid::New();
        mbds->SetBlock(i + 1, this->ExtractSet(i, crvGrid, child, sub, this->Keeper));
        mbds->GetMetaData(i + 1)->Set(
          vtkCompositeDataSet::NAME(), crvGrid->getSet(i)->getName().c_str());
        sub->Delete();
      }
      child->Delete();
      return mbds;
    }
    return this->MakeCrvGrid(crvGrid, vtkStructuredGrid::SafeDownCast(toFill), this->Keeper);
  }

  shared_ptr<XdmfRegularGrid> regGrid = shared_dynamic_cast<XdmfRegularGrid>(item);
  if (regGrid)
  {
    unsigned int nSets = regGrid->getNumberSets();
    if (nSets > 0)
    {
      vtkMultiBlockDataSet* mbds = vtkMultiBlockDataSet::SafeDownCast(toFill);
      vtkImageData* child = vtkImageData::New();
      mbds->SetBlock(0, this->MakeRegGrid(regGrid, child, this->Keeper));
      mbds->GetMetaData((unsigned int)0)
        ->Set(vtkCompositeDataSet::NAME(), regGrid->getName().c_str());
      for (unsigned int i = 0; i < nSets; i++)
      {
        vtkUnstructuredGrid* sub = vtkUnstructuredGrid::New();
        mbds->SetBlock(i + 1, this->ExtractSet(i, regGrid, child, sub, this->Keeper));
        mbds->GetMetaData(i + 1)->Set(
          vtkCompositeDataSet::NAME(), crvGrid->getSet(i)->getName().c_str());
        sub->Delete();
      }
      child->Delete();
      return mbds;
    }
    return this->MakeRegGrid(regGrid, vtkImageData::SafeDownCast(toFill), this->Keeper);
  }
  return nullptr; // already spit a warning out before this
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::Populate(shared_ptr<XdmfGraph> item, vtkDataObject* toFill)
{
  assert(toFill);

  shared_ptr<XdmfGraph> graph = shared_dynamic_cast<XdmfGraph>(item);
  if (graph)
  {
    return this->MakeGraph(graph, vtkMutableDirectedGraph::SafeDownCast(toFill), this->Keeper);
  }
  return nullptr; // already spit a warning out before this
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::ShouldRead(unsigned int piece, unsigned int npieces)
{
  if (this->NumProcs < 1)
  {
    // no parallel information given to us, assume serial
    return true;
  }
  if (npieces == 1)
  {
    return true;
  }
  if (npieces < this->NumProcs)
  {
    return piece == this->Rank;
  }

#if 1
  unsigned int mystart = this->Rank * npieces / this->NumProcs;
  unsigned int myend = (this->Rank + 1) * npieces / this->NumProcs;
  if (piece >= mystart)
  {
    if (piece < myend || (this->Rank == this->NumProcs - 1))
    {
      return true;
    }
  }
  return false;
#else
  if ((piece % this->NumProcs) == this->Rank)
  {
    return true;
  }
  else
  {
    return false;
  }
#endif
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::GridEnabled(shared_ptr<XdmfGrid> grid)
{
  return this->GridsCache->ArrayIsEnabled(grid->getName().c_str());
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::GridEnabled(shared_ptr<XdmfGraph> graph)
{
  return this->GridsCache->ArrayIsEnabled(graph->getName().c_str());
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::SetEnabled(shared_ptr<XdmfSet> set)
{
  return this->SetsCache->ArrayIsEnabled(set->getName().c_str());
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::ForThisTime(shared_ptr<XdmfGrid> grid)
{
  return (!this->doTime || (grid->getTime() && grid->getTime()->getValue() == this->time));
}

//------------------------------------------------------------------------------
bool vtkXdmf3HeavyDataHandler::ForThisTime(shared_ptr<XdmfGraph> graph)
{
  return (!this->doTime || (graph->getTime() && graph->getTime()->getValue() == this->time));
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::MakeUnsGrid(
  shared_ptr<XdmfUnstructuredGrid> grid, vtkUnstructuredGrid* dataSet, vtkXdmf3ArrayKeeper* keeper)
{
  if (dataSet && this->GridEnabled(grid) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfToVTK(
      this->FieldArrays, this->CellArrays, this->PointArrays, grid.get(), dataSet, keeper);
    return dataSet;
  }
  return nullptr;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::MakeRecGrid(
  shared_ptr<XdmfRectilinearGrid> grid, vtkRectilinearGrid* dataSet, vtkXdmf3ArrayKeeper* keeper)
{
  if (dataSet && this->GridEnabled(grid) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfToVTK(
      this->FieldArrays, this->CellArrays, this->PointArrays, grid.get(), dataSet, keeper);
    return dataSet;
  }
  return nullptr;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::MakeCrvGrid(
  shared_ptr<XdmfCurvilinearGrid> grid, vtkStructuredGrid* dataSet, vtkXdmf3ArrayKeeper* keeper)
{
  if (dataSet && this->GridEnabled(grid) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfToVTK(
      this->FieldArrays, this->CellArrays, this->PointArrays, grid.get(), dataSet, keeper);
    return dataSet;
  }
  return nullptr;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::MakeRegGrid(
  shared_ptr<XdmfRegularGrid> grid, vtkImageData* dataSet, vtkXdmf3ArrayKeeper* keeper)
{
  if (dataSet && this->GridEnabled(grid) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfToVTK(
      this->FieldArrays, this->CellArrays, this->PointArrays, grid.get(), dataSet, keeper);
    return dataSet;
  }
  return nullptr;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::MakeGraph(
  shared_ptr<XdmfGraph> grid, vtkMutableDirectedGraph* dataSet, vtkXdmf3ArrayKeeper* keeper)
{
  if (dataSet && this->GridEnabled(grid) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfToVTK(
      this->FieldArrays, this->CellArrays, this->PointArrays, grid.get(), dataSet, keeper);
    return dataSet;
  }
  return nullptr;
}

//------------------------------------------------------------------------------
vtkDataObject* vtkXdmf3HeavyDataHandler::ExtractSet(unsigned int setnum, shared_ptr<XdmfGrid> grid,
  vtkDataSet* dataSet, vtkUnstructuredGrid* subSet, vtkXdmf3ArrayKeeper* keeper)
{
  shared_ptr<XdmfSet> set = grid->getSet(setnum);
  if (dataSet && subSet && SetEnabled(set) && this->ForThisTime(grid))
  {
    vtkXdmf3DataSet::XdmfSubsetToVTK(grid.get(), setnum, dataSet, subSet, keeper);
    return subSet;
  }
  return nullptr;
}