// ================================================================
// $Id$
// ****************************************************************
// 
// Initial coding by
/**
   @author Jiri Bittner
*/

#ifndef __VSPKDTREE_H__
#define __VSPKDTREE_H__

// Standard headers
#include <iomanip>
#include <vector>
#include <functional>
#include <stack>


// User headers
#include "VssRay.h"
#include "AxisAlignedBox3.h"


#define USE_KDNODE_VECTORS 1
#define _RSS_STATISTICS
#define _RSS_TRAVERSAL_STATISTICS


#include "Statistics.h"
#include "Ray.h"

#include "RayInfo.h"
#include "Containers.h"

/**
	Static statistics for vsp kd-tree search
*/
class VspKdStatistics:  public StatisticsBase
{
public:  
	// total number of nodes
	int nodes;
	// number of splits along each of the axes
	int splits[7];
	// totals number of rays
	int rays;
	// initial size of the pvs
	int initialPvsSize;
	// total number of query domains
	int queryDomains;
	// total number of ray references
	int rayRefs;

	// max depth nodes
	int maxDepthNodes;
	// max depth nodes
	int minPvsNodes;
	// nodes with minimum PVS
	int minRaysNodes;
	// max ray contribution nodes
	int maxRayContribNodes;
	// max depth nodes
	int minSizeNodes;
	
	// max number of rays per node
	int maxRayRefs;
	// maximal PVS size / leaf
	int maxPvsSize;

	// number of dynamically added ray refs
	int addedRayRefs;
	// number of dynamically removed ray refs
	int removedRayRefs;
  
	/** Default constructor.
	*/
	VspKdStatistics() {	Reset(); }
	int Nodes() const { return nodes; }
	int Interior() const { return nodes / 2; }
	int Leaves() const { return (nodes / 2) + 1; }

	void Reset() 
	{
		nodes = 0;

		for (int i=0; i<7; i++)
			splits[i] = 0;

		rays = queryDomains = 0;
		rayRefs = 0;
		
		maxDepthNodes = 0;
		minPvsNodes = 0;
		minRaysNodes = 0;
		maxRayContribNodes = 0;
		minSizeNodes = 0;

		maxRayRefs = 0;
		addedRayRefs = removedRayRefs = 0;
		
		initialPvsSize = 0;
		maxPvsSize = 0;
	}

	void Print(ostream &app) const;
	friend ostream &operator<<(ostream &s, const VspKdStatistics &stat) 
	{
	    stat.Print(s);
	    return s;
	}  
};


class VspKdTreeInterior;


/** Abstract superclass of Vsp-Kd-Tree nodes.
*/
class VspKdTreeNode
{
public:

	friend class VspKdTree;

	enum {EInterior, ELeaf};

	/** Constructs new interior node from the parent node.
	*/
	inline VspKdTreeNode(VspKdTreeInterior *p);

	/** Destroys this node and the subtree.
	*/
	virtual ~VspKdTreeNode();

	/** Type of the node (Einterior or ELeaf)
	*/
	virtual int Type() const = 0;
 
	/** Returns true if this node is a leaf.
	*/
	bool IsLeaf() const;
	
	/** Prints node stats.
	*/
	virtual void Print(ostream &s) const = 0;

	/** Returns time needed to access this node.
	*/
	virtual int GetAccessTime(); // NOTE: don't really know how it is used!

	/** Returns parent node.
	*/
	VspKdTreeInterior *GetParent() const;

	/** Sets parent node.
	*/
	void SetParent(VspKdTreeInterior *p);

protected:
	/////////////////////////////////
	// The actual data goes here
  
	/// link to the parent
	VspKdTreeInterior *mParent;

	enum {SPLIT_X = 0, SPLIT_Y, SPLIT_Z, SPLIT_DIRX, SPLIT_DIRY, SPLIT_DIRZ};
  
	/// splitting axis
	char mAxis;
	
	/// depth
	unsigned char mDepth;
};

