/*
    A simple C++ wrapper for John Tsiombikas's kd-tree library.
    
    KDTree<PointT, LabelT, Dims> kdtree;
    PointT: 
    - must contain a constructor to recieve points via a double buffer [double*]
    - must have a .data() method that produces a double buffer

    LabelT:
    -  modifiable

    Dims:
    - Number of dimensions in PointT

    Search methods return ResultSet similar to an iterator with more contraints
*/

#pragma once

#include <memory>

namespace ldrs {
    
    #include <kdtree.h>
    
    template <typename _PointT, typename _LabelT, int _Dims>
    class KDTree {

        public:
        
        struct Pair {
            _PointT point;
            _LabelT label;

            typedef std::shared_ptr<Pair> Ptr;
            typedef std::shared_ptr<Pair const> ConstPtr;
        };

        private:

        struct kdtree * _tree;
        size_t _size;

        static void _delete_node_callback(void * node) {
            delete reinterpret_cast<_LabelT*>(node);
        }

        static Pair toPair(struct kdres const * res) {
            double pos[_Dims];
            void * data = kd_res_item(const_cast<struct kdres *>(res), pos);

            Pair pl {
                _PointT(reinterpret_cast<double *>(pos)),
                *reinterpret_cast<_LabelT *>(data)
            };

            return pl;
        }

        public:

        typedef std::shared_ptr<KDTree> Ptr;
        typedef std::shared_ptr<KDTree const> ConstPtr;

        class ResultSet {
            struct kdres * _res;
        public:
            ResultSet(struct kdres * res)
                : _res(res)
            {}

            ~ResultSet() {
                if(_res)
                    kd_res_free(_res);
            }

            inline bool next() { return kd_res_next(_res) != 0; }

            inline bool end() const { return kd_res_end(_res) != 0; }

            inline void rewind() { kd_res_rewind(_res); }
            
            inline int size() const { return kd_res_size(_res); }

            ResultSet(const ResultSet &) = delete;
            ResultSet& operator=(const ResultSet &) = delete;
            ResultSet(ResultSet && rs) {
                _res = rs._res;
                rs._res = nullptr;
            }

            ResultSet& operator=(ResultSet && rs) {
                _res = rs._res;
                rs._res = nullptr;
                return *this;
            }

            ResultSet& operator++() {
                next();
                return *this;
            }

            _PointT point() const {
                double pos[_Dims];
                kd_res_item(_res, pos);
                return _PointT(pos);
            }

            _LabelT& label() const {
                return *reinterpret_cast<_LabelT*>(kd_res_item_data(_res));
            }

            Pair pair() const {
                return KDTree::toPair(_res);
            }
        };

        KDTree()
            : _tree(kd_create(_Dims))
            , _size(0)
        {
            kd_data_destructor(_tree, &_delete_node_callback);
        }

        ~KDTree() {
            kd_free(_tree);
        }

        size_t size() const {
            return _size;
        }

        void clear() {
            _size = 0;
            kd_clear(_tree);
        }

        int insert(_PointT const & pt, _LabelT const & label) {
            _size++;

            return kd_insert(
                _tree,
                pt.data(), 
                reinterpret_cast<void *>(new _LabelT(label))
            );
        }

        ResultSet nearestRange(_PointT const & pt, double range) const {
            kdres * res = kd_nearest_range(_tree, pt.data(), range);
            return ResultSet(res);
        }

        ResultSet nearestPoint(_PointT const & pt) const {
            kdres * res = kd_nearest(_tree, pt.data());
            return ResultSet(res);
        }

        Pair nearestExistentPoint(_PointT const & pt) const {
            kdres * res = kd_nearest(_tree, pt.data());

            Pair pl = KDTree::toPair(res);
            
            kd_res_free(res);

            return pl;
        }

    };
};