#include "ctree.h"

void k_nn(point p, int k, node *head, vector <point> &results)
{
	vector <node> cset;

	cset.push_back(*head);

	int left = head->below;
	int last_level = head->level;

	vector <point> leftover;

	while(left > k)
	{
		//consider looping through each level, but ignoring distance check if no new nodes added
		int lev = -2000;		//MIN INT
		for(int i = 0; i < (int)cset.size(); i++)
			for(int j = 0; j < (int)cset[i].children.size(); j++)
				if ((cset[i].children[j]->level > lev) && (cset[i].children[j]->level < last_level))
					lev = cset[i].children[j]->level;

		//if no suitable level is found (no children left), move down one level
		if (lev == -2000)		//MIN INT
			lev = last_level-1;

		//add children from the selected level to the current cover set
		for(int i = 0; i < (int)cset.size(); i++)
		{
			for(int j = 0; j < (int)cset[i].children.size(); j++)
			{
				if (cset[i].children[j]->level == lev)
				{
					cset[i].below -= cset[i].children[j]->below;
					cset.push_back(*cset[i].children[j]);
				}
			}
		}

		//prune nodes at distance > d + base ^ level
		float d = set_dist(p, cset);
		for (int i = 0; i < (int)cset.size(); i++)
		{
			if (dist(p, cset[i].pt) > (d + pow(COVER_BASE, lev)))
			{
				left -= cset[i].below;
				cset[i].below = 0;		//below set to 0 flags the node for pruning
			}
		}

		//prune nodes that cannot contain the nearest neighbor
		for (int i = 0; i < (int)cset.size(); i++)
		{
			if (cset[i].below == 0)
			{
				if (left < k)
				{
					add_points(leftover, &cset[i], lev);
				}

				cset.erase(cset.begin()+i);
				i--;
			}
		}

		//done pruning
		if (left < k)
		{
			//naive NN to find remaining neighbors in leftover
			for (int i = 0; i < (k - left); i++)
			{
				float min_d = 100000000;	//MAX FLOAT
				int closest = -1;
				for (int j = 0; j < (int) leftover.size(); j++)
				{
					if (dist(p, leftover[j]) < min_d) closest = j;
				}

				results.push_back(leftover[closest]);
				leftover.erase(leftover.begin()+closest);
			}
		}

		last_level = lev;
	}

	for (int i = 0; i < (int) cset.size(); i++)
	{
		add_points(results, &cset[i], last_level);
	}
}

bool c_insert(point p, node *parent, node *head)
{
	bool success = false;

	float d = dist(p, parent->pt);

	//find level where d > base^level and d < base^(level+1)
	int lev = parent->level;
	if (d >= pow(COVER_BASE, lev)) return false;
	while(d < pow(COVER_BASE, lev)) lev--;

	//create the cover set for this level
	vector <point> cset;
	cover_set(cset, lev, head);

	//if all points in the cover set are at distance > base^level, insert
	if (set_dist(p, cset) > pow(COVER_BASE, lev))
	{
		node *child = new node;
		child->pt = p;
		child->count = 1;
		child->below = 1;
		child->level = lev;
		parent->children.push_back(child);
		parent->below++;
		success = true;
	}
	else if (set_dist(p, cset) == 0)
	{
		//increment count for appropriate node from cover set here
		//success = true;
	}
	else
	{
		for (int i = 0; i < (int)parent->children.size(); i++)
		{
			if (c_insert(p, parent->children[i], head))
			{
				success = true;
				parent->below++;
				break;
			}
		}
	}

	return success;
}

void cover_set(vector <point> &set, int lev, node *head)
{
	if (head->level >= lev)
	{
		set.push_back(head->pt);
		for (int i = 0; i < (int)head->children.size(); i++)
			cover_set(set, lev, head->children[i]);
	}
}

void add_points(vector <point> &set, node *n, int level)
{
	set.push_back(n->pt);
	for (int i = 0; i < (int)n->children.size(); i++)
	{
		if (n->children[i]->level < level) add_points(set, n->children[i], level);
	}
}

float set_dist(point p, const vector <point> &set)
{
	float min = 500000.0;	//MAX FLOAT
	float d;

	for (int i = 0; i < (int)set.size(); i++)
	{
		d = dist(p, set[i]);
		if (d < min) min = d;
	}

	return min;
}

float set_dist(point p, const vector <node> &set)
{
	float min = 500000.0;	//MAX FLOAT
	float d;

	for (int i = 0; i < (int)set.size(); i++)
	{
		d = dist(p, set[i].pt);
		if (d < min) min = d;
	}

	return min;
}

void print_tree(node *head, int lev)
{
	for (int i = 0; i < 100-head->level; i++)
		cout << " ";

	cout << head->level << endl;
	for (int i = 0; i < (int)head->children.size(); i++)
		print_tree(head->children[i], lev+1);
}