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

  Program:   Visualization Toolkit
  Module:    vtkCellDistanceSelector.h

  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 "vtkCellDistanceSelector.h"

#include "vtkCell.h"
#include "vtkCellData.h"
#include "vtkCellLinks.h"
#include "vtkCompositeDataIterator.h"
#include "vtkCompositeDataSet.h"
#include "vtkDataSet.h"
#include "vtkIdList.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkSelection.h"
#include "vtkSelectionNode.h"
#include "vtkSmartPointer.h"
#include "vtkStructuredGrid.h"
#include "vtkUnstructuredGrid.h"

#include <map>
#include <vector>

vtkStandardNewMacro(vtkCellDistanceSelector);

//------------------------------------------------------------------------------
vtkCellDistanceSelector::vtkCellDistanceSelector()
{
  this->Distance = 1;
  this->IncludeSeed = 1;
  this->AddIntermediate = 1;
  this->SetNumberOfInputPorts(2);
}

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

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

//------------------------------------------------------------------------------
int vtkCellDistanceSelector::FillInputPortInformation(int port, vtkInformation* info)
{
  switch (port)
  {
    case INPUT_MESH:
      info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkCompositeDataSet");
      break;
    case INPUT_SELECTION:
      info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkSelection");
      break;
  }
  return 1;
}

//------------------------------------------------------------------------------
void vtkCellDistanceSelector::AddSelectionNode(
  vtkSelection* output, vtkSmartPointer<vtkDataArray> outIndices, int composite_index, int d)
{
  vtkSmartPointer<vtkSelectionNode> outSelNode = vtkSmartPointer<vtkSelectionNode>::New();
  outSelNode->SetContentType(vtkSelectionNode::INDICES);
  outSelNode->SetFieldType(vtkSelectionNode::CELL);
  outSelNode->GetProperties()->Set(vtkSelectionNode::COMPOSITE_INDEX(), composite_index);
  // NB: Use HIERARCHICAL_LEVEL key to store distance to original cells
  outSelNode->GetProperties()->Set(vtkSelectionNode::HIERARCHICAL_LEVEL(), d);
  outSelNode->SetSelectionList(outIndices);
  output->AddNode(outSelNode);
}

