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

  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"

#include "XdmfCurvilinearGrid.hpp"
#include "XdmfDomain.hpp"
#include "XdmfGraph.hpp"
#include "XdmfGrid.hpp"
#include "XdmfGridCollection.hpp"
#include "XdmfGridCollectionType.hpp"
#include "XdmfItem.hpp"
#include "XdmfRectilinearGrid.hpp"
#include "XdmfRegularGrid.hpp"
#include "XdmfSet.hpp"
#include "XdmfUnstructuredGrid.hpp"

#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;
}

//------------------------------------------------------------------------------
vtkXdmf3HeavyDataHandler::vtkXdmf3HeavyDataHandler()
{
}

//------------------------------------------------------------------------------
vtkXdmf3HeavyDataHandler::~vtkXdmf3HeavyDataHandler()
{
}

//------------------------------------------------------------------------------
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?false:true;
  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)
  {
    if (piece == this->Rank)
    {
      return true;
    }
    return false;
  }

#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;
}