// --------------------------------------------------------------
// KD-tree node - interior node
// --------------------------------------------------------------
class VspKdTreeInterior: public VspKdTreeNode
{
public:
	friend class VspKdTree;

	/** Constructs new interior node from parent node.
	*/
	VspKdTreeInterior(VspKdTreeInterior *p);
			
	virtual int GetAccessTime();
	
	virtual int Type() const;
  
	virtual ~VspKdTreeInterior();
    
	virtual void Print(ostream &s) const;

	/** Returns back child.
	*/
	inline VspKdTreeNode *GetBack() const;
	/** Returns front child.
	*/
	inline VspKdTreeNode *GetFront() const;

protected:

	/** Sets pointers to back child and front child.
	*/
	void SetupChildLinks(VspKdTreeNode *b, VspKdTreeNode *f);

	/** Replaces the pointer to oldChild with a pointer to newChild.
	*/
	void ReplaceChildLink(VspKdTreeNode *oldChild, VspKdTreeNode *newChild);
 
	/** Computes intersection of the ray with the node boundaries.
	*/
	int ComputeRayIntersection(const RayInfo &rayData, float &t);

	// plane in local modelling coordinates
	float mPosition;

	// pointers to children
	VspKdTreeNode *mBack;
	VspKdTreeNode *mFront;

	// the bbox of the node
	AxisAlignedBox3 mBox;

	// data for caching
	long mAccesses;
	long mLastAccessTime;
};


// --------------------------------------------------------------
// KD-tree node - leaf node
// --------------------------------------------------------------
class VspKdTreeLeaf: public VspKdTreeNode
{
public:
	friend class VspKdTree;

	/** Constructs leaf from parent node.
		@param p the parent node
		@param nRays preallocates memory to store this number of rays
	*/
	VspKdTreeLeaf(VspKdTreeInterior *p,	const int nRays);
	
	virtual ~VspKdTreeLeaf();

	virtual int Type() const;  

	virtual void Print(ostream &s) const;
  
	/** Adds a ray to the leaf ray container.
	*/
	void AddRay(const RayInfo &data);
	/** Returns size of the leaf PVS as induced by the rays
		@note returns precomputed PVS size, which may not be valid
		anymore. Run UpdatePvsSize to get a valid PVS.
	*/
	int GetPvsSize() const;

	/** If PVS is not valid, this function recomputes the leaf 
		PVS as induced by the rays.
	*/
	void UpdatePvsSize();

	/** Returns stored rays.
	*/
	RayInfoContainer &GetRays();

	/** Returns rays into this ray container.
	*/
	void GetRays(VssRayContainer &rays);
	/** Returns average contribution of a ray to the PVS
	*/
	inline float GetAvgRayContribution() const;
	/** Returns square of average contribution of a ray.
	*/
	inline float GetSqrRayContribution() const;

	/** Extracts PVS from ray set.
	*/
	void ExtractPvs(ObjectContainer &objects) const;

	//-- mailing options
	void Mail();

	bool Mailed() const;
  
	bool Mailed(const int mail);

	static void NewMail();

	static int mailID;


protected:

	/** Manually sets PVS size.
		@param s the PVS size
	*/
	void SetPvsSize(const int s);

	int mMailbox;
  
	RayInfoContainer mRays;
	bool mValidPvs;
	
private:
	int mPvsSize;
};



// ---------------------------------------------------------------
// Main LSDS search class
// ---------------------------------------------------------------
class VspKdTree
{
	// --------------------------------------------------------------
	// For sorting rays
	// -------------------------------------------------------------
	struct SortableEntry
	{
		enum EType 
		{
			ERayMin,
			ERayMax
		};

		int type;
		float value;
		void *data;
  
		SortableEntry() {}
		SortableEntry(const int t, const float v, void *d):type(t),
					  value(v), data(d) 
		{
		}
		
		friend bool operator<(const SortableEntry &a, const SortableEntry &b) 
		{
			return a.value < b.value;
		}
	};

