#ifndef __KDTREE_HPP
#define __KDTREE_HPP

// (c) Matthew B. Kennel, Institute for Nonlinear Science, UCSD (2004)
//
// Licensed under the Academic Free License version 1.1 found in file LICENSE
// with additional provisions in that same file.
//
// Implement a kd tree for fast searching of points in a fixed data base
// in k-dimensional Euclidean space.

#include <vector>
#include <algorithm>

#include <boost/multi_array.hpp>
#include <boost/array.hpp>

namespace kdtree {
  
  typedef boost::multi_array<double, 2>           KDTreeArray;
  typedef boost::const_multi_array_ref<double, 2> KDTreeROArray;
  
  typedef struct {
    double lower, upper;
  } interval;
  
  // let the compiler know that this is a names of classes.
  class KDTreeNode;
  class SearchRecord;  
  
  struct KDTreeResult {
  public:
    double dis;  // square distance
    int idx;    // neighbor index
  };
  
  class KDTreeResultVector : public std::vector<KDTreeResult> {
    // inherit a std::vector<KDTreeResult>
    // but, optionally maintain it in heap form as a priority
    // queue.
  public:
    
    //
    // add one new element to the list of results, and
    // keep it in heap order.  To keep it in ordinary, as inserted,
    // order, then simply use push_back() as inherited
    // via std::vector<>
    
    void push_element_and_heapify(KDTreeResult&);
    double replace_maxpri_elt_return_new_maxpri(KDTreeResult&);
    
    double max_value();
    // return the distance which has the maximum value of all on list,
    // assuming that ALL insertions were made by
    // push_element_and_heapify()
  };
  
  class KDTree {
  public:
    const KDTreeArray& the_data;
    // "the_data" is a reference to the underlying multi_array of the
    // data to be included in the tree.
    //
    // NOTE: this structure does *NOT* own the storage underlying this.
    // Hence, it would be a very bad idea to change the underlying data
    // during use of the search facilities of this tree.
    // Also, the user must deallocate the memory underlying it.
    
    
    const int N;          // number of data points
    int dim;
    bool sort_results;    // sorting result?
    const bool rearrange; // are we rearranging?
    
  public:

    // constructor, has optional 'dim_in' feature, to use only
    // first 'dim_in' components for definition of nearest neighbors.    
    KDTree(KDTreeArray& data_in, bool rearrange_in = true, int dim_in=-1);
    
    // destructor
    ~KDTree();
    
    
  public:
    
    void n_nearest_brute_force(std::vector<double>& qv, int nn, KDTreeResultVector& result);
    // search for n nearest to a given query vector 'qv' usin
    // exhaustive slow search.  For debugging, usually.
    
    void n_nearest(std::vector<double>& qv, int nn, KDTreeResultVector& result);
    // search for n nearest to a given query vector 'qv'.
    
    void n_nearest_around_point(int idxin, int correltime, int nn,
            KDTreeResultVector& result);
    // search for 'nn' nearest to point [idxin] of the input data, excluding
    // neighbors within correltime
    
    void r_nearest(std::vector<double>& qv, double r2, KDTreeResultVector& result);
    // search for all neighbors in ball of size (square Euclidean distance)
    // r2. Return number of neighbors in 'result.size()',
    
    void r_nearest_around_point(int idxin, int correltime, double r2,
            KDTreeResultVector& result);
    // like 'r_nearest', but around existing point, with decorrelation
    // interval.
    
    int r_count(std::vector<double>& qv, double r2);
    // count number of neighbors within square distance r2.
    
    int r_count_around_point(int idxin, int correltime, double r2);
    // like r_count, c
    
    friend class KDTreeNode;
    friend class SearchRecord;
    
  private:
    
    KDTreeNode* root; // the root pointer
    
    const KDTreeArray* data;
    // pointing either to the_data or an internal
    // rearranged data as necessary
    
    std::vector<int> ind;
    // the index for the tree leaves.  Data in a leaf with bounds [l,u] are
    // in  'the_data[ind[l],*] to the_data[ind[u],*]
    
    KDTreeArray rearranged_data;
    // if rearrange is true then this is the rearranged data storage.

    // no of points per node, standard was 12
    static const int bucketsize = 120;  // global constant.
    
  private:
    void set_data(KDTreeArray& din);
    void build_tree(); // builds the tree.  Used upon construction.
    KDTreeNode* build_tree_for_range(int l, int u, KDTreeNode* parent);
    void select_on_coordinate(int c, int k, int l, int u);
    int select_on_coordinate_value(int c, double alpha, int l, int u);
    void spread_in_coordinate(int c, int l, int u, interval& interv);
  };
  
  class KDTreeNode {
  public:
    KDTreeNode(int dim);
    ~KDTreeNode();
    
  private:
    // visible to self and KDTree.
    friend class KDTree;  // allow kdtree to access private data
    
    int cut_dim;                                 // dimension to cut;
    double cut_val, cut_val_left, cut_val_right;  //cut value
    int l, u;  // extents in index array for searching
    
    std::vector<interval> box; // [min,max] of the box enclosing all points
    
    KDTreeNode *left, *right;  // pointers to left and right nodes.
    
    void search(SearchRecord& sr);
    // recursive innermost core routine for searching..
    
    bool box_in_search_range(SearchRecord& sr);
    // return true if the bounding box for this node is within the
    // search range given by the searchvector and maximum ballsize in 'sr'.
    
    void check_query_in_bound(SearchRecord& sr); // debugging only
    
    // for processing final buckets.
    void process_terminal_node(SearchRecord& sr);
    void process_terminal_node_fixedball(SearchRecord& sr); 
  };
  
} // namespace kdtree

#endif