knnGraphRangeQuery.h

/*
 This Source Code Form is subject to the terms of the Mozilla Public
 License, v. 2.0. If a copy of the MPL was not distributed with this
 file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/
 
#pragma once
 
#include "../../query.h"
#include "../Iterator/knnGraphRangeIterator.h"
 
#include <vector>
#include <stack>
#include <set>
 
namespace Ponca
{
    template <typename Traits>
    class KnnGraphBase;
 
    template <typename Traits>
    class KnnGraphRangeQuery : public RangeIndexQuery<typename Traits::IndexType, typename Traits::DataPoint::Scalar>
    {
    protected:
        using QueryType = RangeIndexQuery<typename Traits::IndexType, typename Traits::DataPoint::Scalar>;
        friend class KnnGraphRangeIterator<Traits>; // This type must be equal to KnnGraphRangeQuery::Iterator
 
    public:
        using DataPoint  = typename Traits::DataPoint;
        using IndexType  = typename Traits::IndexType;
        using Scalar     = typename DataPoint::Scalar;
        using VectorType = typename DataPoint::VectorType;
        using Iterator   = KnnGraphRangeIterator<Traits>;
        using Self       = KnnGraphRangeQuery<Traits>;
 
    public:
        inline KnnGraphRangeQuery(const KnnGraphBase<Traits>* graph, Scalar radius, int index)
            : QueryType(radius, index), m_graph(graph), m_flag(), m_stack()
        {
        }
 
        inline Self& operator()(int index, Scalar radius)
        {
            return QueryType::template operator()<Self>(index, radius);
        }
 
        inline Self& operator()(int index) { return QueryType::template operator()<Self>(index); }
 
        inline Iterator begin()
        {
            QueryType::reset();
            Iterator it(this);
            this->initialize(it);
            this->advance(it);
            return it;
        }
 
        inline Iterator end() { return Iterator(this, static_cast<int>(m_graph->size())); }
 
    protected:
        inline void initialize(Iterator& iterator)
        {
            m_flag.clear();
            m_flag.insert(QueryType::input());
 
            PONCA_DEBUG_ASSERT(m_stack.empty());
            m_stack.push(QueryType::input());
 
            iterator.m_index = -1;
        }
 
        inline void advance(Iterator& iterator)
        {
            const auto& points = m_graph->m_kdTreePoints;
            const auto& point  = points[QueryType::input()].pos();
 
            if (!(iterator != end()))
                return;
 
            if (m_stack.empty())
            {
                iterator = end();
            }
            else
            {
                int idx_current = m_stack.top();
                m_stack.pop();
 
                PONCA_DEBUG_ASSERT((point - points[idx_current].pos()).squaredNorm() < QueryType::squaredRadius());
 
                iterator.m_index = idx_current;
 
                for (int idx_nei : m_graph->kNearestNeighbors(idx_current))
                {
                    PONCA_DEBUG_ASSERT(idx_nei >= 0);
                    Scalar d  = (point - points[idx_nei].pos()).squaredNorm();
                    Scalar th = QueryType::descentDistanceThreshold();
                    if ((point - points[idx_nei].pos()).squaredNorm() < QueryType::descentDistanceThreshold() &&
                        m_flag.insert(idx_nei).second)
                    {
                        m_stack.push(idx_nei);
                    }
                }
                if (iterator.m_index == QueryType::input())
                    advance(iterator); // query is not included in returned set
            }
        }
 
    protected:
        const KnnGraphBase<Traits>* m_graph{nullptr};
        std::set<int> m_flag;    
        std::stack<int> m_stack; 
    };
 
} // namespace Ponca