// ================================================================
// $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"
#include "ViewCell.h"

class VspKdLeaf;
class ViewCellsManager;
class ViewCellsStatistics;

namespace GtpVisibilityPreprocessor {

/**
	Candidate for leaf merging based on priority.
*/
class VspKdMergeCandidate
{  
public:

	VspKdMergeCandidate(VspKdLeaf *l1, VspKdLeaf *l2);

	/** Computes PVS difference between the leaves.
	*/
	//int ComputePvsDifference() const;

	/** If this merge pair is still valid.
	*/
	bool Valid() const;

	/** Sets this merge candidate to be valid.
	*/
	void SetValid();

	friend bool operator<(const VspKdMergeCandidate &leafa, const VspKdMergeCandidate &leafb)
	{
		return leafb.GetMergeCost() < leafa.GetMergeCost();
	}

	void SetLeaf1(VspKdLeaf *l);
	void SetLeaf2(VspKdLeaf *l);

	VspKdLeaf *GetLeaf1();
	VspKdLeaf *GetLeaf2();

	/** Merge cost of this candidate pair.
	*/
	float GetMergeCost() const;

	/** Returns cost of leaf 1.
	*/
	float GetLeaf1Cost() const;
	/** Returns cost of leaf 2.
	*/
	float GetLeaf2Cost() const;

	/// maximal pvs size
	static int sMaxPvsSize;
	/// overall cost used to normalize cost ratio
	static float sOverallCost;

protected:
	
	/** Evaluates the merge costs of the leaves.
	*/
	void EvalMergeCost();

	int mLeaf1Id;
	int mLeaf2Id;

	float mMergeCost;
	
	VspKdLeaf *mLeaf1;
	VspKdLeaf *mLeaf2;
};



/**
	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[3];
	// 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;

	// max cost ratio
	int maxCostNodes;

	// number of dynamically added ray refs
	int addedRayRefs;
	// number of dynamically removed ray refs
	int removedRayRefs;
  
	/// for average pvs
	int accPvsSize;

	/** 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<3; i++)
			splits[i] = 0;

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

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

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


class VspKdInterior;


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

	friend class VspKdTree;

	enum {EInterior, EIntermediate, ELeaf};

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

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

	/** 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.
		@NOTE: don't really know how it is used!
	*/
	virtual int GetAccessTime(); 

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

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

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

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

// --------------------------------------------------------------
// KD-tree node - interior node
// --------------------------------------------------------------
class VspKdInterior: public VspKdNode
{
public:
	friend class VspKdTree;

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

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

protected:

