// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-FileCopyrightText: Copyright (c) Sandia Corporation
// SPDX-License-Identifier: BSD-3-Clause

#include "vtkKdTreeSelector.h"

#include "vtkGraph.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkKdTree.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPointSet.h"
#include "vtkSelection.h"
#include "vtkSelectionNode.h"
#include "vtkSmartPointer.h"

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkKdTreeSelector);

vtkKdTreeSelector::vtkKdTreeSelector()
{
  this->SelectionBounds[0] = 0.0;
  this->SelectionBounds[1] = -1.0;
  this->SelectionBounds[2] = 0.0;
  this->SelectionBounds[3] = -1.0;
  this->SelectionBounds[4] = VTK_DOUBLE_MIN;
  this->SelectionBounds[5] = VTK_DOUBLE_MAX;
  this->KdTree = nullptr;
  this->BuildKdTreeFromInput = true;
  this->SelectionFieldName = nullptr;
  this->SingleSelection = false;
  this->SingleSelectionThreshold = 1.0;
  this->SelectionAttribute = -1;
}

vtkKdTreeSelector::~vtkKdTreeSelector()
{
  this->SetKdTree(nullptr);
  this->SetSelectionFieldName(nullptr);
}

void vtkKdTreeSelector::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "KdTree: " << (this->KdTree ? "" : "(null)") << endl;
  if (this->KdTree)
  {
    this->KdTree->PrintSelf(os, indent.GetNextIndent());
  }
  os << indent
     << "SelectionFieldName: " << (this->SelectionFieldName ? this->SelectionFieldName : "(null)")
     << endl;
  os << indent << "BuildKdTreeFromInput: " << (this->BuildKdTreeFromInput ? "on" : "off") << endl;
  os << indent << "SelectionBounds: " << endl;
  os << indent << "  xmin, xmax = (" << this->SelectionBounds[0] << "," << this->SelectionBounds[1]
     << ")" << endl;
  os << indent << "  ymin, ymax = (" << this->SelectionBounds[2] << "," << this->SelectionBounds[3]
     << ")" << endl;
  os << indent << "  zmin, zmax = (" << this->SelectionBounds[4] << "," << this->SelectionBounds[5]
     << ")" << endl;
  os << indent << "SingleSelection: " << (this->SingleSelection ? "on" : "off") << endl;
  os << indent << "SingleSelectionThreshold: " << this->SingleSelectionThreshold << endl;
  os << indent << "SelectionAttribute: " << this->SelectionAttribute << endl;
}

void vtkKdTreeSelector::SetKdTree(vtkKdTree* arg)
{
  vtkDebugMacro(<< this->GetClassName() << " (" << this << "): setting KdTree to " << arg);
  if (this->KdTree != arg)
  {
    vtkKdTree* tempSGMacroVar = this->KdTree;
    this->KdTree = arg;
    if (this->KdTree != nullptr)
    {
      this->BuildKdTreeFromInput = false;
      this->KdTree->Register(this);
    }
    else
    {
      this->BuildKdTreeFromInput = true;
    }
    if (tempSGMacroVar != nullptr)
    {
      tempSGMacroVar->UnRegister(this);
    }
    this->Modified();
  }
}

vtkMTimeType vtkKdTreeSelector::GetMTime()
{
  vtkMTimeType mTime = this->Superclass::GetMTime();
  if (this->KdTree != nullptr)
  {
    vtkMTimeType time = this->KdTree->GetMTime();
    mTime = (time > mTime ? time : mTime);
  }
  return mTime;
}