//------------------------------------------------------------------------------
int vtkCellDistanceSelector::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{

  // Retrieve input mesh as composite object
  vtkInformation* inDataObjectInfo = inputVector[INPUT_MESH]->GetInformationObject(0);
  vtkCompositeDataSet* compositeInput =
    vtkCompositeDataSet::SafeDownCast(inDataObjectInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Retrieve input selection
  vtkInformation* inSelectionInfo = inputVector[INPUT_SELECTION]->GetInformationObject(0);
  vtkSelection* inputSelection =
    vtkSelection::SafeDownCast(inSelectionInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Retrieve output selection
  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  vtkSelection* output = vtkSelection::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  if (!compositeInput)
  {
    vtkErrorMacro(<< "Missing input data object");
    return 0;
  }

  if (!inputSelection)
  {
    vtkErrorMacro(<< "Missing input selection");
    return 0;
  }

  std::map<int, std::vector<vtkSelectionNode*>> partSelections;
  int nSelNodes = inputSelection->GetNumberOfNodes();
  for (int i = 0; i < nSelNodes; ++i)
  {
    vtkSelectionNode* sn = inputSelection->GetNode(i);
    int composite_index = sn->GetProperties()->Get(vtkSelectionNode::COMPOSITE_INDEX());
    partSelections[composite_index].push_back(sn);
  }

  vtkCompositeDataIterator* inputIterator = compositeInput->NewIterator();
  inputIterator->SkipEmptyNodesOn();
  inputIterator->InitTraversal();
  inputIterator->GoToFirstItem();
  while (!inputIterator->IsDoneWithTraversal())
  {
    vtkDataSet* input = vtkDataSet::SafeDownCast(inputIterator->GetCurrentDataObject());
    // NB: composite indices start at 1
    int composite_index = inputIterator->GetCurrentFlatIndex();
    inputIterator->GoToNextItem();

    std::vector<vtkSelectionNode*>::iterator selNodeIt = partSelections[composite_index].begin();
    while (selNodeIt != partSelections[composite_index].end())
    {
      vtkSelectionNode* selectionNode = *selNodeIt;
      ++selNodeIt;

      vtkDataArray* selectionList =
        vtkArrayDownCast<vtkDataArray>(selectionNode->GetSelectionList());
      vtkIdType numSeeds = selectionList->GetNumberOfTuples();
      if (numSeeds > 0 && selectionNode->GetContentType() == vtkSelectionNode::INDICES &&
        selectionNode->GetFieldType() == vtkSelectionNode::CELL && input->GetNumberOfCells() > 0)
      {
        vtkIdType numCells = input->GetNumberOfCells();

        vtkUnstructuredGrid* ug_input = vtkUnstructuredGrid::SafeDownCast(input);
        vtkStructuredGrid* sg_input = vtkStructuredGrid::SafeDownCast(input);
        vtkPolyData* pd_input = vtkPolyData::SafeDownCast(input);

        if (ug_input)
        {
          if (!ug_input->GetCellLinks())
          {
            ug_input->BuildLinks();
          }
        }

        std::vector<int> flags(numCells, 0);

        vtkSmartPointer<vtkIdTypeArray> outIndices = vtkSmartPointer<vtkIdTypeArray>::New();
        outIndices->SetNumberOfTuples(numSeeds);

        int seedCount = 0;
        for (int i = 0; i < numSeeds; ++i)
        {
          vtkIdType cellIndex = static_cast<vtkIdType>(selectionList->GetTuple1(i));
          if (cellIndex >= 0 && cellIndex < numCells)
          {
            flags[cellIndex] = true;
            outIndices->SetTuple1(seedCount++, cellIndex);
          }
          else
          {
            vtkWarningMacro(<< "Cell index out of bounds in selection (" << cellIndex << "/"
                            << numCells << ")\n");
          }
        }
        outIndices->SetNumberOfTuples(seedCount);

        vtkSmartPointer<vtkIdTypeArray> finalIndices = vtkSmartPointer<vtkIdTypeArray>::New();
        vtkSmartPointer<vtkIntArray> cellDistance = vtkSmartPointer<vtkIntArray>::New();
        cellDistance->SetName("Cell Distance");

        // Iterate over increasing topological distance until desired distance is met
        for (int d = 0; d < this->Distance; ++d)
        {
          vtkSmartPointer<vtkIdTypeArray> nextIndices = vtkSmartPointer<vtkIdTypeArray>::New();

          if (ug_input)
          {
            int nIndices = outIndices->GetNumberOfTuples();
            for (int i = 0; i < nIndices; ++i)
            {
              vtkIdType cellIndex = static_cast<vtkIdType>(outIndices->GetTuple1(i));
              const vtkIdType* points;
              vtkIdType n;
              ug_input->GetCellPoints(cellIndex, n, points);
              for (int k = 0; k < n; ++k)
              {
                vtkIdType pid = points[k];
                vtkIdType np, *cells;
                ug_input->GetPointCells(pid, np, cells);
                for (int j = 0; j < np; ++j)
                {
                  vtkIdType cid = cells[j];
                  if (cid >= 0 && cid < numCells)
                  {
                    if (!flags[cid])
                    {
                      flags[cid] = true;
                      nextIndices->InsertNextValue(cid);
                    }
                  }
                  else
                  {
                    vtkWarningMacro(<< "Selection's cell index out of bounds (" << cid << "/"
                                    << numCells << ")\n");
                  }
                }
              }
            }
          } // if ( ug_input )
          else if (pd_input)
          {
            pd_input->BuildLinks();
            int nIndices = outIndices->GetNumberOfTuples();
            for (int i = 0; i < nIndices; ++i)
            {
              vtkIdType cellIndex = static_cast<vtkIdType>(outIndices->GetTuple1(i));
              const vtkIdType* points;
              vtkIdType n;
              pd_input->GetCellPoints(cellIndex, n, points);
              for (int k = 0; k < n; ++k)
              {
                vtkIdType pid = points[k];
                vtkIdType np;
                vtkIdType* cells;
                pd_input->GetPointCells(pid, np, cells);
                for (int j = 0; j < np; j++)
                {
                  vtkIdType cid = cells[j];
                  if (cid >= 0 && cid < numCells)
                  {
                    if (!flags[cid])
                    {
                      flags[cid] = true;
                      nextIndices->InsertNextValue(cid);
                    }
                  }
                  else
                  {
                    vtkWarningMacro(<< "Selection's cell index out of bounds (" << cid << "/"
                                    << numCells << ")\n");
                  }
                }
              }
            }
          } // else if ( ug_input )
          else if (sg_input)
          {
            int dim[3];
            sg_input->GetDimensions(dim);
            --dim[0];
            --dim[1];
            --dim[2];

            int nIndices = outIndices->GetNumberOfTuples();
            for (int idx = 0; idx < nIndices; ++idx)
            {
              vtkIdType cellIndex = static_cast<vtkIdType>(outIndices->GetTuple1(idx));
              vtkIdType cellId = cellIndex;
              vtkIdType ijk[3];
              for (int c = 0; c < 3; ++c)
              {
                if (dim[c] <= 1)
                {
                  ijk[c] = 0;
                }
                else
                {
                  ijk[c] = cellId % dim[c];
                  cellId /= dim[c];
                }
              } // c
              for (int k = -1; k <= 1; ++k)
              {
                for (int j = -1; j <= 1; ++j)
                {
                  for (int i = -1; i <= 1; ++i)
                  {
                    int I = ijk[0] + i;
                    int J = ijk[1] + j;
                    int K = ijk[2] + k;
                    if (I >= 0 && I < dim[0] && J >= 0 && J < dim[1] && K >= 0 && K < dim[2])
                    {
                      cellId = I + J * dim[0] + K * dim[0] * dim[1];
                      if (cellId >= 0 && cellId < numCells)
                      {
                        if (!flags[cellId])
                        {
                          flags[cellId] = true;
                          nextIndices->InsertNextValue(cellId);
                        }
                      }
                      else
                      {
                        vtkWarningMacro(<< "Selection's cell index out of bounds (" << cellId << "/"
                                        << numCells << ")\n");
                      }
                    }
                  } // i
                }   // j
              }     // k
            }       // idx
          }         // else if ( sg_input )
          else
          {
            vtkErrorMacro(<< "Unsupported data type : " << input->GetClassName() << "\n");
          }

          if ((!d && this->IncludeSeed) || (d > 0 && this->AddIntermediate))
          {
            int ni = outIndices->GetNumberOfTuples();
            for (int i = 0; i < ni; ++i)
            {
              cellDistance->InsertNextTuple1(d);
              finalIndices->InsertNextTuple1(outIndices->GetTuple1(i));
            } // i
          }   // if ( ( ! d && this->IncludeSeed ) || ( d > 0 && this->AddIntermediate ) )

          outIndices = nextIndices;
        } // for ( int d = 0; d < this->Distance; ++ d )

        //        if( ( ! this->Distance && this->IncludeSeed ) || ( this->Distance > 0 &&
        //        this->AddIntermediate ) )
        if ((!this->Distance && this->IncludeSeed) || this->Distance > 0)
        {
          int ni = outIndices->GetNumberOfTuples();
          cerr << "There are " << ni << " tuples\n";
          for (int i = 0; i < ni; ++i)
          {
            cellDistance->InsertNextTuple1(this->Distance);
            finalIndices->InsertNextTuple1(outIndices->GetTuple1(i));
          } // i
        }

        // Store selected cells for given seed cell
        if (finalIndices->GetNumberOfTuples() > 0)
        {
          vtkSmartPointer<vtkSelectionNode> outSelNode = vtkSmartPointer<vtkSelectionNode>::New();
          outSelNode->SetContentType(vtkSelectionNode::INDICES);
          outSelNode->SetFieldType(vtkSelectionNode::CELL);
          outSelNode->GetProperties()->Set(vtkSelectionNode::COMPOSITE_INDEX(), composite_index);
          outSelNode->SetSelectionList(finalIndices);
          outSelNode->GetSelectionData()->AddArray(cellDistance);
          output->AddNode(outSelNode);
        }
      } // if numSeeds > 0 etc.
    }   // while selNodeIt
  }     // while inputIterator

  // Clean up
  inputIterator->Delete();

  return 1;
}