	/** Sets pointers to back child and front child.
	*/
	void SetupChildLinks(VspKdNode *b, VspKdNode *f);
	/** Replaces the pointer to oldChild with a pointer to newChild.
	*/
	void ReplaceChildLink(VspKdNode *oldChild, VspKdNode *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
	VspKdNode *mBack;
	VspKdNode *mFront;

	// the bbox of the node
	AxisAlignedBox3 mBox;

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


/**
	Node type just before leaf holding abitrary split plane
*/
class VspKdIntermediate: public VspKdNode
{
public:
	friend class VspKdTree;

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

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

protected:

	/** Sets pointers to back child and front child.
	*/
	void SetupChildLinks(VspKdLeaf *b, VspKdLeaf *f);
	/** Computes intersection of the ray with the node boundaries.
	*/
	int ComputeRayIntersection(const RayInfo &rayData, float &t);

	// plane in local modelling coordinates
	Plane3 mSplitPlane;

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

	// the bbox of the node
	AxisAlignedBox3 mBox;

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

// --------------------------------------------------------------
// KD-tree node - leaf node
// --------------------------------------------------------------
class VspKdLeaf: public VspKdNode
{
public:

	friend class VspKdTree;

	/** Constructs leaf from parent node.
		@param p the parent node
		@param nRays preallocates memory to store this number of rays
		@parma maxMisses how many times the max cost ratio was missed on the path to the leaf
	*/
	VspKdLeaf(VspKdNode *p,	const int nRays, const int maxCostMisses = 0);
	
	virtual ~VspKdLeaf();

	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;
  
	/** Returns true if mail equals the leaf mail box.
	*/
	bool Mailed(const int mail) const;

	void SetViewCell(VspKdViewCell *viewCell);

	/** Returns mail box of this leaf.
	*/
	int GetMailbox() const;

	/** Returns view cell associated with this leaf.
	*/
	VspKdViewCell *GetViewCell();

	/** Returns number of times the max cost ratio was missed until
		this leaf.
	*/
	int GetMaxCostMisses();

	////////////////////////////////////////////
	
	static void NewMail();
	static int sMailId;

	ObjectPvs *mPvs;
	float mVolume;

protected:

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

	int mMailbox;
  
	/// rays intersecting this leaf.
	RayInfoContainer mRays;
	/// true if mPvsSize is valid => PVS does not have to be updated
	bool mValidPvs;
	/// the view cell associated with this leaf
	VspKdViewCell *mViewCell;
	/// number of times the max cost ratio was missed on the way to the leaf.
	int mMaxCostMisses;

//private:
	/// stores PVS size so we have to evaluate PVS size only once
	int mPvsSize;
};


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

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

	struct TraversalData
	{  
		VspKdNode *mNode;
		AxisAlignedBox3 mBox;
		//TODO PolygonContainer *mPolys;

		int mDepth;
		//float mPriority;
        
		TraversalData() {}

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

		TraversalData(VspKdNode *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();
			VspKdLeaf *leafa = dynamic_cast<VspKdLeaf *>(a.mNode);
			VspKdLeaf *leafb = dynamic_cast<VspKdLeaf *>(b.mNode);
#if 0
			return
				leafa->rays.size() * a.mBox.GetVolume()
				<
				leafb->rays.size() * b.mBox.GetVolume();
#endif
#if 0
			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
#if 1
			return a.mDepth > b.mDepth;
#endif
		}
	};
  
	/** Simplified data for ray traversal only.
	*/
	struct RayTraversalData 
	{
		RayInfo mRayData;
		VspKdNode *mNode;

		RayTraversalData() {}
	  
		RayTraversalData(VspKdNode *n, const RayInfo &data):
		mRayData(data), mNode(n) {}
	};
	
	struct LineTraversalData 
	{
		VspKdNode *mNode;
		Vector3 mExitPoint;
		
		float mMaxT;
    
		LineTraversalData () {}
		LineTraversalData (VspKdNode *n, const Vector3 &p, const float maxt):
		mNode(n), mExitPoint(p), mMaxT(maxt) {}
	};

public:

	VspKdTree();
	virtual ~VspKdTree();

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

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

	const VspKdStatistics &GetStatistics() const;

	/** Get the root of the tree.
	*/
	VspKdNode *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<VspKdLeaf *> &leaves) const;

	/** Merges view cells created with this tree according to 
		some (global) cost heuristics.
	*/
	int MergeViewCells(const VssRayContainer &rays);

	/** Finds neighbours of this node.
		@param n the input node
		@param neighbours the neighbors of the input node
		@param onlyUnmailed if only unmailed neighbors are collected
	*/
	int FindNeighbors(VspKdLeaf *n, 
					  vector<VspKdLeaf *> &neighbors, 
					  bool onlyUnmailed);

	
	/** Sets pointer to view cells manager.
	*/
	void SetViewCellsManager(ViewCellsManager *vcm);

	/** A ray is cast possible intersecting the tree.
		@param the ray that is cast.
		@returns the number of intersections with objects stored in the tree.
	*/
	int CastLineSegment(const Vector3 &origin,
						const Vector3 &termination,
						vector<ViewCell *> &viewcells);

	/** Collects view cells generated by this tree.
	*/
	void CollectViewCells(ViewCellContainer &viewCells) const;
	
	/** Refines view cells using shuffling, i.e., border leaves 
		of two view cells are exchanged if the resulting view cells
		are tested to be "better" than the old ones.
		@returns number of refined view cells
	*/
	int RefineViewCells(const VssRayContainer &rays);

	/** Collects candidates for the merge in the merge queue.
	*/
	void CollectMergeCandidates();

	/** Collapses the tree with respect to the view cell partition.
		@returns number of collapsed nodes
	*/
	int CollapseTree();

protected:

	/** Collapses the tree with respect to the view cell partition,
		i.e. leaves having the same view cell are collapsed.
		@param node the root of the subtree to be collapsed
		@param collapsed returns the number of collapsed nodes
		@returns node of type leaf if the node could be collapsed, 
		this node otherwise
	*/
	VspKdNode *CollapseTree(VspKdNode *node, int &collapsed);

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

	virtual void AddRays(VssRayContainer &add);

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

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

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


	VspKdNode * SubdivideLeaf(VspKdLeaf *leaf,
								  const float SAThreshold);

	void RemoveRay(VssRay *ray,
				   vector<VspKdLeaf *> *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(VspKdNode *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(VspKdNode *sroot, const int time);
	
	int ReleaseMemory(const int time);

	bool TerminationCriteriaMet(const VspKdLeaf *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(VspKdNode *node, 
					   const RayInfoContainer &globalRays) const;


	/** Generates view cell for this leaf taking the ray contributions.
	*/
	void GenerateViewCell(VspKdLeaf *leaf);

	/** Merges view cells of the two leaves.
	*/
	bool MergeViewCells(VspKdLeaf *l1, VspKdLeaf *l2);
	
	/** Helper function revalidating the view cell leaf list after merge.
	*/
	void RepairViewCellsLeafLists();

	/** Shuffles the leaves, i.e., tests if exchanging
		the view cells the leaves belong to helps in improving the view cells.
	*/
	bool ShuffleLeaves(VspKdLeaf *leaf1, VspKdLeaf *leaf2) const;

	/** Shuffles, i.e. takes border leaf from view cell 1 and adds it 
		to view cell 2.
	*/
	void ShuffleLeaf(VspKdLeaf *leaf,
					 VspKdViewCell *vc1, 
					 VspKdViewCell *vc2) const;
protected:
	

	/////////////////////////////
	// The core pointer
	VspKdNode *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;
	
	ViewCellsManager *mViewCellsManager;

	priority_queue<VspKdMergeCandidate> mMergeQueue;


	/////////////////////////////
	// 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 acceptable cost ratio to continue subdivision
	float mTermMaxCostRatio;
	/// maximal contribution per ray to subdivide the node
	float mTermMaxRayContribution;
	/// tolerance value indicating how often the max cost ratio can be failed
	int mTermMissTolerance;

	/// maximal numbers of view cells
	int mMaxViewCells;

	/// maximal cost ratio of a merge
	float mMergeMaxCostRatio;
	
	/// minimal number of view cells
	int mMergeMinViewCells;

	/// if only the "driving axis", i.e., the axis with the biggest extent
	/// should be used for subdivision
	bool mOnlyDrivingAxis;

	/////////////////////////////
	VspKdStatistics mStat;	
};

}

#endif // __VSP_KDTREE_H__