//----------------------------------------------------------------------
// File:			kd_pr_search.cpp
// Programmer:		Sunil Arya and David Mount
// Description:		Priority search for kd-trees
// Last modified:	01/04/05 (Version 1.0)
//----------------------------------------------------------------------
// Copyright (c) 1997-2005 University of Maryland and Sunil Arya and
// David Mount.  All Rights Reserved.
// 
// This software and related documentation is part of the Approximate
// Nearest Neighbor Library (ANN).  This software is provided under
// the provisions of the Lesser GNU Public License (LGPL).  See the
// file ../ReadMe.txt for further information.
// 
// The University of Maryland (U.M.) and the authors make no
// representations about the suitability or fitness of this software for
// any purpose.  It is provided "as is" without express or implied
// warranty.
//----------------------------------------------------------------------
// History:
//	Revision 0.1  03/04/98
//		Initial release
//----------------------------------------------------------------------

#include "kd_pr_search.h"				// kd priority search declarations

//----------------------------------------------------------------------
//	Approximate nearest neighbor searching by priority search.
//		The kd-tree is searched for an approximate nearest neighbor.
//		The point is returned through one of the arguments, and the
//		distance returned is the SQUARED distance to this point.
//
//		The method used for searching the kd-tree is called priority
//		search.  (It is described in Arya and Mount, ``Algorithms for
//		fast vector quantization,'' Proc. of DCC '93: Data Compression
//		Conference}, eds. J. A. Storer and M. Cohn, IEEE Press, 1993,
//		381--390.)
//
//		The cell of the kd-tree containing the query point is located,
//		and cells are visited in increasing order of distance from the
//		query point.  This is done by placing each subtree which has
//		NOT been visited in a priority queue, according to the closest
//		distance of the corresponding enclosing rectangle from the
//		query point.  The search stops when the distance to the nearest
//		remaining rectangle exceeds the distance to the nearest point
//		seen by a factor of more than 1/(1+eps). (Implying that any
//		point found subsequently in the search cannot be closer by more
//		than this factor.)
//
//		The main entry point is annkPriSearch() which sets things up and
//		then call the recursive routine ann_pri_search().  This is a
//		recursive routine which performs the processing for one node in
//		the kd-tree.  There are two versions of this virtual procedure,
//		one for splitting nodes and one for leaves. When a splitting node
//		is visited, we determine which child to continue the search on
//		(the closer one), and insert the other child into the priority
//		queue.  When a leaf is visited, we compute the distances to the
//		points in the buckets, and update information on the closest
//		points.
//
//		Some trickery is used to incrementally update the distance from
//		a kd-tree rectangle to the query point.  This comes about from
//		the fact that which each successive split, only one component
//		(along the dimension that is split) of the squared distance to
//		the child rectangle is different from the squared distance to
//		the parent rectangle.
//----------------------------------------------------------------------

//----------------------------------------------------------------------
//		To keep argument lists short, a number of global variables
//		are maintained which are common to all the recursive calls.
//		These are given below.
//----------------------------------------------------------------------

double			ANNprEps;				// the error bound
int				ANNprDim;				// dimension of space
ANNpoint		ANNprQ;					// query point
double			ANNprMaxErr;			// max tolerable squared error
ANNpointArray	ANNprPts;				// the points
ANNpr_queue		*ANNprBoxPQ;			// priority queue for boxes
ANNmin_k		*ANNprPointMK;			// set of k closest points

//----------------------------------------------------------------------
//	annkPriSearch - priority search for k nearest neighbors
//----------------------------------------------------------------------

void ANNkd_tree::annkPriSearch(
	ANNpoint			q,				// query point
	int					k,				// number of near neighbors to return
	ANNidxArray			nn_idx,			// nearest neighbor indices (returned)
	ANNdistArray		dd,				// dist to near neighbors (returned)
	double				eps)			// error bound (ignored)
{
										// max tolerable squared error
	switch(ANN::MetricType)
	{
	case ANN_METRIC0:
		ANNprMaxErr = ANN_POW0(1.0 + eps);
		break;
	case ANN_METRIC1:
		ANNprMaxErr = ANN_POW1(1.0 + eps);
		break;
	case ANN_METRIC2:
		ANNprMaxErr = ANN_POW(1.0 + eps);
		break;
	case ANN_METRICP:
		ANNprMaxErr = ANN_POWp(1.0 + eps);
		break;
	}
	ANN_FLOP(2)							// increment floating ops

	ANNprDim = dim;						// copy arguments to static equivs
	ANNprQ = q;
	ANNprPts = pts;
	ANNptsVisited = 0;					// initialize count of points visited

	ANNprPointMK = new ANNmin_k(k);		// create set for closest k points

										// distance to root box
	ANNdist box_dist = annBoxDistance(q,
				bnd_box_lo, bnd_box_hi, dim);

	ANNprBoxPQ = new ANNpr_queue(n_pts);// create priority queue for boxes
	ANNprBoxPQ->insert(box_dist, root); // insert root in priority queue

	while (ANNprBoxPQ->non_empty() &&
		(!(ANNmaxPtsVisited != 0 && ANNptsVisited > ANNmaxPtsVisited))) {
		ANNkd_ptr np;					// next box from prior queue

										// extract closest box from queue
		ANNprBoxPQ->extr_min(box_dist, (void *&) np);

		ANN_FLOP(2)						// increment floating ops
		if (box_dist*ANNprMaxErr >= ANNprPointMK->max_key())
			break;

		np->ann_pri_search(box_dist);	// search this subtree.
	}

	for (int i = 0; i < k; i++) {		// extract the k-th closest points
		dd[i] = ANNprPointMK->ith_smallest_key(i);
		nn_idx[i] = ANNprPointMK->ith_smallest_info(i);
	}

	delete ANNprPointMK;				// deallocate closest point set
	delete ANNprBoxPQ;					// deallocate priority queue
}

