// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkExtractCellsByType.h"

#include "vtkCellArray.h"
#include "vtkCellData.h"
#include "vtkIdList.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPointSet.h"
#include "vtkPoints.h"
#include "vtkPolyData.h"
#include "vtkRectilinearGrid.h"
#include "vtkStructuredGrid.h"
#include "vtkUniformGrid.h"
#include "vtkUnstructuredGrid.h"

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkExtractCellsByType);

VTK_ABI_NAMESPACE_END
#include <set>

// Special token marks any cell type
#define VTK_ANY_CELL_TYPE 1000000

VTK_ABI_NAMESPACE_BEGIN
struct vtkCellTypeSet : public std::set<unsigned int>
{
};

//------------------------------------------------------------------------------
vtkExtractCellsByType::vtkExtractCellsByType()
{
  this->CellTypes = new vtkCellTypeSet;
}

//------------------------------------------------------------------------------
vtkExtractCellsByType::~vtkExtractCellsByType()
{
  delete this->CellTypes;
}

//------------------------------------------------------------------------------
void vtkExtractCellsByType::AddCellType(unsigned int cellType)
{
  auto prevSize = this->CellTypes->size();
  this->CellTypes->insert(cellType);
  if (prevSize != this->CellTypes->size())
  {
    this->Modified();
  }
}

//------------------------------------------------------------------------------
void vtkExtractCellsByType::RemoveCellType(unsigned int cellType)
{
  auto prevSize = this->CellTypes->size();
  this->CellTypes->erase(cellType);
  this->CellTypes->erase(VTK_ANY_CELL_TYPE);
  if (prevSize != this->CellTypes->size())
  {
    this->Modified();
  }
}

//------------------------------------------------------------------------------
void vtkExtractCellsByType::RemoveAllCellTypes()
{
  if (!this->CellTypes->empty())
  {
    this->CellTypes->clear();
    this->Modified();
  }
}

//------------------------------------------------------------------------------
// Special value indicates that all cells are to be selected. This is better
// than populating from the list vtkCellType.h due to the associated
// maintenance burden.
void vtkExtractCellsByType::AddAllCellTypes()
{
  auto prevSize = this->CellTypes->size();
  this->CellTypes->insert(VTK_ANY_CELL_TYPE);
  if (prevSize != this->CellTypes->size())
  {
    this->Modified();
  }
}

//------------------------------------------------------------------------------
bool vtkExtractCellsByType::ExtractCellType(unsigned int cellType)
{
  return this->CellTypes->find(cellType) != this->CellTypes->end() ||
    this->CellTypes->find(VTK_ANY_CELL_TYPE) != this->CellTypes->end();
}

//------------------------------------------------------------------------------
void vtkExtractCellsByType::ExtractUnstructuredData(vtkDataSet* inDS, vtkDataSet* outDS)
{
  vtkPointData* inPD = inDS->GetPointData();
  vtkPointData* outPD = outDS->GetPointData();

  vtkIdType numPts = inDS->GetNumberOfPoints();

  // We are going some maps to indicate where the points and cells originated
  // from. Values <0 mean that the points or cells are not mapped to the
  // output.
  vtkIdType* ptMap = new vtkIdType[numPts];
  std::fill_n(ptMap, numPts, -1);

  // Now dispatch to specific unstructured type
  vtkIdType numNewPts = 0;
  if (inDS->GetDataObjectType() == VTK_POLY_DATA)
  {
    this->ExtractPolyDataCells(inDS, outDS, ptMap, numNewPts);
  }

  else if (inDS->GetDataObjectType() == VTK_UNSTRUCTURED_GRID)
  {
    this->ExtractUnstructuredGridCells(inDS, outDS, ptMap, numNewPts);
  }

  // Define points using point mapping for extracted cells
  if (numNewPts > 0)
  {
    // Copy referenced input points to new points array
    outPD->CopyAllocate(inPD);
    vtkPointSet* inPtSet = vtkPointSet::SafeDownCast(inDS);
    vtkPointSet* outPtSet = vtkPointSet::SafeDownCast(outDS);
    vtkPoints* inPts = inPtSet->GetPoints();
    vtkNew<vtkPoints> outPts;
    outPts->SetNumberOfPoints(numNewPts);
    for (vtkIdType ptId = 0; ptId < numPts; ++ptId)
    {
      if (ptMap[ptId] >= 0)
      {
        outPts->SetPoint(ptMap[ptId], inPts->GetPoint(ptId));
        outPD->CopyData(inPD, ptId, ptMap[ptId]);
      }
    }
    outPtSet->SetPoints(outPts);
  }

  // Clean up
  delete[] ptMap;
}

