#ifndef KD_TREE_H
#define KD_TREE_H

#include <vector>
#include <utility>

#include <vtkIdList.h>

// https://rosettacode.org/wiki/K-d_tree#C.2B.2B

class KdTree {
public:
    KdTree() = default;

    struct Point {
        Point();
        Point(double x, double y, double z);
        Point(double *position);
        double x, y, z;
        double operator[] (std::size_t i) const;
        double dist2(Point const &other) const;

        friend ostream& operator<<(ostream& os, Point const & point);
    };

    using Tuple = std::pair<Point, vtkIdType>;
    void fill(std::vector<Tuple> &points);
    Point query(Point const &point);
    Point query(double *point);

private:
    struct Node {
        Point position;
        vtkIdType index;

        Node *leftChild;
        Node *rightChild;
    };

    std::vector<Node> nodes;
    Node *root;

    double bestDist;
    Node *bestNode;

    Node *fillRec(std::size_t begin, std::size_t end, int axis);
    void queryRec(Node *node, Point const &point, int axis);
};

#endif