#include <algorithm>

#include "kdtree.h"

void Container::addShape(Shape* aShape)
{
	shapes.push_back(aShape);
	if (shapes.size() == 0) {
		/* initialize bounding box */
		bbox = aShape->get_bbox();
	} else {
		/* adjust bounding box */
		BBox shapebb = aShape->get_bbox();
		if (shapebb.L.x < bbox.L.x)  bbox.L.x = shapebb.L.x;
		if (shapebb.L.y < bbox.L.y)  bbox.L.y = shapebb.L.y;
		if (shapebb.L.z < bbox.L.z)  bbox.L.z = shapebb.L.z;
		if (shapebb.R.x > bbox.R.x)  bbox.R.x = shapebb.R.x;
		if (shapebb.R.y > bbox.R.y)  bbox.R.y = shapebb.R.y;
		if (shapebb.R.z > bbox.R.z)  bbox.R.z = shapebb.R.z;
	}
};

void KdNode::subdivide(BBox bbox, int depth)
{
	if (depth >= 10 || shapes.size() <= 2)
	{
		leaf = true;
		return;
	}

	// choose split axis
	axis = 0;
	if (bbox.h() > bbox.w() && bbox.h() > bbox.d())
		axis = 1;
	if (bbox.d() > bbox.w() && bbox.d() > bbox.h())
		axis = 2;

	// *** find optimal split position
	SortableShapeList sslist(shapes, axis);
	sort(sslist.begin(), sslist.end());

	SplitList splitlist;

	SortableShapeList::iterator sh;
	for (sh = sslist.begin(); sh != sslist.end(); sh++)
	{
		splitlist.push_back(SplitPos(sh->bbox.L.cell[axis]));
		splitlist.push_back(SplitPos(sh->bbox.R.cell[axis]));
	}
	sort(splitlist.begin(), splitlist.end());
	// remove duplicate splits
	splitlist.resize(unique(splitlist.begin(), splitlist.end()) - splitlist.begin());

	// find all posible splits
	SplitList::iterator spl;
	int rest;
	for (spl = splitlist.begin(); spl != splitlist.end(); spl++)
	{
		for (sh = sslist.begin(), rest = sslist.size(); sh != sslist.end(); sh++, rest--)
		{
			if (sh->hasMark())
			{
				spl->lnum++;
				continue;
			}
			// if shape is completely contained in split plane
			if (spl->pos == sh->bbox.L.cell[axis] == sh->bbox.R.cell[axis])
			{
				if (spl->pos - bbox.L.cell[axis] < bbox.R.cell[axis] - spl->pos)
				{
					// left subcell is smaller -> if not empty, put shape here
					if (spl->lnum)
						spl->lnum++;
					else
						spl->rnum++;
				} else {
					// right subcell is smaller
					if (spl->rnum)
						spl->rnum++;
					else
						spl->lnum++;
				}
				sh->setMark();
			} else
			// if shape is on left side of split plane
			if (sh->bbox.R.cell[axis] <= spl->pos)
			{
				spl->lnum++;
				sh->setMark();
			} else
			// if shape occupies both sides of split plane
			if (sh->bbox.L.cell[axis] < spl->pos && sh->bbox.R.cell[axis] > spl->pos)
			{
				spl->lnum++;
				spl->rnum++;
			} else
			// if shape is on right side of split plane
			if (sh->bbox.L.cell[axis] >= spl->pos)
			{
				spl->rnum += rest;
				break;
			}
		}
	}

	// choose best split pos
	const float K = 1.4; // constant, K = cost of traversal / cost of ray-triangle intersection
	float SAV = 2*(bbox.w()*bbox.h() + bbox.w()*bbox.d() + bbox.h()*bbox.d()); // surface area of node
	float cost = SAV * (K + shapes.size()); // initial cost = non-split cost
	SplitPos *splitpos = NULL;
	leaf = true;
	BBox lbb = bbox;
	BBox rbb = bbox;
	for (spl = splitlist.begin(); spl != splitlist.end(); spl++)
	{
		// calculate SAH cost of this split
		lbb.R.cell[axis] = spl->pos;
		rbb.L.cell[axis] = spl->pos;
		float SAL = 2*(lbb.w()*lbb.h() + lbb.w()*lbb.d() + lbb.h()*lbb.d());
		float SAR = 2*(rbb.w()*rbb.h() + rbb.w()*rbb.d() + rbb.h()*rbb.d());
		float splitcost = K + SAL/SAV*(K+spl->lnum) + SAR/SAV*(K+spl->rnum);

		if (splitcost < cost)
		{
			leaf = false;
			cost = splitcost;
			splitpos = &*spl;
		}
	}

	if (leaf)
		return;

#if 1
// export kd-tree as .obj for visualization
// this would be hard to reconstruct later
	static ofstream F("kdtree.obj");
	Vector3 v;
	static int f=0;
	v.cell[axis] = splitpos->pos;
	v.cell[(axis+1)%3] = bbox.L.cell[(axis+1)%3];
	v.cell[(axis+2)%3] = bbox.L.cell[(axis+2)%3];
	F << "v " << v.x << " " << v.y << " " << v.z << endl;
	v.cell[(axis+1)%3] = bbox.L.cell[(axis+1)%3];
	v.cell[(axis+2)%3] = bbox.R.cell[(axis+2)%3];
	F << "v " << v.x << " " << v.y << " " << v.z << endl;
	v.cell[(axis+1)%3] = bbox.R.cell[(axis+1)%3];
	v.cell[(axis+2)%3] = bbox.R.cell[(axis+2)%3];
	F << "v " << v.x << " " << v.y << " " << v.z << endl;
	v.cell[(axis+1)%3] = bbox.R.cell[(axis+1)%3];
	v.cell[(axis+2)%3] = bbox.L.cell[(axis+2)%3];
	F << "v " << v.x << " " << v.y << " " << v.z << endl;
	F << "f " << f+1 << " " << f+2 << " " << f+3 << " " << f+4 << endl;
	f += 4;
#endif

	split = splitpos->pos;
	float lnum = splitpos->lnum;
	float rnum = splitpos->rnum;

	// split this node
	children = new KdNode[2];
	int state = 0;
	for (sh = sslist.begin(); sh != sslist.end(); sh++)
	{
		// if shape is on left side of split plane
		if (state == 1)
		{ // only right
			children[1].addShape(sh->shape);
			continue;
		}
		if (state == 0)
		{
			if (sh->bbox.R.cell[axis] < split)
			{ // left
				children[0].addShape(sh->shape);
			} else
			if (sh->bbox.R.cell[axis] > split)
			{
				if (sh->bbox.L.cell[axis] < split)
				{ // both
					children[0].addShape(sh->shape);
					children[1].addShape(sh->shape);
				} else
				{ // right
					children[1].addShape(sh->shape);
					state = 1;
				}
			} else
			{ // R == split
				if (sh->bbox.L.cell[axis] < split)
				{ // left
					children[0].addShape(sh->shape);
				} else
				{ // contained
					if (split - bbox.L.cell[axis] < bbox.R.cell[axis] - split)
					{
						// left subcell is smaller -> if not empty, put shape here
						if (lnum)
							children[0].addShape(sh->shape);
						else
							children[1].addShape(sh->shape);
					} else {
						// right subcell is smaller
						if (rnum)
							children[1].addShape(sh->shape);
						else
							children[0].addShape(sh->shape);
					}
				}
			}
		}
	}

	lbb.R.cell[axis] = split;
	rbb.L.cell[axis] = split;

	children[0].subdivide(lbb, depth+1);
	children[1].subdivide(rbb, depth+1);
}

void KdTree::optimize()
{
	root = new KdNode();
	root->shapes = shapes;
	root->subdivide(bbox, 0);
	built = true;
}

void KdTree::save(ostream &str, KdNode *node)
{
	if (!built)
		return;
	if (node == NULL)
		node = root;
	if (node->isLeaf())
		str << "(leaf: " << node->shapes.size() << " shapes)";
	else
	{
		str << "(split " << (char)('x'+node->getAxis()) << " at " << node->getSplit() << "; L=";
		save(str, node->getLeftChild());
		str << "; R=";
		save(str, node->getRightChild());
		str << ";)";
	}
}