//------------------------------------------------------------------------------
void vtkExtractCellsByType::ExtractPolyDataCells(
  vtkDataSet* inDS, vtkDataSet* outDS, vtkIdType* ptMap, vtkIdType& numNewPts)
{
  vtkPolyData* input = vtkPolyData::SafeDownCast(inDS);
  vtkCellData* inCD = input->GetCellData();
  vtkPolyData* output = vtkPolyData::SafeDownCast(outDS);
  vtkCellData* outCD = output->GetCellData();

  // Treat the four cell arrays separately. If the array might have cells of
  // the specified types, then traverse it, copying the input cells to the
  // output cells. Also keep track of the point map.

  // The cellIds are numbered across the four arrays: verts, lines, polys,
  // strips. Have to carefully coordinate the cell ids with traversal of each
  // array.
  vtkIdType cellId, currentCellId = 0;
  vtkIdList* ptIds = vtkIdList::New();

  // Verts
  outCD->CopyAllocate(inCD);
  vtkIdType i;
  vtkIdType npts;
  const vtkIdType* pts;
  vtkCellArray* inVerts = input->GetVerts();
  vtkIdType progressCounter = 0;
  vtkIdType checkAbortInterval = std::min(inVerts->GetNumberOfCells() / 10 + 1, (vtkIdType)1000);
  if (this->ExtractCellType(VTK_VERTEX) || this->ExtractCellType(VTK_POLY_VERTEX))
  {
    vtkCellArray* verts = vtkCellArray::New();
    for (inVerts->InitTraversal(); inVerts->GetNextCell(npts, pts); ++currentCellId)
    {
      if (progressCounter % checkAbortInterval == 0 && this->CheckAbort())
      {
        break;
      }
      progressCounter++;
      if (this->ExtractCellType(input->GetCellType(currentCellId)))
      {
        ptIds->Reset();
        for (i = 0; i < npts; ++i)
        {
          if (ptMap[pts[i]] < 0)
          {
            ptMap[pts[i]] = numNewPts++;
          }
          ptIds->InsertId(i, ptMap[pts[i]]);
        }
        cellId = verts->InsertNextCell(ptIds);
        outCD->CopyData(inCD, currentCellId, cellId);
      }
    }
    output->SetVerts(verts);
    verts->Delete();
  }
  else
  {
    currentCellId += inVerts->GetNumberOfCells();
  }

  // Lines
  vtkCellArray* inLines = input->GetLines();
  if (this->ExtractCellType(VTK_LINE) || this->ExtractCellType(VTK_POLY_LINE))
  {
    vtkCellArray* lines = vtkCellArray::New();
    progressCounter = 0;
    checkAbortInterval = std::min(inLines->GetNumberOfCells() / 10 + 1, (vtkIdType)1000);
    for (inLines->InitTraversal(); inLines->GetNextCell(npts, pts); ++currentCellId)
    {
      if (progressCounter % checkAbortInterval == 0 && this->CheckAbort())
      {
        break;
      }
      progressCounter++;
      if (this->ExtractCellType(input->GetCellType(currentCellId)))
      {
        ptIds->Reset();
        for (i = 0; i < npts; ++i)
        {
          if (ptMap[pts[i]] < 0)
          {
            ptMap[pts[i]] = numNewPts++;
          }
          ptIds->InsertId(i, ptMap[pts[i]]);
        }
        cellId = lines->InsertNextCell(ptIds);
        outCD->CopyData(inCD, currentCellId, cellId);
      }
    }
    output->SetLines(lines);
    lines->Delete();
  }
  else
  {
    currentCellId += inLines->GetNumberOfCells();
  }

  // Polys
  vtkCellArray* inPolys = input->GetPolys();
  if (this->ExtractCellType(VTK_TRIANGLE) || this->ExtractCellType(VTK_QUAD) ||
    this->ExtractCellType(VTK_POLYGON))
  {
    vtkCellArray* polys = vtkCellArray::New();
    progressCounter = 0;
    checkAbortInterval = std::min(inPolys->GetNumberOfCells() / 10 + 1, (vtkIdType)1000);
    for (inPolys->InitTraversal(); inPolys->GetNextCell(npts, pts); ++currentCellId)
    {
      if (progressCounter % checkAbortInterval == 0 && this->CheckAbort())
      {
        break;
      }
      progressCounter++;
      if (this->ExtractCellType(input->GetCellType(currentCellId)))
      {
        ptIds->Reset();
        for (i = 0; i < npts; ++i)
        {
          if (ptMap[pts[i]] < 0)
          {
            ptMap[pts[i]] = numNewPts++;
          }
          ptIds->InsertId(i, ptMap[pts[i]]);
        }
        cellId = polys->InsertNextCell(ptIds);
        outCD->CopyData(inCD, currentCellId, cellId);
      }
    }
    output->SetPolys(polys);
    polys->Delete();
  }
  else
  {
    currentCellId += inPolys->GetNumberOfCells();
  }

  // Triangle strips
  vtkCellArray* inStrips = input->GetStrips();
  if (this->ExtractCellType(VTK_TRIANGLE_STRIP))
  {
    vtkCellArray* strips = vtkCellArray::New();
    checkAbortInterval = std::min(inStrips->GetNumberOfCells() / 10 + 1, (vtkIdType)1000);
    progressCounter++;
    // All cells are of type VTK_TRIANGLE_STRIP
    for (inStrips->InitTraversal(); inStrips->GetNextCell(npts, pts); ++currentCellId)
    {
      if (progressCounter % checkAbortInterval == 0 && this->CheckAbort())
      {
        break;
      }
      progressCounter++;
      ptIds->Reset();
      for (i = 0; i < npts; ++i)
      {
        if (ptMap[pts[i]] < 0)
        {
          ptMap[pts[i]] = numNewPts++;
        }
        ptIds->InsertId(i, ptMap[pts[i]]);
      }
      cellId = strips->InsertNextCell(ptIds);
      outCD->CopyData(inCD, currentCellId, cellId);
    }
    output->SetStrips(strips);
    strips->Delete();
  }

  // Clean up
  ptIds->Delete();
}

