#include <algorithm>

#include <stack>

#include "common.h"

#include "kdtree.h"

class SortableShape

{

public:

	Shape *shape;

	BBox bbox;

	short axis;

	short mark;

	SortableShape(Shape *aShape, short &aAxis): shape(aShape), axis(aAxis), mark(0)

		{ bbox = shape->get_bbox(); };

	friend bool operator<(const SortableShape& a, const SortableShape& b)

		{ return a.bbox.L.cell[a.axis] < b.bbox.L.cell[b.axis]; };

	void setAxis(short aAxis) { axis = aAxis; };

	void setMark() { mark = 1; };

	short hasMark() { return mark; };

};

class SortableShapeList: public vector<SortableShape>

{

public:

	SortableShapeList(ShapeList *shapes, short axis)

	{

		ShapeList::iterator shape;

		for (shape = shapes->begin(); shape != shapes->end(); shape++)

			push_back(SortableShape(*shape, axis));

	};

};

class SplitPos

{

public:

	Float pos;

	int lnum, rnum;

	SplitPos(): pos(0.0), lnum(0), rnum(0) {};

	SplitPos(Float &aPos): pos(aPos), lnum(0), rnum(0) {};

	friend bool operator<(const SplitPos& a, const SplitPos& b)

		{ return a.pos < b.pos; };

	friend bool operator==(const SplitPos& a, const SplitPos& b)

		{ return a.pos == b.pos; };

};

class SplitList: public vector<SplitPos>

{

};

// stack element for kd-tree traversal

class StackElem

{

public:

	KdNode* node; /* pointer to far child */

	Float t; /* the entry/exit signed distance */

	Vector3 pb; /* the coordinates of entry/exit point */

	StackElem(KdNode *anode, const Float &at, const Vector3 &apb):

		node(anode), t(at), pb(apb) {};

};

// ----------------------------------------

KdNode::~KdNode()

{

	if (isLeaf())

		delete shapes;

	else

		delete[] children;

}

void KdNode::subdivide(BBox bbox, int maxdepth)

{

	if (maxdepth <= 0 || shapes->size() <= 2)

	{

		setLeaf();

		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.H.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.H.cell[axis])

			{

				if (spl->pos - bbox.L.cell[axis] < bbox.H.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.H.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.H.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

	bool leaf = true;

	BBox lbb = bbox;

	BBox rbb = bbox;

	for (spl = splitlist.begin(); spl != splitlist.end(); spl++)

	{

		// calculate SAH cost of this split

		lbb.H.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;

			split = spl->pos;

		}

	}

	if (leaf)

	{

		setLeaf();

		return;

	}

#if 0

// 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.H.cell[(axis+2)%3];

	F << "v " << v.x << " " << v.y << " " << v.z << endl;

	v.cell[(axis+1)%3] = bbox.H.cell[(axis+1)%3];

	v.cell[(axis+2)%3] = bbox.H.cell[(axis+2)%3];

	F << "v " << v.x << " " << v.y << " " << v.z << endl;

	v.cell[(axis+1)%3] = bbox.H.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 this node

	delete shapes;

	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.H.cell[axis] < split)

			{ // left

				children[0].addShape(sh->shape);

			} else

			if (sh->bbox.H.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

			{ // H == split

				if (sh->bbox.L.cell[axis] < split)

				{ // left

					children[0].addShape(sh->shape);

				} else

				{ // contained

					if (split - bbox.L.cell[axis] < bbox.H.cell[axis] - split)

					{

						// left subcell is smaller -> if not empty, put shape here

						if (children[0].shapes->size())

							children[0].addShape(sh->shape);

						else

							children[1].addShape(sh->shape);

					} else {

						// right subcell is smaller

						if (children[1].shapes->size())

							children[1].addShape(sh->shape);

						else

							children[0].addShape(sh->shape);

					}

				}

			}

		}

	}

	lbb.H.cell[axis] = split;

	rbb.L.cell[axis] = split;

	children[0].subdivide(lbb, maxdepth-1);

	children[1].subdivide(rbb, maxdepth-1);

}

void KdTree::build()

