/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

// Software Guide : BeginLatex
//
// \index{itk::Statistics::KdTree}
// \index{itk::Statistics::KdTree\-Generator}
// \index{itk::Statistics::Weighted\-Centroid\-KdTree\-Generator}
//
// The \subdoxygen{Statistics}{KdTree} implements a data structure that
// separates samples in a $k$-dimension space.  The \code{std::vector} class
// is used here as the container for the measurement vectors from a sample.
//
// Software Guide : EndLatex

// Software Guide : BeginCodeSnippet
#include "itkVector.h"
#include "itkMath.h"
#include "itkListSample.h"
#include "itkWeightedCentroidKdTreeGenerator.h"
#include "itkEuclideanDistanceMetric.h"
// Software Guide : EndCodeSnippet

int
main()
{
  // Software Guide : BeginLatex
  //
  // We define the measurement vector type and instantiate a
  // \subdoxygen{Statistics}{ListSample} object, and then put 1000
  // measurement vectors in the object.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  using MeasurementVectorType = itk::Vector<float, 2>;

  using SampleType = itk::Statistics::ListSample<MeasurementVectorType>;
  SampleType::Pointer sample = SampleType::New();
  sample->SetMeasurementVectorSize(2);

  MeasurementVectorType mv;
  for (unsigned int i = 0; i < 1000; ++i)
  {
    mv[0] = (float)i;
    mv[1] = (float)((1000 - i) / 2);
    sample->PushBack(mv);
  }
  // Software Guide : EndCodeSnippet

  // Software Guide : BeginLatex
  //
  // The following code snippet shows how to create two KdTree objects. The
  // first object \subdoxygen{Statistics}{KdTreeGenerator} has a minimal set
  // of information (partition dimension, partition value, and pointers to
  // the left and right child nodes). The second tree from the
  // \subdoxygen{Statistics}{WeightedCentroidKdTreeGenerator} has additional
  // information such as the number of children under each node, and the
  // vector sum of the measurement vectors belonging to children of a
  // particular node.  WeightedCentroidKdTreeGenerator and the resulting k-d
  // tree structure were implemented based on the description given in the
  // paper by Kanungo et al \cite{Kanungo2000}.
  //
  // The instantiation and input variables are exactly the same for both
  // tree generators. Using the \code{SetSample()} method we plug-in the
  // source sample. The bucket size input specifies the limit on the
  // maximum number of measurement vectors that can be stored in a
  // terminal (leaf) node. A bigger bucket size results in a smaller number of
  // nodes in a tree. It also affects the efficiency of search. With
  // many small leaf nodes, we might experience slower search
  // performance because of excessive boundary comparisons.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  using TreeGeneratorType = itk::Statistics::KdTreeGenerator<SampleType>;
  TreeGeneratorType::Pointer treeGenerator = TreeGeneratorType::New();

  treeGenerator->SetSample(sample);
  treeGenerator->SetBucketSize(16);
  treeGenerator->Update();

  using CentroidTreeGeneratorType =
    itk::Statistics::WeightedCentroidKdTreeGenerator<SampleType>;

  CentroidTreeGeneratorType::Pointer centroidTreeGenerator =
    CentroidTreeGeneratorType::New();

  centroidTreeGenerator->SetSample(sample);
  centroidTreeGenerator->SetBucketSize(16);
  centroidTreeGenerator->Update();
  // Software Guide : EndCodeSnippet

  // Software Guide : BeginLatex
  //
  // After the generation step, we can get the pointer to the kd-tree
  // from the generator by calling the \code{GetOutput()} method. To
  // traverse a kd-tree, we have to use the \code{GetRoot()} method. The
  // method will return the root node of the tree. Every node in a tree
  // can have its left and/or right child node. To get the child node,
  // we call the \code{Left()} or the \code{Right()} method of a node
  // (these methods do not belong to the kd-tree but to the nodes).
  //
  // We can get other information about a node by calling the methods
  // described below in addition to the child node pointers.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  using TreeType = TreeGeneratorType::KdTreeType;
  using NodeType = TreeType::KdTreeNodeType;

  TreeType::Pointer tree = treeGenerator->GetOutput();
  TreeType::Pointer centroidTree = centroidTreeGenerator->GetOutput();

  NodeType * root = tree->GetRoot();

  if (root->IsTerminal())
  {
    std::cout << "Root node is a terminal node." << std::endl;
  }
  else
  {
    std::cout << "Root node is not a terminal node." << std::endl;
  }

  unsigned int partitionDimension;
  float        partitionValue;
  root->GetParameters(partitionDimension, partitionValue);
  std::cout << "Dimension chosen to split the space = " << partitionDimension
            << std::endl;
  std::cout << "Split point on the partition dimension = " << partitionValue
            << std::endl;

  std::cout << "Address of the left chile of the root node = " << root->Left()
            << std::endl;

  std::cout << "Address of the right chile of the root node = "
            << root->Right() << std::endl;

  root = centroidTree->GetRoot();
  std::cout << "Number of the measurement vectors under the root node"
            << " in the tree hierarchy = " << root->Size() << std::endl;

  NodeType::CentroidType centroid;
  root->GetWeightedCentroid(centroid);
  std::cout << "Sum of the measurement vectors under the root node = "
            << centroid << std::endl;

  std::cout << "Number of the measurement vectors under the left child"
            << " of the root node = " << root->Left()->Size() << std::endl;
  // Software Guide : EndCodeSnippet

  // Software Guide : BeginLatex
  //
  // In the following code snippet, we query the three nearest neighbors of
  // the \code{queryPoint} on the two tree. The results and procedures are
  // exactly the same for both. First we define the point from which distances
  // will be measured.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  MeasurementVectorType queryPoint;
  queryPoint[0] = 10.0;
  queryPoint[1] = 7.0;
  // Software Guide : EndCodeSnippet


  // Software Guide : BeginLatex
  //
  // Then we instantiate the type of a distance metric, create an object of
  // this type and set the origin of coordinates for measuring distances.
  // The \code{GetMeasurementVectorSize()} method returns the length of
  // each measurement vector stored in the sample.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  using DistanceMetricType =
    itk::Statistics::EuclideanDistanceMetric<MeasurementVectorType>;
  DistanceMetricType::Pointer distanceMetric = DistanceMetricType::New();

  DistanceMetricType::OriginType origin(2);
  for (unsigned int i = 0; i < sample->GetMeasurementVectorSize(); ++i)
  {
    origin[i] = queryPoint[i];
  }
  distanceMetric->SetOrigin(origin);
  // Software Guide : EndCodeSnippet


  // Software Guide : BeginLatex
  //
  // We can now set the number of neighbors to be located and the point
  // coordinates to be used as a reference system.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  unsigned int                           numberOfNeighbors = 3;
  TreeType::InstanceIdentifierVectorType neighbors;
  tree->Search(queryPoint, numberOfNeighbors, neighbors);

  std::cout
    << "\n*** kd-tree knn search result using an Euclidean distance metric:"
    << std::endl
    << "query point = [" << queryPoint << "]" << std::endl
    << "k = " << numberOfNeighbors << std::endl;
  std::cout << "measurement vector : distance from querry point "
            << std::endl;
  std::vector<double> distances1(numberOfNeighbors);
  for (unsigned int i = 0; i < numberOfNeighbors; ++i)
  {
    distances1[i] =
      distanceMetric->Evaluate(tree->GetMeasurementVector(neighbors[i]));
    std::cout << "[" << tree->GetMeasurementVector(neighbors[i])
              << "] : " << distances1[i] << std::endl;
  }
  // Software Guide : EndCodeSnippet

  // Software Guide : BeginLatex
  //
  // Instead of using an Euclidean distance metric, Tree itself can also
  // return the distance vector. Here we get the distance values from tree and
  // compare them with previous values.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  std::vector<double> distances2;
  tree->Search(queryPoint, numberOfNeighbors, neighbors, distances2);

  std::cout << "\n*** kd-tree knn search result directly from tree:"
            << std::endl
            << "query point = [" << queryPoint << "]" << std::endl
            << "k = " << numberOfNeighbors << std::endl;
  std::cout << "measurement vector : distance from querry point "
            << std::endl;
  for (unsigned int i = 0; i < numberOfNeighbors; ++i)
  {
    std::cout << "[" << tree->GetMeasurementVector(neighbors[i])
              << "] : " << distances2[i] << std::endl;
    if (itk::Math::NotAlmostEquals(distances2[i], distances1[i]))
    {
      std::cerr << "Mismatched distance values by tree." << std::endl;
      return EXIT_FAILURE;
    }
  }
  // Software Guide : EndCodeSnippet


  // Software Guide : BeginLatex
  //
  // As previously indicated, the interface for finding nearest neighbors in
  // the centroid tree is very similar.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  std::vector<double> distances3;
  centroidTree->Search(queryPoint, numberOfNeighbors, neighbors, distances3);

  centroidTree->Search(queryPoint, numberOfNeighbors, neighbors);
  std::cout << "\n*** Weighted centroid kd-tree knn search result:"
            << std::endl
            << "query point = [" << queryPoint << "]" << std::endl
            << "k = " << numberOfNeighbors << std::endl;
  std::cout
    << "measurement vector : distance_by_distMetric : distance_by_tree"
    << std::endl;
  std::vector<double> distances4(numberOfNeighbors);
  for (unsigned int i = 0; i < numberOfNeighbors; ++i)
  {
    distances4[i] = distanceMetric->Evaluate(
      centroidTree->GetMeasurementVector(neighbors[i]));
    std::cout << "[" << centroidTree->GetMeasurementVector(neighbors[i])
              << "]       :       " << distances4[i] << "            :       "
              << distances3[i] << std::endl;
    if (itk::Math::NotAlmostEquals(distances2[i], distances1[i]))
    {
      std::cerr << "Mismatched distance values by centroid tree."
                << std::endl;
      return EXIT_FAILURE;
    }
  }
  // Software Guide : EndCodeSnippet


  // Software Guide : BeginLatex
  //
  // KdTree also supports searching points within a hyper-spherical
  // kernel. We specify the radius and call the \code{Search()} method.  In
  // the case of the KdTree, this is done with the following lines of code.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  double radius = 437.0;

  tree->Search(queryPoint, radius, neighbors);

  std::cout << "\nSearching points within a hyper-spherical kernel:"
            << std::endl;
  std::cout << "*** kd-tree radius search result:" << std::endl
            << "query point = [" << queryPoint << "]" << std::endl
            << "search radius = " << radius << std::endl;
  std::cout << "measurement vector : distance" << std::endl;
  for (auto neighbor : neighbors)
  {
    std::cout << "[" << tree->GetMeasurementVector(neighbor) << "] : "
              << distanceMetric->Evaluate(
                   tree->GetMeasurementVector(neighbor))
              << std::endl;
  }
  // Software Guide : EndCodeSnippet


  // Software Guide : BeginLatex
  //
  // In the case of the centroid KdTree, the \code{Search()} method is used as
  // illustrated by the following code.
  //
  // Software Guide : EndLatex

  // Software Guide : BeginCodeSnippet
  centroidTree->Search(queryPoint, radius, neighbors);
  std::cout << "\n*** Weighted centroid kd-tree radius search result:"
            << std::endl
            << "query point = [" << queryPoint << "]" << std::endl
            << "search radius = " << radius << std::endl;
  std::cout << "measurement vector : distance" << std::endl;
  for (auto neighbor : neighbors)
  {
    std::cout << "[" << centroidTree->GetMeasurementVector(neighbor) << "] : "
              << distanceMetric->Evaluate(
                   centroidTree->GetMeasurementVector(neighbor))
              << std::endl;
  }
  // Software Guide : EndCodeSnippet
  return EXIT_SUCCESS;
}