// ================================================================
// $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"

// --------------------------------------------------------------
// Static statistics for 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;
  int minRaysNodes;
	// max ray contribution nodes
  int maxRayContribNodes;
	// max depth nodes
  int minSizeNodes;
	
  // max number of rays per node
  int maxRayRefs;
  // number of dynamically added ray refs
  int addedRayRefs;
  // number of dynamically removed ray refs
  int removedRayRefs;
  
  // 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; 
	  maxRayRefs = 0;
	  addedRayRefs = removedRayRefs = 0;
	  initialPvsSize = 0;
	  maxRayContribNodes = 0;
	  minSizeNodes = 0;
  }
  
  void
  Print(ostream &app) const;

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


class VspKdTreeInterior;


// --------------------------------------------------------------
// KD-tree node - abstract class
// --------------------------------------------------------------
class VspKdTreeNode
{
public:

#define USE_FIXEDPOINT_T 1

struct RayInfo
{
	// pointer to the actual ray
	VssRay *mRay;
	
	// endpoints  - do we need them?
#if USE_FIXEDPOINT_T
	short mMinT, mMaxT;
#else
	float mMinT, mMaxT;
#endif
	
	RayInfo():mRay(NULL) {}
	
	RayInfo(VssRay *r):mRay(r), mMinT(0),
#if USE_FIXEDPOINT_T
#define FIXEDPOINT_ONE 0x7FFE
	// mMaxT(0xFFFF)
		mMaxT(FIXEDPOINT_ONE)
#else
		mMaxT(1.0f)
#endif
	{}
	
	RayInfo(VssRay *r,
			const float _min,
			const float _max
			):mRay(r) 
	{
		SetMinT(_min);
		SetMaxT(_max);
	}
	
	RayInfo(VssRay *r,
			const short _min,
			const float _max):mRay(r), mMinT(_min) 
	{
		SetMaxT(_max);
	}
	
	RayInfo(VssRay *r,
			const float _min,
			const short _max):
	mRay(r), mMaxT(_max) 
	{
		SetMinT(_min);
	}
	
	friend bool operator<(const RayInfo &a, const RayInfo &b) {
		return a.mRay < b.mRay;
	}
  
	
	float ExtrapOrigin(const int axis) const {
		return mRay->GetOrigin(axis) + GetMinT()*mRay->GetDir(axis);
	}
	
	float ExtrapTermination(const int axis) const {
		return mRay->GetOrigin(axis) + GetMaxT()*mRay->GetDir(axis);
	}
	
#if USE_FIXEDPOINT_T
	float GetMinT () const { return mMinT/(float)(FIXEDPOINT_ONE); }
	float GetMaxT () const { return mMaxT/(float)(FIXEDPOINT_ONE); }
	
	void SetMinT (const float t) {
		mMinT = (short) (t*(float)(FIXEDPOINT_ONE));
	}
	
	void SetMaxT (const float t) {
		mMaxT = (short) (t*(float)(FIXEDPOINT_ONE));
		mMaxT++;
		//      if (mMaxT!=0xFFFF)
		//	mMaxT++;
	}
#else
	float GetMinT () const { return mMinT; }
	float GetMaxT () const { return mMaxT; }
	
	void SetMinT (const float t) { mMinT = t; }
	void SetMaxT (const float t) { mMaxT = t; }
#endif 


	int ComputeRayIntersection(const int axis, const float position, float &t) const 
	{		
		// intersect the ray with the plane
		float denom = mRay->GetDir(axis);
    
		if (fabs(denom) < 1e-20)
			//if (denom == 0.0f)
			return (mRay->GetOrigin(axis) > position) ? 1 : -1;
    
		t = (position - mRay->GetOrigin(axis))/denom;

		if (t < GetMinT())
			return (denom > 0) ? 1 : -1;

		if (t > GetMaxT())
			return (denom > 0) ? -1 : 1;

		return 0;
	}


};


typedef vector<RayInfo> RayInfoContainer;

enum {EInterior, ELeaf};

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

  enum {SPLIT_X=0, SPLIT_Y, SPLIT_Z, SPLIT_DIRX, SPLIT_DIRY, SPLIT_DIRZ};
  
  // splitting axis
  char axis;
	
  // depth
  unsigned char depth;
  
  //  short depth;
  //
  /////////////////////////////////
  
  inline VspKdTreeNode(VspKdTreeInterior *p);

  
  virtual ~VspKdTreeNode() {};
  virtual int Type() const  = 0;
  

  bool IsLeaf() const { return axis == -1; }
  
  virtual void Print(ostream &s) const = 0;

