#ifndef KDTREE_H
#define KDTREE_H

#include <vector>
#include <iostream>
#include <fstream>

#include "scene.h"

using namespace std;

class ShapeList: public vector<Shape*>
{
};

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>
{
};

class Container
{
protected:
	BBox bbox;
public:
	ShapeList shapes;
	Container(): bbox(), shapes() {};
	virtual void addShape(Shape* aShape);
	//void addShapeNoExtend(Shape* aShape) { shapes.push_back(aShape); };
	virtual void optimize() {};
};

class KdNode
{
	float split;
	bool leaf; /* is this node a leaf? */
	short axis; /* 0,1,2 => x,y,z */
	KdNode *children;
public:
	ShapeList shapes;

	KdNode() : leaf(true), axis(0), shapes() {};

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

	void setSplit(float aSplit) { split = aSplit; };
	float getSplit() { return split; };

	void setLeaf(bool aLeaf) { leaf = aLeaf; };
	bool isLeaf() { return leaf; };

	KdNode *getLeftChild() { return children; };
	KdNode *getRightChild() { return children+1; };

	void addShape(Shape* aShape) { shapes.push_back(aShape); };

	void subdivide(BBox bbox, int depth);
};

class KdTree: public Container
{
	KdNode *root;
	bool built;
public:
	KdTree() : Container(), built(false) {};
	void addShape(Shape* aShape) { Container::addShape(aShape); built = false; };
	void optimize(); // build kd-tree
	void save(ostream &str, KdNode *node = NULL);
};

#endif