{

	infomsg("* building kd-tree\n");

	root = new KdNode();

	ShapeList::iterator shape;

	for (shape = shapes.begin(); shape != shapes.end(); shape++)

		root->addShape(*shape);

	root->subdivide(bbox, max_depth);

	built = true;

}

/* algorithm by Vlastimil Havran, Heuristic Ray Shooting Algorithms, appendix C */

Shape *KdTree::nearest_intersection(const Shape *origin_shape, const Ray &ray,

	Float &nearest_distance)

{

	Float a, b; /* entry/exit signed distance */

	Float t;    /* signed distance to the splitting plane */

	/* if we have no tree, fall back to naive test */

	if (!built)

		return Container::nearest_intersection(origin_shape, ray, nearest_distance);

	if (!bbox.intersect(ray, a, b))

		return NULL;

	/* pointers to the far child node and current node */

	KdNode *farchild, *node;

	node = root; /* start from the kd-tree root node */

	/* std vector is much faster than stack */

	vector<StackElem*> st;

	StackElem *enPt = new StackElem(NULL, a,

		/* distinguish between internal and external origin of a ray*/

		a >= 0.0 ?

			ray.o + ray.dir * a : /* external */

			ray.o);               /* internal */

	/* setup initial exit point in the stack */

	StackElem *exPt = new StackElem(NULL, b, ray.o + ray.dir * b);

	st.push_back(exPt);

	/* loop, traverse through the whole kd-tree, until an object is intersected or ray leaves the scene */

	while (node)

	{

		exPt = st.back();

		/* loop until a leaf is found */

		while (!node->isLeaf())

		{

			/* retrieve position of splitting plane */

			Float splitVal = node->getSplit();

			short axis = node->getAxis();

			if (enPt->pb[axis] <= splitVal)

			{

				if (exPt->pb[axis] <= splitVal)

				{ /* case N1, N2, N3, P5, Z2, and Z3 */

					node = node->getLeftChild();

					continue;

				}

				if (exPt->pb[axis] == splitVal)

				{ /* case Z1 */

					node = node->getRightChild();

					continue;

				}

				/* case N4 */

				farchild = node->getRightChild();

				node = node->getLeftChild();

			}

			else

			{ /* (enPt->pb[axis] > splitVal) */

				if (splitVal < exPt->pb[axis])

				{

					/* case P1, P2, P3, and N5 */

					node = node->getRightChild();

					continue;

				}

				/* case P4*/

				farchild = node->getLeftChild();

				node = node->getRightChild();

			}

			/* case P4 or N4 . . . traverse both children */

			/* signed distance to the splitting plane */

			t = (splitVal - ray.o.cell[axis]) / ray.dir.cell[axis];

			/* setup the new exit point and push it onto stack */

			exPt = new StackElem(farchild, t, Vector3());

			exPt->pb[axis] = splitVal;

			exPt->pb[(axis+1)%3] = ray.o.cell[(axis+1)%3] + t * ray.dir.cell[(axis+1)%3];

			exPt->pb[(axis+2)%3] = ray.o.cell[(axis+2)%3] + t * ray.dir.cell[(axis+2)%3];

			st.push_back(exPt);

		} /* while */

		/* current node is the leaf . . . empty or full */

		/* "intersect ray with each object in the object list, discarding "

		"those lying before stack[enPt].t or farther than stack[exPt].t" */

		Shape *nearest_shape = NULL;

		Float dist = exPt->t;

		ShapeList::iterator shape;

		for (shape = node->shapes->begin(); shape != node->shapes->end(); shape++)

			if (*shape != origin_shape && (*shape)->intersect(ray, dist)

			&& dist >= enPt->t)

				nearest_shape = *shape;

		delete enPt;

		if (nearest_shape)

		{

			while (!st.empty())

			{

				delete st.back();

				st.pop_back();

			}

			nearest_distance = dist;

			return nearest_shape;

		}

		enPt = exPt;

		/* retrieve the pointer to the next node, it is possible that ray traversal terminates */

		node = enPt->node;

		st.pop_back();

	} /* while */

	delete enPt;

	/* ray leaves the scene */

	return NULL;

}

// this should save whole kd-tree with triangles distributed into leaves

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 << ";)";

	}

}

// load kd-tree from file/stream

void KdTree::load(istream &str, KdNode *node)

{

}