  virtual int GetAccessTime() {
    return 0x7FFFFFF;
  }

  
	
};

// --------------------------------------------------------------
// KD-tree node - interior node
// --------------------------------------------------------------
class VspKdTreeInterior: public VspKdTreeNode
{
public:
  // plane in local modelling coordinates
  float position;

  // pointers to children
  VspKdTreeNode *back, *front;

  // the bbox of the node
  AxisAlignedBox3 bbox;

  // the bbox of the node
  AxisAlignedBox3 dirBBox;
  
  // data for caching
  long accesses;
  long lastAccessTime;
  
  VspKdTreeInterior(VspKdTreeInterior *p):VspKdTreeNode(p),
					back(NULL),
					front(NULL),
					accesses(0),
					lastAccessTime(-1)
  { }

  virtual int GetAccessTime() {
    return lastAccessTime;
  }

  void SetupChildLinks(VspKdTreeNode *b, VspKdTreeNode *f) {
    back = b;
    front = f;
    b->parent = f->parent = this;
  }

  void ReplaceChildLink(VspKdTreeNode *oldChild, VspKdTreeNode *newChild) {
    if (back == oldChild)
      back = newChild;
    else
      front = newChild;
  }

  virtual int Type() const  { return EInterior; }
  
  virtual ~VspKdTreeInterior()
  {
	  DEL_PTR(back);
	  DEL_PTR(front);
  }
  
  virtual void Print(ostream &s) const 
  {
	  if (axis == 0)
		  s << "x ";
	  else
		  if (axis == 1)
			  s << "y ";
		  else
			  s << "z ";
	  s << position << " ";
	  back->Print(s);
	  front->Print(s);
  }
	
  int ComputeRayIntersection(const RayInfo &rayData, float &t) 
  {
	  return rayData.ComputeRayIntersection(axis, position, t);
  }

};


// --------------------------------------------------------------
// KD-tree node - leaf node
// --------------------------------------------------------------
class VspKdTreeLeaf : public VspKdTreeNode
{
private:
	int mPvsSize;
public:
	static int mailID;
	int mailbox;
  
	RayInfoContainer rays;
	bool mValidPvs;
	
	VspKdTreeLeaf(VspKdTreeInterior *p,	const int nRays):
	VspKdTreeNode(p), rays(), mPvsSize(0), mValidPvs(false) 
	{
		rays.reserve(nRays);
	}
	
	virtual ~VspKdTreeLeaf() { }

	virtual int Type() const  
	{ 
		return ELeaf; 
	}

	virtual void Print(ostream &s) const 
	{
		s << endl << "L: r = " << rays.size() << endl;
	};
  
	void AddRay(const RayInfo &data) 
	{
		mValidPvs = false;
		rays.push_back(data);
		data.mRay->Ref();
	}
	
	int GetPvsSize() const 
	{
		return mPvsSize;
	}
	void SetPvsSize(const int s) 
	{
		mPvsSize = s;
	}

	void UpdatePvsSize();

	void Mail() 
	{ 
		mailbox = mailID; 
	}
  
	static void NewMail() 
	{ 
		++ mailID; 
	}
  
	bool Mailed() const 
	{ 
		return mailbox == mailID; 
	}
  
	bool Mailed(const int mail) 
	{
		return mailbox >= mailID + mail;
	}

	float GetAvgRayContribution() const 
	{
		return GetPvsSize()/((float)rays.size() + Limits::Small);
	}

	float GetSqrRayContribution() const 
	{
		return sqr(GetPvsSize()/((float)rays.size() + Limits::Small));
	}

};

// Inline functions
inline VspKdTreeNode::VspKdTreeNode(VspKdTreeInterior *p):
parent(p), axis(-1), depth(p ? p->depth + 1 : 0) 
{}



// ---------------------------------------------------------------
// 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 *node;
		AxisAlignedBox3 bbox;
		int depth;
		float priority;
        
		TraversalData() {}

		TraversalData(VspKdTreeNode *n, const float p):
		node(n), priority(p)
		{}

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

		//    ~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.node;
			VspKdTreeLeaf *leafb = (VspKdTreeLeaf *) b.node;
#if 0
			return
				leafa->rays.size()*a.bbox.GetVolume()
				<
				leafb->rays.size()*b.bbox.GetVolume();
#endif
#if 1
			return
				leafa->GetPvsSize()*a.bbox.GetVolume()
				<
				leafb->GetPvsSize()*b.bbox.GetVolume();
#endif
#if 0
			return
				leafa->GetPvsSize()
				<
				leafb->GetPvsSize();
#endif
#if 0
			return
				leafa->GetPvsSize() / (float)(leafa->rays.size() + 1)
				>
				leafb->GetPvsSize() / (float)(leafb->rays.size() + 1);