	struct TraversalData
	{  
		VspKdTreeNode *mNode;
		AxisAlignedBox3 mBox;
		int mDepth;
		//float mPriority;
        
		TraversalData() {}

		TraversalData(VspKdTreeNode *n, const float p):
		mNode(n)//, mPriority(p)
		{}

		TraversalData(VspKdTreeNode *n,	const AxisAlignedBox3 &b, const int d):
		mNode(n), mBox(b), mDepth(d) {}
    
		// comparator for the priority queue
		/*struct less_priority : public binary_function<const TraversalData, const TraversalData, bool> 
		{
			bool operator()(const TraversalData a, const TraversalData b) {			
			return a.mPriority < b.mPriority;		}	};*/

		//    ~TraversalData() {}
		//    TraversalData(const TraversalData &s):node(s.node), bbox(s.bbox), depth(s.depth) {}
    
		friend bool operator<(const TraversalData &a, const TraversalData &b) 
		{
			// return a.node->queries.size() < b.node->queries.size();
			VspKdTreeLeaf *leafa = (VspKdTreeLeaf *) a.mNode;
			VspKdTreeLeaf *leafb = (VspKdTreeLeaf *) b.mNode;
#if 0
			return
				leafa->rays.size()*a.mBox.GetVolume()
				<
				leafb->rays.size()*b.mBox.GetVolume();
#endif
#if 1
			return
				leafa->GetPvsSize()*a.mBox.GetVolume()
				<
				leafb->GetPvsSize()*b.mBox.GetVolume();
#endif
#if 0
			return
				leafa->GetPvsSize()
				<
				leafb->GetPvsSize();
#endif
#if 0
			return
				leafa->GetPvsSize() / (float)(leafa->rays.size() + Limits::Small())
				>
				leafb->GetPvsSize() / (float)(leafb->rays.size() + Limits::Small());
#endif
#if 0
			return
				leafa->GetPvsSize() * (float)leafa->rays.size()
				<
				leafb->GetPvsSize() * (float)leafb->rays.size();
#endif
    }
  };
  
  /** Simplified data for ray traversal only.
  */
  struct RayTraversalData 
  {
	  RayInfo mRayData;
	  VspKdTreeNode *mNode;
      
	  RayTraversalData() {}
	  
	  RayTraversalData(VspKdTreeNode *n, const RayInfo &data):
	  mRayData(data), mNode(n) {}
  };
	
public:

	VspKdTree();
	virtual ~VspKdTree();

	virtual void Construct(const VssRayContainer &rays,
						   AxisAlignedBox3 *forcedBoundingBox = NULL);

	/** Returns bounding box of the specified node.
	*/
	AxisAlignedBox3 GetBBox(VspKdTreeNode *node) const;

	const VspKdStatistics &GetStatistics() const;

	/** Get the root of the tree.
	*/
	VspKdTreeNode *GetRoot() const;

	/** Returns average PVS size in tree.
	*/
	float GetAvgPvsSize();

	/** Returns memory usage in MB.
	*/
	float GetMemUsage() const;
	
	float GetRayMemUsage() const;

	/** Collects leaves of this tree.
	*/
	void CollectLeaves(vector<VspKdTreeLeaf *> &leaves) const;

protected:

	// incremental construction
	virtual void UpdateRays(VssRayContainer &remove, VssRayContainer &add);

	virtual void AddRays(VssRayContainer &add);

	VspKdTreeNode *Locate(const Vector3 &v);
	
	VspKdTreeNode *SubdivideNode(VspKdTreeLeaf *leaf,
				 				 const AxisAlignedBox3 &box,
								 AxisAlignedBox3 &backBox,
								 AxisAlignedBox3 &frontBox);
	
	VspKdTreeNode *Subdivide(const TraversalData &tdata);
	
	int SelectPlane(VspKdTreeLeaf *leaf,
					const AxisAlignedBox3 &box,
					float &position,
					int &raysBack,
					int &raysFront,
					int &pvsBack,
					int &pvsFront);

