GNN.h

#ifndef MTF_GNN_H
#define MTF_GNN_H

#include "mtf/SM/GNNParams.h"

#include <vector>
#include <memory>

_MTF_BEGIN_NAMESPACE

namespace gnn{
    struct Node{
        VectorXi nns_inds;
        int size;
        int capacity;
    };
    struct IndxDist{
        double dist;
        int idx;
    };
    inline int cmpQsort(const void *a, const void *b){
        IndxDist *a1 = (IndxDist *)a;
        IndxDist *b1 = (IndxDist *)b;
        // Ascending
        if(a1->dist > b1->dist) return 1;
        if(a1->dist == b1->dist) return 0;
        return -1;
    }

    template<class DistType>
    class GNN{
    public:

        typedef GNNParams ParamType;
        typedef std::shared_ptr<const DistType> DistTypePtr;

        GNN(DistTypePtr _dist_func, int _n_samples, int _n_dims,
            bool _is_symmetrical = true, const ParamType *gnn_params = nullptr);
        ~GNN(){}
        void computeDistances(const double *dataset);
        void buildGraph(const double *dataset);
        void searchGraph(const double *query, const double *dataset, 
            int *nn_ids, double *nn_dists, int K = 1);
        void saveGraph(const char* file_name);
        void loadGraph(const char* file_name);

        void  buildGraph(const double *X, int k);
        int searchGraph(const double *Xq, const double *X, 
            int NNs, int K);

    protected:

        DistTypePtr dist_func;
        const int n_samples, n_dims;
        const bool is_symmetrical;
        ParamType params;
        std::vector<Node> nodes;
        MatrixXd dataset_distances;

        int start_node_idx;
        bool dist_computed;

        int getRandNum(int lb, int ub){
            //  time_t sec;
            //  time(&sec);
            //  srand((unsigned int) sec);
            return (rand() % (ub - lb + 1) + lb);
        }
        template<typename ScalarT>
        void swap(ScalarT *i, ScalarT *j){
            ScalarT temp;
            temp = *i;
            *i = *j;
            *j = temp;
        }
        void knnSearch2(const double *Q, IndxDist *dists, const double *X,
            int rows, int cols, int k);
        void knnSearch11(const double *Q, IndxDist *dists, const double *X, int rows,
            int cols, int k, int *X_inds);

        int min(int a, int b){ return a < b ? a : b; }

        void pickKNNs(IndxDist *vis_nodes, int visited, IndxDist **gnn_dists,
            int K, int *gnns_cap);

        void addNode(Node *node_i, int nn);
    };
}
_MTF_END_NAMESPACE

#endif