#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 
  {
	  VspKdTreeNode::RayInfo rayData;
	  VspKdTreeNode *node;
      
	  RayTraversalData() {}
	  
	  RayTraversalData(VspKdTreeNode *n, const VspKdTreeNode::RayInfo &data):
	  rayData(data), node(n) {}
  };
	
public:
	
	/////////////////////////////
	// The core pointer
	VspKdTreeNode *root;
  
	/////////////////////////////
	// Basic properties

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

	// axis aligned bounding box of directions
	AxisAlignedBox3 dirBBox;

	const VspKdStatistics &GetStatistics() const {
		return mStat;
	}
	/////////////////////////////
	// Construction parameters

	// epsilon used for the construction
	float epsilon;

	// ratio between traversal and intersection costs
	float ct_div_ci;
	
	// max depth of the tree
	int termMaxDepth;
	// minimal ratio of the volume of the cell and the query volume
	float termMinSize;

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

	// minimal ray number per node to still get subdivided
	int termMinRays;
	
	// maximal cost ration to subdivide a node
	float termMaxCostRatio;
	
	// maximal contribution per ray to subdivide the node
	float termMaxRayContribution;

	// randomized construction
	bool randomize;

	// type of the splitting to use for the tree construction
	enum {ESplitRegular, ESplitHeuristic };
	int splitType;
	
	// maximal size of the box on which the refdir splitting can be performed
	// (relative to the scene bbox
	float refDirBoxMaxSize;
  
	// maximum alovable memory in MB
	float maxTotalMemory;

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

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

	// minimal acess time for collapse
	int accessTimeThreshold;

	// minimal depth at which to perform collapse
	int minCollapseDepth;
	
	// reusable array of split candidates
	vector<SortableEntry> *splitCandidates;
	
	/////////////////////////////
	VspKdStatistics mstat;
	
	VspKdTree();
	virtual ~VspKdTree();

	virtual void Construct(VssRayContainer &rays,
						   AxisAlignedBox3 *forcedBoundingBox = NULL);
	
	// incemental construction
	virtual void UpdateRays(VssRayContainer &remove, VssRayContainer &add);

	virtual void AddRays(VssRayContainer &add)
	{
		VssRayContainer remove;
		UpdateRays(remove, 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);
	
  
	// return memory usage in MB
	float GetMemUsage() const 
	{
		return 
			(sizeof(VspKdTree) + 
			 stat.Leaves() * sizeof(VspKdTreeLeaf) +
			 stat.Interior() * sizeof(VspKdTreeInterior) +
			 stat.rayRefs * sizeof(VspKdTreeNode::RayInfo)) / (1024.0f*1024.0f);
	}
	
	float GetRayMemUsage() const 
	{
		return stat.rays * (sizeof(VssRay))/(1024.0f * 1024.0f);
	}
  
	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);

	AxisAlignedBox3 GetBBox(const VspKdTreeNode *node) 
	{
		if (node->parent == NULL)
			return bbox;

		if (!node->IsLeaf())
			return ((VspKdTreeInterior *)node)->bbox;

		if (node->parent->axis >= 3)
			return node->parent->bbox;
      
		AxisAlignedBox3 box(node->parent->bbox);
		if (node->parent->front == node)
			box.SetMin(node->parent->axis, node->parent->position);
		else
			box.SetMax(node->parent->axis, node->parent->position);
		return box;
	}

	AxisAlignedBox3 GetDirBBox(const VspKdTreeNode *node) 
	{
		if (node->parent == NULL)
			return dirBBox;
		if (!node->IsLeaf())
			return ((VspKdTreeInterior *)node)->dirBBox;
		if (node->parent->axis < 3)
			return node->parent->dirBBox;
    
		AxisAlignedBox3 dBBox(node->parent->dirBBox);

		if (node->parent->front == node)
			dBBox.SetMin(node->parent->axis - 3, node->parent->position);
		else
			dBBox.SetMax(node->parent->axis - 3, node->parent->position);
		return dBBox;
	}
  
	int	ReleaseMemory(const int time);

	int	CollapseSubtree(VspKdTreeNode *node, const int time);
	
	void CountAccess(VspKdTreeInterior *node, const long time) 
	{
		++ node->accesses;
		node->lastAccessTime = time;
	}

	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 
	{
		return GetPvsSize(root, bbox);
	}
	
	int	GetPvsSize(VspKdTreeNode *node, const AxisAlignedBox3 &box) const;

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

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

	float GetAvgPvsSize();
};


#endif // __LSDS_KDTREE_H__