	void SortSplitCandidates(VspKdTreeLeaf *node,
							 const int axis);	
  
	float BestCostRatioHeuristic(VspKdTreeLeaf *node,
								 int &axis,
								 float &position,
								 int &raysBack,
								 int &raysFront,
								 int &pvsBack,
								 int &pvsFront);

	float BestCostRatioRegular(VspKdTreeLeaf *node, 
							   int &axis,
							   float &position,
							   int &raysBack,
							   int &raysFront,
							   int &pvsBack,
							   int &pvsFront);
	
	float EvalCostRatio(VspKdTreeLeaf *node,
						const int axis,
						const float position,
						int &raysBack,
						int &raysFront,
						int &pvsBack,
						int &pvsFront);


	VspKdTreeNode * SubdivideLeaf(VspKdTreeLeaf *leaf,
								  const float SAThreshold);

	void RemoveRay(VssRay *ray,
				   vector<VspKdTreeLeaf *> *affectedLeaves,
				   const bool removeAllScheduledRays);

	void AddRay(VssRay *ray);
	
	void TraverseInternalNode(RayTraversalData &data,
							 stack<RayTraversalData> &tstack);

	void EvaluateLeafStats(const TraversalData &data);


	int	GetRootPvsSize() const;
	
	int	GetPvsSize(VspKdTreeNode *node, const AxisAlignedBox3 &box) const;

	void GetRayContributionStatistics(float &minRayContribution,
									  float &maxRayContribution,
									  float &avgRayContribution);

	int GenerateRays(const float ratioPerLeaf,
					 SimpleRayContainer &rays);

	/** Collapses this subtree and releases the memory.
		@returns number of collapsed rays.
	*/
	int CollapseSubtree(VspKdTreeNode *sroot, const int time);
	
	int ReleaseMemory(const int time);

	bool TerminationCriteriaMet(const VspKdTreeLeaf *leaf, 
								const AxisAlignedBox3 &box) const;

	/** Computes PVS size of this node given a global ray set.
		@returns PVS size of the intersection of this node bounding box with the rays.
	*/
	int ComputePvsSize(VspKdTreeNode *node, const RayInfoContainer &globalRays) const;

protected:
	/////////////////////////////
	// The core pointer
	VspKdTreeNode *mRoot;
  
	/////////////////////////////
	// Basic properties

	// total number of nodes of the tree
	int mNumNodes;
	
	// axis aligned bounding box of the scene
	AxisAlignedBox3 mBox;

	// epsilon used for the construction
	float mEpsilon;

	// ratio between traversal and intersection costs
	float mCtDivCi;
	
	// type of the splitting to use for the tree construction
	enum {ESplitRegular, ESplitHeuristic };
	int splitType;
	
	// maximum alovable memory in MB
	float mMaxTotalMemory;

	// maximum alovable memory for static kd tree in MB
	float mMaxStaticMemory;

	// this is used during the construction depending
	// on the type of the tree and queries...
	float mMaxMemory;

	// minimal acess time for collapse
	int mAccessTimeThreshold;

	// minimal depth at which to perform collapse
	int mMinCollapseDepth;
	
	// reusable array of split candidates
	vector<SortableEntry> *mSplitCandidates;
	
	/////////////////////////////
	// Construction parameters

	/// max depth of the tree
	int mTermMaxDepth;

	/// minimal ratio of the volume of the cell and the query volume
	float mTermMinSize;

	/// minimal pvs per node to still get subdivided
	int mTermMinPvs;

	/// minimal ray number per node to still get subdivided
	int mTermMinRays;
	
	/// maximal cost ration to subdivide a node
	float mTermMaxCostRatio;
	
	/// maximal contribution per ray to subdivide the node
	float mTermMaxRayContribution;


	bool mOnlyDrivingAxis;
	/////////////////////////////
	VspKdStatistics mStat;	
};


#endif // __LSDS_KDTREE_H__