//------------------------------------------------------------------------------
// Helper
namespace
{
struct ExtractPolyVisitor
{
  // Insert full cell
  template <typename CellStateT>
  vtkIdType operator()(CellStateT& state, vtkIdType* ptMap, vtkIdType numNewPts)
  {
    using ValueType = typename CellStateT::ValueType;
    auto* conn = state.GetConnectivity();
    const vtkIdType nids = conn->GetNumberOfValues();
    for (vtkIdType i = 0; i < nids; ++i)
    {
      ValueType ptId = conn->GetValue(i);
      if (ptMap[ptId] < 0)
      {
        ptMap[ptId] = numNewPts++;
      }
      conn->SetValue(i, ptMap[ptId]);
    }
    return numNewPts;
  }
};
}
//------------------------------------------------------------------------------
void vtkExtractCellsByType::ExtractUnstructuredGridCells(
  vtkDataSet* inDS, vtkDataSet* outDS, vtkIdType* ptMap, vtkIdType& numNewPts)
{
  vtkUnstructuredGrid* input = vtkUnstructuredGrid::SafeDownCast(inDS);
  vtkCellData* inCD = input->GetCellData();
  vtkUnstructuredGrid* output = vtkUnstructuredGrid::SafeDownCast(outDS);
  vtkCellData* outCD = output->GetCellData();

  vtkIdType numCells = input->GetNumberOfCells();

  // Check for trivial cases: either all in or all out
  if (input->IsHomogeneous())
  {
    if (this->ExtractCellType(input->GetCellType(0)))
    {
      output->ShallowCopy(input);
    }
    else
    {
      output->Initialize();
    }
    return;
  }

  // Mixed collection of cells so simply loop over all cells, copying
  // appropriate types to the output. Along the way keep track of the points
  // that are used.
  vtkIdType i, cellId, newCellId, npts, ptId;
  vtkNew<vtkIdList> ptIds;
  vtkNew<vtkCellArray> faces;
  int cellType;
  output->Allocate(numCells);
  outCD->CopyAllocate(inCD);

  vtkIdType checkAbortInterval = std::min(numCells / 10 + 1, (vtkIdType)1000);

  for (cellId = 0; cellId < numCells; ++cellId)
  {
    if (cellId % checkAbortInterval == 0 && this->CheckAbort())
    {
      break;
    }
    cellType = input->GetCellType(cellId);
    if (this->ExtractCellType(cellType))
    {
      if (cellType == VTK_POLYHEDRON)
      {
        faces->Reset();
        input->GetPolyhedronFaces(cellId, faces);
        numNewPts = faces->Visit(ExtractPolyVisitor{}, ptMap, numNewPts);
        input->GetCellPoints(cellId, ptIds);
        npts = ptIds->GetNumberOfIds();
        for (i = 0; i < npts; ++i)
        {
          ptId = ptIds->GetId(i);
          ptIds->InsertId(i, ptMap[ptId]);
        }
        newCellId = output->InsertNextCell(VTK_POLYHEDRON, npts, ptIds->GetPointer(0), faces);
      }
      else
      {
        input->GetCellPoints(cellId, ptIds);
        npts = ptIds->GetNumberOfIds();
        for (i = 0; i < npts; ++i)
        {
          ptId = ptIds->GetId(i);
          if (ptMap[ptId] < 0)
          {
            ptMap[ptId] = numNewPts++;
          }
          ptIds->InsertId(i, ptMap[ptId]);
        }
        newCellId = output->InsertNextCell(cellType, ptIds);
      }
      outCD->CopyData(inCD, cellId, newCellId);
      ptIds->Reset();
    }
  }
}