//----------------------------------------------------------------------
//	kd_split::ann_pri_search - search a splitting node
//----------------------------------------------------------------------

void ANNkd_split::ann_pri_search(ANNdist box_dist)
{
	ANNdist new_dist;					// distance to child visited later
										// distance to cutting plane
	ANNcoord cut_diff = ANNprQ[cut_dim] - cut_val;

	if (cut_diff < 0) {					// left of cutting plane
		ANNcoord box_diff = cd_bnds[ANN_LO] - ANNprQ[cut_dim];
		if (box_diff < 0)				// within bounds - ignore
			box_diff = 0;
										// distance to further box
		switch(ANN::MetricType)
		{
		case ANN_METRIC0:
			new_dist = (ANNdist) ANN_SUM0(box_dist, ANN_DIFF0(ANN_POW0(box_diff), ANN_POW0(cut_diff)));
			break;
		case ANN_METRIC1:
			new_dist = (ANNdist) ANN_SUM1(box_dist, ANN_DIFF1(ANN_POW1(box_diff), ANN_POW1(cut_diff)));
			break;
		case ANN_METRIC2:
			new_dist = (ANNdist) ANN_SUM(box_dist, ANN_DIFF(ANN_POW(box_diff), ANN_POW(cut_diff)));
			break;
		case ANN_METRICP:
			new_dist = (ANNdist) ANN_SUMp(box_dist, ANN_DIFFp(ANN_POWp(box_diff), ANN_POWp(cut_diff)));
			break;
		}
		if (child[ANN_HI] != KD_TRIVIAL)// enqueue if not trivial
			ANNprBoxPQ->insert(new_dist, child[ANN_HI]);
										// continue with closer child
		child[ANN_LO]->ann_pri_search(box_dist);
	}
	else {								// right of cutting plane
		ANNcoord box_diff = ANNprQ[cut_dim] - cd_bnds[ANN_HI];
		if (box_diff < 0)				// within bounds - ignore
			box_diff = 0;
										// distance to further box
		switch(ANN::MetricType)
		{
		case ANN_METRIC0:
			new_dist = (ANNdist) ANN_SUM0(box_dist, ANN_DIFF0(ANN_POW0(box_diff), ANN_POW0(cut_diff)));
			break;
		case ANN_METRIC1:
			new_dist = (ANNdist) ANN_SUM1(box_dist, ANN_DIFF1(ANN_POW1(box_diff), ANN_POW1(cut_diff)));
			break;
		case ANN_METRIC2:
			new_dist = (ANNdist) ANN_SUM(box_dist, ANN_DIFF(ANN_POW(box_diff), ANN_POW(cut_diff)));
			break;
		case ANN_METRICP:
			new_dist = (ANNdist) ANN_SUMp(box_dist, ANN_DIFFp(ANN_POWp(box_diff), ANN_POWp(cut_diff)));
			break;
		}
		if (child[ANN_LO] != KD_TRIVIAL)// enqueue if not trivial
			ANNprBoxPQ->insert(new_dist, child[ANN_LO]);
										// continue with closer child
		child[ANN_HI]->ann_pri_search(box_dist);
	}
	ANN_SPL(1)							// one more splitting node visited
	ANN_FLOP(8)							// increment floating ops
}

//----------------------------------------------------------------------
//	kd_leaf::ann_pri_search - search points in a leaf node
//
//		This is virtually identical to the ann_search for standard search.
//----------------------------------------------------------------------

void ANNkd_leaf::ann_pri_search(ANNdist box_dist)
{
	register ANNdist dist;				// distance to data point
	register ANNcoord* pp;				// data coordinate pointer
	register ANNcoord* qq;				// query coordinate pointer
	register ANNdist min_dist;			// distance to k-th closest point
	register ANNcoord t;
	register int d;

	min_dist = ANNprPointMK->max_key(); // k-th smallest distance so far

	for (int i = 0; i < n_pts; i++) {	// check points in bucket

		pp = ANNprPts[bkt[i]];			// first coord of next data point
		qq = ANNprQ;					// first coord of query point
		dist = 0;

		for(d = 0; d < ANNprDim; d++) {
			ANN_COORD(1)				// one more coordinate hit
			ANN_FLOP(4)					// increment floating ops

			t = *(qq++) - *(pp++);		// compute length and adv coordinate
										// exceeds dist to k-th smallest?
			switch(ANN::MetricType)
			{
			case ANN_METRIC0:
				dist = ANN_SUM0(dist, ANN_POW0(t));
				break;
			case ANN_METRIC1:
				dist = ANN_SUM1(dist, ANN_POW1(t));
				break;
			case ANN_METRIC2:
				dist = ANN_SUM(dist, ANN_POW(t));
				break;
			case ANN_METRICP:
				dist = ANN_SUMp(dist, ANN_POWp(t));
				break;
			}
			if( dist > min_dist) break;
		}

		if (d >= ANNprDim &&					// among the k best?
		   (ANN_ALLOW_SELF_MATCH || dist!=0)) { // and no self-match problem
												// add it to the list
			ANNprPointMK->insert(dist, bkt[i]);
			min_dist = ANNprPointMK->max_key();
		}
	}
	ANN_LEAF(1)							// one more leaf node visited
	ANN_PTS(n_pts)						// increment points visited
	ANNptsVisited += n_pts;				// increment number of points visited
}