#include "pch.h" // use stdafx.h in Visual Studio 2017 and earlier
#include "KDtree.h"

KDtree kdtree;

bool cmp0(const point &a, const point &b)
{
	return a.x[0] < b.x[0];
}
bool cmp1(const point &a, const point &b)
{
	return a.x[1] < b.x[1];
}
bool cmp2(const point &a, const point &b)
{
	return a.x[2] < b.x[2];
}

point::point()
{
	x[0] = x[1] = x[2] = 0;
	id = 0;
}

bool node::operator<(node b) const
{
	return dis < b.dis || (dis == b.dis && id < b.id);
}
node::node()
{
	dis = 0;
	id = 0;
}
node::node(int id1, double dis1)
{
	dis = dis1;
	id = id1;
}

tree::tree()
{
	ls = rs = id = 0;
}

KDtree::KDtree()
{
	X = Y = Z = 0;
	n = tot = root = 0;
}

double KDtree::dis(tree &x)
{
	double P = (x.p.x[0] - X) * (x.p.x[0] - X);
	double Q = (x.p.x[1] - Y) * (x.p.x[1] - Y);
	double O = (x.p.x[2] - Z) * (x.p.x[2] - Z);
	return P + Q + O;
}
double KDtree::mndis(tree &x)
{
	double P = (x.mn[0] - X) * (x.mn[0] - X);
	double M = (x.mx[0] - X) * (x.mx[0] - X);
	if (X >= x.mn[0] && X <= x.mx[0])
		P = M = 0;
	double Q = (x.mn[1] - Y) * (x.mn[1] - Y);
	double N = (x.mx[1] - Y) * (x.mx[1] - Y);
	if (Y >= x.mn[1] && Y <= x.mx[1])
		Q = N = 0;
	double O = (x.mn[2] - Z) * (x.mn[2] - Z);
	double U = (x.mx[2] - Z) * (x.mx[2] - Z);
	if (Z >= x.mn[2] && Z <= x.mx[2])
		O = U = 0;
	return min(P, M) + min(Q, N) + min(O, U);
}
void KDtree::update(int x)
{
	if (!x)
		return;
	int l = t[x].ls, r = t[x].rs;
	if (l)
		t[x].mn[0] = min(t[x].mn[0], t[l].mn[0]),
		t[x].mn[1] = min(t[x].mn[1], t[l].mn[1]),
		t[x].mn[2] = min(t[x].mn[2], t[l].mn[2]),
		t[x].mx[0] = max(t[x].mx[0], t[l].mx[0]),
		t[x].mx[1] = max(t[x].mx[1], t[l].mx[1]),
		t[x].mx[2] = max(t[x].mx[2], t[l].mx[2]);
	if (r)
		t[x].mn[0] = min(t[x].mn[0], t[r].mn[0]),
		t[x].mn[1] = min(t[x].mn[1], t[r].mn[1]),
		t[x].mn[2] = min(t[x].mn[2], t[r].mn[2]),
		t[x].mx[0] = max(t[x].mx[0], t[r].mx[0]),
		t[x].mx[1] = max(t[x].mx[1], t[r].mx[1]),
		t[x].mx[2] = max(t[x].mx[2], t[r].mx[2]);
}

void KDtree::query(int x)
{
	if (!x)
		return;
	double res = dis(t[x]);
	if (res < q.top().dis || (res == q.top().dis && t[x].id < q.top().id))
		q.pop(), q.push(node(t[x].id, res));
	int l = t[x].ls, r = t[x].rs;
	double ld = 0, rd = 0;
	if (l)
		ld = mndis(t[l]);
	if (r)
		rd = mndis(t[r]);
	// cout<<x<<"  "<<ld<<"  "<<rd<<"    "<<res<<"  "<<q.top().id<<" "<<q.top().dis<<endl;
	if (ld < rd)
	{
		if (ld <= q.top().dis)
			query(l);
		if (rd <= q.top().dis)
			query(r);
	}
	else
	{
		if (rd <= q.top().dis)
			query(r);
		if (ld <= q.top().dis)
			query(l);
	}
}

void KDtree::queryd(int x, double d)
{
	if (!x)
		return;
	double res = dis(t[x]);
	// cout<<t[x].p.id<<"  "<<res<<endl;
	if (res <= d * d)
		vec.push_back(t[x].p.id);
	int l = t[x].ls, r = t[x].rs;
	double ld = 0, rd = 0;
	if (l)
		ld = mndis(t[l]);
	if (r)
		rd = mndis(t[r]);
	// cout<<x<<"  "<<ld<<"  "<<rd<<"    "<<res<<"  "<<q.top().id<<" "<<q.top().dis<<endl;
	if (ld < rd)
	{
		if (ld <= d * d)
			queryd(l, d);
		if (rd <= d * d)
			queryd(r, d);
	}
	else
	{
		if (rd <= d * d)
			queryd(r, d);
		if (ld <= d * d)
			queryd(l, d);
	}
}

void KDtree::add(P &s)
{
	n++;
	p[n].x[0] = s.x;
	p[n].x[1] = s.y;
	p[n].x[2] = s.z;
	p[n].id = n;
}

void KDtree::build(int &x, int l, int r, int k)
{
	if (l > r)
		return;
	x = ++tot;
	int mid = (l + r) >> 1;
	// cout<<"stage1  "<<x<<endl;
	if (k == 0)
		nth_element(p + l, p + mid, p + r + 1, cmp0);
	else if (k == 1)
		nth_element(p + l, p + mid, p + r + 1, cmp1);
	else if (k == 2)
		nth_element(p + l, p + mid, p + r + 1, cmp2);
	// cout<<"stage2"<<endl;
	t[x].p = p[mid];
	t[x].id = t[x].p.id;
	t[x].mn[0] = t[x].mx[0] = t[x].p.x[0];
	t[x].mn[1] = t[x].mx[1] = t[x].p.x[1];
	t[x].mn[2] = t[x].mx[2] = t[x].p.x[2];
	build(t[x].ls, l, mid - 1, (k + 1) % 3);
	build(t[x].rs, mid + 1, r, (k + 1) % 3);
	// cout<<"stage3"<<endl;
	update(x);
	// cout<<x<<"  "<<t[x].ls<<"  "<<t[x].rs<<"           "<<t[x].mx[0]<<" "<<t[x].mx[1]<<" "<<t[x].mx[2]<<endl;
}

void KDtree::build()
{
	cout << root << "  " << n << endl;
	build(root, 1, n, 0);
}

// 返回离s最近的k个点的编号
vector<int> KDtree::search_by_k(P &s, int k)
{
	X = s.x;
	Y = s.y;
	Z = s.z;
	while (q.size())
		q.pop();
	for (int j = 1; j <= k; j++)
		q.push(node(0, 1e9));
	query(root);
	vector<int> veck;
	while (q.size())
	{
		if (q.top().id != 0)
			veck.push_back(q.top().id);
		q.pop();
	}
	return veck;
}

// 返回距离s点d以内的所有点的编号
vector<int> KDtree::search_by_dis(P &s, double d)
{
	X = s.x;
	Y = s.y;
	Z = s.z;
	vec.clear();
	queryd(root, d);
	return vec;
}