//------------------------------------------------------------------------------
int vtkExtractCellsByType::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  // get the input and output
  vtkDataSet* input = vtkDataSet::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkDataSet* output = vtkDataSet::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Handle the trivial case
  vtkIdType numCells = input->GetNumberOfCells();
  if (this->CellTypes->empty() || numCells <= 0)
  {
    output->Initialize(); // output is empty
    return 1;
  }

  // Dispatch to appropriate type. This filter does not directly handle
  // composite dataset types, composite types should be looped over by
  // the pipeline executive.
  if (input->GetDataObjectType() == VTK_POLY_DATA ||
    input->GetDataObjectType() == VTK_UNSTRUCTURED_GRID)
  {
    this->ExtractUnstructuredData(input, output);
  }

  // Structured data has only one cell type per dataset
  else if (input->GetDataObjectType() == VTK_IMAGE_DATA ||
    input->GetDataObjectType() == VTK_STRUCTURED_POINTS ||
    input->GetDataObjectType() == VTK_RECTILINEAR_GRID ||
    input->GetDataObjectType() == VTK_STRUCTURED_GRID ||
    input->GetDataObjectType() == VTK_UNIFORM_GRID ||
    input->GetDataObjectType() == VTK_HYPER_TREE_GRID)
  {
    if (this->ExtractCellType(input->GetCellType(0)))
    {
      output->ShallowCopy(input);
    }
    else
    {
      output->Initialize(); // output is empty
    }
  }

  else
  {
    vtkErrorMacro("Unknown dataset type");
    output->Initialize(); // output is empty
  }

  this->CheckAbort();

  return 1;
}

//------------------------------------------------------------------------------
int vtkExtractCellsByType::FillInputPortInformation(int, vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkDataSet");
  return 1;
}

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

  // Output the number of types specified
  os << indent << "Number of types specified: " << this->CellTypes->size() << "\n";
}
VTK_ABI_NAMESPACE_END