int vtkKdTreeSelector::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  vtkAbstractArray* field = nullptr;
  vtkGraph* graph = nullptr;

  if (this->BuildKdTreeFromInput)
  {
    vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
    if (inInfo == nullptr)
    {
      vtkErrorMacro("No input, but building kd-tree from input");
      return 0;
    }
    vtkDataObject* input = inInfo->Get(vtkDataObject::DATA_OBJECT());
    if (input == nullptr)
    {
      vtkErrorMacro("Input is nullptr");
      return 0;
    }
    graph = vtkGraph::SafeDownCast(input);
    vtkPointSet* pointSet = vtkPointSet::SafeDownCast(input);
    if (!graph && !pointSet)
    {
      vtkErrorMacro("Input must be a graph or point set");
      return 0;
    }

    vtkPoints* points = nullptr;
    if (graph)
    {
      points = graph->GetPoints();
    }
    else
    {
      points = pointSet->GetPoints();
    }

    // If no points, there is nothing to do
    if (points == nullptr || points->GetNumberOfPoints() == 0)
    {
      return 1;
    }

    // Construct the kd-tree if we need to
    if (this->KdTree == nullptr || this->KdTree->GetMTime() < input->GetMTime())
    {
      if (this->KdTree == nullptr)
      {
        this->KdTree = vtkKdTree::New();
      }
      this->KdTree->Initialize();
      this->KdTree->BuildLocatorFromPoints(points);
    }

    // Look for selection field
    if (this->SelectionAttribute == vtkDataSetAttributes::GLOBALIDS ||
      this->SelectionAttribute == vtkDataSetAttributes::PEDIGREEIDS)
    {
      if (graph)
      {
        field = graph->GetVertexData()->GetAbstractAttribute(this->SelectionAttribute);
      }
      else
      {
        field = pointSet->GetPointData()->GetAbstractAttribute(this->SelectionAttribute);
      }
      if (field == nullptr)
      {
        vtkErrorMacro("Could not find attribute " << this->SelectionAttribute);
        return 0;
      }
    }
    if (this->SelectionFieldName)
    {
      if (graph)
      {
        field = graph->GetVertexData()->GetAbstractArray(this->SelectionFieldName);
      }
      else
      {
        field = pointSet->GetPointData()->GetAbstractArray(this->SelectionFieldName);
      }
      if (field == nullptr)
      {
        vtkErrorMacro("SelectionFieldName field not found");
        return 0;
      }
    }
  }

  // If no kd-tree, there is nothing to do
  if (this->KdTree == nullptr)
  {
    return 1;
  }

  // Use the kd-tree to find the selected points
  vtkIdTypeArray* ids = vtkIdTypeArray::New();
  if (this->SingleSelection)
  {
    double center[3];
    for (int c = 0; c < 3; c++)
    {
      center[c] = (this->SelectionBounds[2 * c] + this->SelectionBounds[2 * c + 1]) / 2.0;
    }
    double dist;
    vtkIdType closestToCenter = this->KdTree->FindClosestPoint(center, dist);
    if (dist < this->SingleSelectionThreshold)
    {
      ids->InsertNextValue(closestToCenter);
    }
  }
  else
  {
    this->KdTree->FindPointsInArea(this->SelectionBounds, ids);
  }

  // Fill the selection with the found ids
  vtkSelection* output = vtkSelection::GetData(outputVector);
  vtkSmartPointer<vtkSelectionNode> node = vtkSmartPointer<vtkSelectionNode>::New();
  output->AddNode(node);
  if (graph)
  {
    node->SetFieldType(vtkSelectionNode::VERTEX);
  }
  else
  {
    node->SetFieldType(vtkSelectionNode::POINT);
  }
  if (field)
  {
    vtkAbstractArray* arr = vtkAbstractArray::CreateArray(field->GetDataType());
    arr->SetName(field->GetName());
    for (vtkIdType i = 0; i < ids->GetNumberOfTuples(); i++)
    {
      arr->InsertNextTuple(ids->GetValue(i), field);
    }
    if (this->SelectionAttribute == vtkDataSetAttributes::GLOBALIDS ||
      this->SelectionAttribute == vtkDataSetAttributes::PEDIGREEIDS)
    {
      if (this->SelectionAttribute == vtkDataSetAttributes::GLOBALIDS)
      {
        node->SetContentType(vtkSelectionNode::GLOBALIDS);
      }
      else if (this->SelectionAttribute == vtkDataSetAttributes::PEDIGREEIDS)
      {
        node->SetContentType(vtkSelectionNode::PEDIGREEIDS);
      }
    }
    else
    {
      node->SetContentType(vtkSelectionNode::VALUES);
    }
    node->SetSelectionList(arr);
    arr->Delete();
  }
  else
  {
    node->SetContentType(vtkSelectionNode::INDICES);
    node->SetSelectionList(ids);
  }

  // Clean up
  ids->Delete();

  return 1;
}

int vtkKdTreeSelector::FillInputPortInformation(int vtkNotUsed(port), vtkInformation* info)
{
  // Input (if specified) may be a point set or graph.
  info->Remove(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE());
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkPointSet");
  info->Append(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkGraph");
  info->Set(vtkAlgorithm::INPUT_IS_OPTIONAL(), 1);
  return 1;
}
VTK_ABI_NAMESPACE_END
