
// ##################################################
// ##    kd-tree - vicerozmerna datova struktura   ##
// ##                                              ##
// ##################################################

#pragma once
#ifndef KD_TREE_H
#define KD_TREE_H

#include "common.h"

#include <memory>
#include <vector>
#include <algorithm>

template<typename Point, typename DistanceMeasure>
class k2tree
{
private:
    // --------------------------------

    /* uzel */
    struct k2_node;
    struct k2_node 
    {
        unsigned int dimension;
        DistanceMeasure splittingPoint;
        Point* data;

        k2_node* leftSubtree;
        k2_node* rightSubtree;

        k2_node(Point* d, DistanceMeasure splittingBarrier, int axis_)
        {
            splittingPoint = splittingBarrier;
            data = d, dimension = axis_;
        }

        k2_node()  { leftSubtree = NULL, rightSubtree = NULL; };
		~k2_node() { delete leftSubtree; delete rightSubtree; };
    };

    // --------------------------------
    // --------------------------------

    /* porovnavaci objekt podle dimenze */
    struct comparator
    {
        int dim;

        comparator(int dimension)
        {
            dim = dimension;
        }

        bool operator()(Point* a, Point* b)
        {
            return dim == 0 ? a->getPosition().x < b->getPosition().x :
                              a->getPosition().y < b->getPosition().y;
        }
    };

    // --------------------------------
    // --------------------------------

    k2_node* root;

    const inline DistanceMeasure pt2dm(Point* pt, int d)
    {
        return d == 0 ? pt->getPosition().x : pt->getPosition().y;
    }

    // --------------------------------

    /* rozdeli body podle dane dimenze a bariery do dvou mnozin */
    void split(std :: vector<Point*>& source, std :: vector<Point*>& left, std :: vector<Point*>& right, 
        int dimension, DistanceMeasure splittingPoint)
    {
        for (std :: vector<Point*> :: iterator iter = source.begin(); iter != source.end(); iter++)
        {
            if (pt2dm(*iter, dimension) < splittingPoint) 
                left.push_back(*iter);
            else
                right.push_back(*iter);
        }
    }

    // --------------------------------

    k2_node* __buildTree(std :: vector<Point*>& items, const int dimension)
    {
        /* jestlize mame prazdny seznam, nedelame nic */
        if (items.size() == 0) return NULL;

        /* setridime prvky podle dimenze */
        sort(items.begin(), items.end(), comparator(dimension));

        /* najdeme rozdelovaci prvek podle dimenze */
        Point *current = items[items.size() / 2];
        items.erase(items.begin() + (items.size() / 2));
        DistanceMeasure barrier = pt2dm(current, dimension);
        
        /* vytvorime novy uzel */
        k2_node *newNode = new k2_node(current, barrier, dimension);

        /* body podle delici nadroviny */
        std :: vector<Point*> left, right;
        split(items, left, right, dimension, barrier);

        /* rekurzivne vytvor podstromy k2-stromu */
        newNode->leftSubtree  = __buildTree(left,  dimension ^ 1);
        newNode->rightSubtree = __buildTree(right, dimension ^ 1);

        return newNode;
    }

    // --------------------------------

    /* doda spravne deti */
    void getChildren(k2_node* parent, const vector_2d<DistanceMeasure>& point, k2_node **nearnd, k2_node **farnd)
    {
        DistanceMeasure spPoint = (parent->dimension == 0) ? point.x : point.y;

        if (spPoint < parent->splittingPoint)
        {
            *farnd  = parent->rightSubtree;
            *nearnd = parent->leftSubtree;
        } else
        {
            *nearnd = parent->rightSubtree;
            *farnd  = parent->leftSubtree;
        }
    }

    // --------------------------------

    /* absolutni hodnota */
    DistanceMeasure my_abs(const DistanceMeasure val)
    {
        return val >= 0 ? val : -val;
    }

    // --------------------------------

    /* Lua-NN algoritmus z wikipedie, http://en.wikipedia.org/wiki/Kd-tree */
    k2_node* _nearest(const vector_2d<DistanceMeasure>& point, k2_node* current, k2_node* best)
    {
        if (current == NULL)
            return best;

        if (best == NULL)
            best = current;

        /* pokud je vzdalenost tohoto uzlu lepsi, nez vzdalenost nejlepsiho, tak ho prepis */
        DistanceMeasure currentPointDist = (point - current->data->getPosition()).squaredNorm();
        DistanceMeasure bestDist         = (point - best->data->getPosition()).squaredNorm();
        if (currentPointDist < bestDist)
        {
            best = current;
        }

        /* prohledame deti */
        k2_node* nearChild;
        k2_node* farChild;
        getChildren(current, point, &nearChild, &farChild);

        best = _nearest(point, nearChild, best);

        /* zjistime, zda musime prohledat vzdaleneho potomka a pokud ano, prohledame ho */
        DistanceMeasure da = (current->dimension == 0) ? my_abs((point.x - current->data->getPosition().x)) :
                                                         my_abs((point.y - current->data->getPosition().y));
                  bestDist = (point - best->data->getPosition()).squaredNorm();
        if (da < bestDist)
        {
            best = _nearest(point, farChild, best);
        }

        return best;
    }

    // --------------------------------

    bool isInRect(Point* data, const DistanceMeasure xmin, const DistanceMeasure xmax, 
        const DistanceMeasure ymin, const DistanceMeasure ymax)
    {
        return (data->getPosition().x >= xmin) &&
            (data->getPosition().x <= xmax) &&
            (data->getPosition().y >= ymin) &&
            (data->getPosition().y <= ymax);
    }

    // --------------------------------

    void _region(k2_node* current, const DistanceMeasure xmin, const DistanceMeasure xmax, 
        const DistanceMeasure ymin, const DistanceMeasure ymax,
        std :: list<Point*>& result)
    {
        if (current == NULL)
            return;

        if (isInRect(current->data, xmin, xmax, ymin, ymax))
            result.push_back(current->data);

        const bool dimIsX = current->dimension == 0;
        const DistanceMeasure dmin = dimIsX ? xmin : ymin;
        const DistanceMeasure dmax = dimIsX ? xmax : ymax;
        const DistanceMeasure d    = dimIsX ? current->data->getPosition().x : current->data->getPosition().y;

        if (dmin <= d && dmax >= d)
        {
            _region(current->leftSubtree,  xmin, xmax, ymin, ymax, result);
            _region(current->rightSubtree, xmin, xmax, ymin, ymax, result);
        } else
        {
            if (dmax <= d)
            {
                _region(current->leftSubtree,  xmin, xmax, ymin, ymax, result);
            }
            if (dmin >= d)
            {
                 _region(current->rightSubtree, xmin, xmax, ymin, ymax, result);
            }
        }
    }
public:
	k2tree() { root = NULL;  };

    /* vybuduje kd-strom */
    void buildTree(std :: vector<Point*>& items)
    {
        clear();
        root = __buildTree(items, 0);
    }

    // --------------------------------

    void clear(void)
    {
        delete root; 
        root = NULL;
    }

    // --------------------------------

    Point* nearest(const vector_2d<DistanceMeasure>& pt)
    {
        k2_node* bestNode = _nearest(pt, root, root);
        return bestNode == NULL ? NULL : bestNode->data;
    }

    // --------------------------------

    void region(const DistanceMeasure xmin, const DistanceMeasure xmax, 
        const DistanceMeasure ymin, const DistanceMeasure ymax,
        std :: list<Point*>& result)
    {
        _region(root, xmin, xmax, ymin, ymax, result);
    }
};

#endif 