#include <stack>
#include <algorithm>
#include <queue>
#include "Environment.h"
#include "Mesh.h"
#include "KdTree.h"


int KdNode::mailID = 1;


KdNode::KdNode(KdInterior *parent):mParent(parent), mailbox(0)
{
  if (parent)
    mDepth = parent->mDepth+1;
  else
    mDepth = 0;
}

KdTree::KdTree()
{
  mRoot = new KdLeaf(NULL, 0);
  environment->GetIntValue("KdTree.Termination.maxDepth", mTermMaxDepth);
  environment->GetIntValue("KdTree.Termination.minCost", mTermMinCost);
  environment->GetFloatValue("KdTree.Termination.maxCostRatio", mMaxCostRatio);
  environment->GetFloatValue("KdTree.Termination.ct_div_ci", mCt_div_ci);
  environment->GetFloatValue("KdTree.splitBorder", mSplitBorder);

  environment->GetBoolValue("KdTree.sahUseFaces", mSahUseFaces);

  char splitType[64];
  environment->GetStringValue("KdTree.splitMethod", splitType);
  
  mSplitMethod = SPLIT_SPATIAL_MEDIAN;
  if (strcmp(splitType, "spatialMedian") == 0)
    mSplitMethod = SPLIT_SPATIAL_MEDIAN;
  else
    if (strcmp(splitType, "objectMedian") == 0)
      mSplitMethod = SPLIT_OBJECT_MEDIAN;
    else
      if (strcmp(splitType, "SAH") == 0)
	mSplitMethod = SPLIT_SAH;
      else {
	cerr<<"Wrong kd split type "<<splitType<<endl;
	exit(1);
      }
  splitCandidates = NULL;
}

bool
KdTree::Construct()
{
  if (!splitCandidates)
    splitCandidates = new vector<SortableEntry>;

  // first construct a leaf that will get subdivide
  KdLeaf *leaf = (KdLeaf *) mRoot;

  mStat.nodes = 1;

  mBox.Initialize();
  
  ObjectContainer::const_iterator mi;
  for ( mi = leaf->mObjects.begin();
	mi != leaf->mObjects.end();
	mi++) {
    mBox.Include((*mi)->GetBox());
  }

  cout <<"KdTree Root Box:"<< mBox<<endl;
  mRoot = Subdivide(TraversalData(leaf, mBox, 0));

  // remove the allocated array
  delete splitCandidates;

  return true;
}

KdNode *
KdTree::Subdivide(const TraversalData &tdata)
{

  KdNode *result = NULL;

  priority_queue<TraversalData> tStack;
  //  stack<STraversalData> tStack;
  
  tStack.push(tdata);
  AxisAlignedBox3 backBox, frontBox;

  
  while (!tStack.empty()) {

#if 0
    if ( GetMemUsage() > maxMemory ) {
      // count statistics on unprocessed leafs
      while (!tStack.empty()) {
	EvaluateLeafStats(tStack.top());
	tStack.pop();
      }
      break;
    }
#endif

    TraversalData data = tStack.top();
    tStack.pop();
    
    KdNode *node = SubdivideNode((KdLeaf *) data.mNode,
				 data.mBox,
				 backBox,
				 frontBox
				 );
    if (result == NULL)
      result = node;
    
    if (!node->IsLeaf()) {

      KdInterior *interior = (KdInterior *) node;
      // push the children on the stack
      tStack.push(TraversalData(interior->mBack, backBox, data.mDepth+1));
      tStack.push(TraversalData(interior->mFront, frontBox, data.mDepth+1));
      
    } else {
      EvaluateLeafStats(data);
    }
  }
  
  return result;

}



bool
KdTree::TerminationCriteriaMet(const KdLeaf *leaf)
{
  //  cerr<<"\n OBJECTS="<<leaf->mObjects.size()<<endl;
  return
    (leaf->mObjects.size() <= mTermMinCost) ||
    (leaf->mDepth >= mTermMaxDepth);
  
}


int
KdTree::SelectPlane(KdLeaf *leaf,
		    const AxisAlignedBox3 &box,
		    float &position
		    )
{
  int axis = -1;
  
  switch (mSplitMethod)
    {
    case SPLIT_SPATIAL_MEDIAN: {
      axis = box.Size().DrivingAxis();
      position = (box.Min()[axis] + box.Max()[axis])*0.5f;
      break;
    }
    case SPLIT_SAH: {
      int objectsBack, objectsFront;
      float costRatio;
      bool mOnlyDrivingAxis = false;
      if (mOnlyDrivingAxis) {
				axis = box.Size().DrivingAxis();
				costRatio = BestCostRatio(leaf,
																	box,
																	axis,
																	position,
																	objectsBack,
																	objectsFront);
      } else {
				costRatio = MAX_FLOAT;
				for (int i=0; i < 3; i++) {
					float p;
					float r = BestCostRatio(leaf,
																	box,
																	i,
																	p,
																	objectsBack,
																	objectsFront);
					if (r < costRatio) {
						costRatio = r;
						axis = i;
						position = p;
					}
				}
      }
      
      if (costRatio > mMaxCostRatio) {
				//	cout<<"Too big cost ratio "<<costRatio<<endl;
				axis = -1;
      }
      break;
    }
    
    }
  return axis;
}

KdNode *
KdTree::SubdivideNode(
		      KdLeaf *leaf,
		      const AxisAlignedBox3 &box,
		      AxisAlignedBox3 &backBBox,
		      AxisAlignedBox3 &frontBBox
		      )
{
  
  if (TerminationCriteriaMet(leaf))
    return leaf;
  
  float position;
  
  // select subdivision axis
  int axis = SelectPlane( leaf, box, position );

  if (axis == -1) {
    return leaf;
  }
  
  mStat.nodes+=2;
  mStat.splits[axis]++;
  
  // add the new nodes to the tree
  KdInterior *node = new KdInterior(leaf->mParent);

  node->mAxis = axis;
  node->mPosition = position;
  node->mBox = box;
  
  backBBox = box;
  frontBBox = box;
  
  // first count ray sides
  int objectsBack = 0;
  int objectsFront = 0;
  
  backBBox.SetMax(axis, position);
  frontBBox.SetMin(axis, position);

  ObjectContainer::const_iterator mi;

  for ( mi = leaf->mObjects.begin();
	mi != leaf->mObjects.end();
	mi++) {
    // determine the side of this ray with respect to the plane
    AxisAlignedBox3 box = (*mi)->GetBox();
    if (box.Max(axis) > position ) 
      objectsFront++;
    
    if (box.Min(axis) < position )
      objectsBack++;
  }

  
  KdLeaf *back = new KdLeaf(node, objectsBack);
  KdLeaf *front = new KdLeaf(node, objectsFront);

  // replace a link from node's parent
  if (  leaf->mParent )
    leaf->mParent->ReplaceChildLink(leaf, node);

  // and setup child links
  node->SetupChildLinks(back, front);
  
  for (mi = leaf->mObjects.begin();
       mi != leaf->mObjects.end();
       mi++) {
    // determine the side of this ray with respect to the plane
    AxisAlignedBox3 box = (*mi)->GetBox();

    if (box.Max(axis) >= position )
      front->mObjects.push_back(*mi);
    
    if (box.Min(axis) < position )
      back->mObjects.push_back(*mi);
    
    mStat.objectRefs -= leaf->mObjects.size();
    mStat.objectRefs += objectsBack + objectsFront;
  }
  
  delete leaf;
  return node;
}



void
KdTreeStatistics::Print(ostream &app) const
{
  app << "===== KdTree statistics ===============\n";

  app << "#N_RAYS Number of rays )\n"
      << rays <<endl;
  app << "#N_DOMAINS  ( Number of query domains )\n"
      << queryDomains <<endl;
  
  app << "#N_NODES ( Number of nodes )\n" << nodes << "\n";

  app << "#N_LEAVES ( Number of leaves )\n" << Leaves() << "\n";

  app << "#N_SPLITS ( Number of splits in axes x y z dx dy dz \n";
  for (int i=0; i<7; i++)
    app << splits[i] <<" ";
  app <<endl;

  app << "#N_RAYREFS ( Number of rayRefs )\n" <<
    rayRefs << "\n";

  app << "#N_RAYRAYREFS  ( Number of rayRefs / ray )\n" <<
    rayRefs/(double)rays << "\n";

  app << "#N_LEAFRAYREFS  ( Number of rayRefs / leaf )\n" <<
    rayRefs/(double)Leaves() << "\n";

  app << "#N_MAXOBJECTREFS  ( Max number of rayRefs / leaf )\n" <<
    maxObjectRefs << "\n";

  app << "#N_NONEMPTYRAYREFS  ( Number of rayRefs in nonEmpty leaves / non empty leaf )\n" <<
    rayRefsNonZeroQuery/(double)(Leaves() - zeroQueryNodes) << "\n";

  app << "#N_LEAFDOMAINREFS  ( Number of query domain Refs / leaf )\n" <<
    objectRefs/(double)Leaves() << "\n";

  //  app << setprecision(4);

  app << "#N_PEMPTYLEAVES  ( Percentage of leaves with zero query domains )\n"<<
    zeroQueryNodes*100/(double)Leaves()<<endl;

  app << "#N_PMAXDEPTHLEAVES ( Percentage of leaves at maxdepth )\n"<<
    maxDepthNodes*100/(double)Leaves()<<endl;

  app << "#N_PMINCOSTLEAVES  ( Percentage of leaves with minCost )\n"<<
    minCostNodes*100/(double)Leaves()<<endl;

  app << "#N_ADDED_RAYREFS  (Number of dynamically added ray references )\n"<<
    addedRayRefs<<endl;

  app << "#N_REMOVED_RAYREFS  (Number of dynamically removed ray references )\n"<<
    removedRayRefs<<endl;

  //  app << setprecision(4);

  //  app << "#N_CTIME  ( Construction time [s] )\n"
  //      << Time() << " \n";

  app << "===== END OF KdTree statistics ==========\n";

}



void
KdTree::EvaluateLeafStats(const TraversalData &data)
{

  // the node became a leaf -> evaluate stats for leafs
  KdLeaf *leaf = (KdLeaf *)data.mNode;

  if (data.mDepth > mTermMaxDepth)
    mStat.maxDepthNodes++;
  
  if ( (int)(leaf->mObjects.size()) < mTermMinCost)
    mStat.minCostNodes++;
  
  
  if ( (int)(leaf->mObjects.size()) > mStat.maxObjectRefs)
    mStat.maxObjectRefs = leaf->mObjects.size();
  
}



void
KdTree::SortSplitCandidates(
			    KdLeaf *node,
			    const int axis
			    )
{
  splitCandidates->clear();
  
  int requestedSize = 2*node->mObjects.size();
  // creates a sorted split candidates array
  if (splitCandidates->capacity() > 500000 &&
      requestedSize < (int)(splitCandidates->capacity()/10) ) {
    delete splitCandidates;
    splitCandidates = new vector<SortableEntry>;
  }
  
  splitCandidates->reserve(requestedSize);
  
  // insert all queries 
  for(ObjectContainer::const_iterator mi = node->mObjects.begin();
      mi != node->mObjects.end();
      mi++) {
    AxisAlignedBox3 box = (*mi)->GetBox();

    splitCandidates->push_back(SortableEntry(SortableEntry::BOX_MIN,
					     box.Min(axis),
					     *mi)
			       );
    
    
    splitCandidates->push_back(SortableEntry(SortableEntry::BOX_MAX,
					     box.Max(axis),
					     *mi)
			       );
  }
  
  stable_sort(splitCandidates->begin(), splitCandidates->end());
}


float
KdTree::BestCostRatio(
											KdLeaf *node,
											const AxisAlignedBox3 &box,
											const int axis,
											float &position,
											int &objectsBack,
											int &objectsFront
											)
{

  SortSplitCandidates(node, axis);
  
  // go through the lists, count the number of objects left and right
  // and evaluate the following cost funcion:
  // C = ct_div_ci  + (ol + or)/queries

  float totalIntersections = 0.0f;
  vector<SortableEntry>::const_iterator ci;

  for(ci = splitCandidates->begin();
      ci < splitCandidates->end();
      ci++) 
    if ((*ci).type == SortableEntry::BOX_MIN) {
      totalIntersections += (*ci).intersectable->IntersectionComplexity();
    }
	
  float intersectionsLeft = 0;
  float intersectionsRight = totalIntersections;
	
  int objectsLeft = 0, objectsRight = node->mObjects.size();
  
  float minBox = box.Min(axis);
  float maxBox = box.Max(axis);
  float boxArea = box.SurfaceArea();
  
  float minBand = minBox + mSplitBorder*(maxBox - minBox);
  float maxBand = minBox + (1.0f - mSplitBorder)*(maxBox - minBox);
  
  float minSum = 1e20;
  
  for(ci = splitCandidates->begin();
      ci < splitCandidates->end();
      ci++) {
    switch ((*ci).type) {
    case SortableEntry::BOX_MIN:
      objectsLeft++;
      intersectionsLeft += (*ci).intersectable->IntersectionComplexity();
      break;
    case SortableEntry::BOX_MAX:
      objectsRight--;
      intersectionsRight -= (*ci).intersectable->IntersectionComplexity();
      break;
    }

    if ((*ci).value > minBand && (*ci).value < maxBand) {
      AxisAlignedBox3 lbox = box;
      AxisAlignedBox3 rbox = box;
      lbox.SetMax(axis, (*ci).value);
      rbox.SetMin(axis, (*ci).value);

      float sum;
      if (mSahUseFaces)
	sum = intersectionsLeft*lbox.SurfaceArea() + intersectionsRight*rbox.SurfaceArea();
      else
	sum = objectsLeft*lbox.SurfaceArea() + objectsRight*rbox.SurfaceArea();
      
      //      cout<<"pos="<<(*ci).value<<"\t q=("<<ql<<","<<qr<<")\t r=("<<rl<<","<<rr<<")"<<endl;
      //      cout<<"cost= "<<sum<<endl;
      
      if (sum < minSum) {
	minSum = sum;
	position = (*ci).value;
	
	objectsBack = objectsLeft;
	objectsFront = objectsRight;
      }
    }
  }
  
  float oldCost = mSahUseFaces ? totalIntersections : node->mObjects.size();
  float newCost = mCt_div_ci + minSum/boxArea;
  float ratio = newCost/oldCost;
  
#if 0
  cout<<"===================="<<endl;
  cout<<"costRatio="<<ratio<<" pos="<<position<<" t="<<(position - minBox)/(maxBox - minBox)
      <<"\t o=("<<objectsBack<<","<<objectsFront<<")"<<endl;
#endif
  return ratio;
}

int
KdTree::CastRay(
		Ray &ray
		)
{
  int hits = 0;
  
  stack<RayTraversalData> tStack;
  
  float maxt = 1e6;
  float mint = 0;

  Intersectable::NewMail();

  if (!mBox.GetMinMaxT(ray, &mint, &maxt))
    return 0;

  if (mint < 0)
    mint = 0;
  
  maxt += Limits::Threshold;
  
  Vector3 entp = ray.Extrap(mint);
  Vector3 extp = ray.Extrap(maxt);
  
  KdNode *node = mRoot;
  KdNode *farChild;
  float position;
  int axis;
  
  while (1) {
    if (!node->IsLeaf()) {
      KdInterior *in = (KdInterior *) node;
      position = in->mPosition;
      axis = in->mAxis;

      if (entp[axis] <= position) {
	if (extp[axis] <= position) {
	  node = in->mBack;
	  // cases N1,N2,N3,P5,Z2,Z3
	  continue;
	} else {
	  // case N4
	  node = in->mBack;
	  farChild = in->mFront;
	}
      }
      else {
	if (position <= extp[axis]) {
	  node = in->mFront;
	  // cases P1,P2,P3,N5,Z1
	  continue;
	} else {
	  node = in->mFront;
	  farChild = in->mBack;
	  // case P4
	}
	  }
      // $$ modification 3.5.2004 - hints from Kamil Ghais
      // case N4 or P4
      float tdist = (position - ray.GetLoc(axis)) / ray.GetDir(axis);
      tStack.push(RayTraversalData(farChild, extp, maxt));
      extp = ray.GetLoc() + ray.GetDir()*tdist;
      maxt = tdist;
	} else {
      // compute intersection with all objects in this leaf
      KdLeaf *leaf = (KdLeaf *) node;
      ray.leaves.push_back(leaf);
      
      ObjectContainer::const_iterator mi;
      for ( mi = leaf->mObjects.begin();
	    mi != leaf->mObjects.end();
	    mi++) {
	Intersectable *object = *mi;
	if (!object->Mailed() ) {
	  object->Mail();
	  //ray.meshes.push_back(mesh);
	  hits += object->CastRay(ray);
	}
      }

      if (hits && ray.GetType() == Ray::LOCAL_RAY)
	if (ray.intersections[0].mT <= maxt)
	  break;
      
      // get the next node from the stack
      if (tStack.empty())
	break;
      
      entp = extp;
      mint = maxt;
      if (ray.GetType() == Ray::LINE_SEGMENT && mint > 1.0f)
	break;
      
      RayTraversalData &s  = tStack.top();
      node = s.mNode;
      extp = s.mExitPoint;
      maxt = s.mMaxT;
      tStack.pop();
    }
  }
  
  
  return hits;
}

void
KdTree::CollectObjects(KdNode *n, ObjectContainer &objects)
{
  stack<KdNode *> nodeStack;

  nodeStack.push(n);

  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      KdLeaf *leaf = (KdLeaf *)node;
      for (int j=0; j < leaf->mObjects.size(); j++) {
	Intersectable *object = leaf->mObjects[j];
	if (!object->Mailed()) {
	  object->Mail();
	  objects.push_back(object);
	}
      }
    } else {
      KdInterior *interior = (KdInterior *)node;
      nodeStack.push(interior->mFront);
      nodeStack.push(interior->mBack);
    }
  }
}

// Find random neighbor which was not mailed
KdNode *
KdTree::FindRandomNeighbor(KdNode *n,
													 bool onlyUnmailed
													 )
{
  stack<KdNode *> nodeStack;
  
  nodeStack.push(mRoot);

  AxisAlignedBox3 box = GetBox(n);
  int mask = rand();

  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      if ( node != n && (!onlyUnmailed || !node->Mailed()) )
	return node;
    } else {
      KdInterior *interior = (KdInterior *)node;
      if (interior->mPosition > box.Max(interior->mAxis))
	nodeStack.push(interior->mBack);
      else
	if (interior->mPosition < box.Min(interior->mAxis))
	  nodeStack.push(interior->mFront);
	else {
	  // random decision
	  if (mask&1)
	    nodeStack.push(interior->mBack);
	  else
	    nodeStack.push(interior->mFront);
	  mask = mask>>1;
	}
    }
  }
  
  return NULL;
}

int
KdTree::FindNeighbors(KdNode *n,
		      vector<KdNode *> &neighbors,
		      bool onlyUnmailed
		      )
{
  stack<KdNode *> nodeStack;
  
  nodeStack.push(mRoot);

  AxisAlignedBox3 box = GetBox(n);

  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      if ( node != n && (!onlyUnmailed || !node->Mailed()) ) 
	neighbors.push_back(node);
    } else {
      KdInterior *interior = (KdInterior *)node;
      if (interior->mPosition > box.Max(interior->mAxis))
	nodeStack.push(interior->mBack);
      else
	if (interior->mPosition < box.Min(interior->mAxis))
	  nodeStack.push(interior->mFront);
	else {
	  // random decision
	    nodeStack.push(interior->mBack);
	    nodeStack.push(interior->mFront);
	}
    }
  }
  
  return neighbors.size();
}

// Find random neighbor which was not mailed
KdNode *
KdTree::GetRandomLeaf(const Plane3 &plane)
{
  stack<KdNode *> nodeStack;
  
  nodeStack.push(mRoot);
  
  int mask = rand();
  
  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      return node;
    } else {
      KdInterior *interior = (KdInterior *)node;
      KdNode *next;
	if (GetBox(interior->mBack).Side(plane) < 0)
	  next = interior->mFront;
	else
	  if (GetBox(interior->mFront).Side(plane) < 0)
	    next = interior->mBack;
	  else {
	    // random decision
	    if (mask&1)
	      next = interior->mBack;
	    else
	      next = interior->mFront;
	    mask = mask>>1;
	  }
	nodeStack.push(next);
    }
  }
  
  
  return NULL;
}

void
KdTree::CollectLeaves(vector<KdLeaf *> &leaves)
{
  stack<KdNode *> nodeStack;
  nodeStack.push(mRoot);

  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      KdLeaf *leaf = (KdLeaf *)node;
      leaves.push_back(leaf);
    } else {
      KdInterior *interior = (KdInterior *)node;
      nodeStack.push(interior->mBack);
      nodeStack.push(interior->mFront);
    }
  }
}


int
KdTree::CollectLeafPvs()
{
  int totalPvsSize = 0;
  stack<KdNode *> nodeStack;
  
  nodeStack.push(mRoot);
  
  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      KdLeaf *leaf = (KdLeaf *)node;
      for (int j=0; j < leaf->mObjects.size(); j++) {
	Intersectable *object = leaf->mObjects[j];
	if (!object->Mailed()) {
	  object->Mail();
	  // add this node to pvs of all nodes it can see
	  KdPvsMap::iterator ni = object->mKdPvs.mEntries.begin();
	  for (; ni != object->mKdPvs.mEntries.end(); ni++) {
	    KdNode *node = (*ni).first;
	    // $$ JB TEMPORARY solution -> should add object PVS or explictly computed
	    // kd tree PVS
	    if (leaf->mKdPvs.AddSample(node))
	      totalPvsSize++;
	  }
	}
      }
    } else {
      KdInterior *interior = (KdInterior *)node;
      nodeStack.push(interior->mFront);
      nodeStack.push(interior->mBack);
    }
  }

  return totalPvsSize;
}


KdNode *
KdTree::GetRandomLeaf(const bool onlyUnmailed)
{
	stack<KdNode *> nodeStack;
  nodeStack.push(mRoot);

	int mask = rand();
	
  while (!nodeStack.empty()) {
    KdNode *node = nodeStack.top();
    nodeStack.pop();
    if (node->IsLeaf()) {
      if ( (!onlyUnmailed || !node->Mailed()) )
				return node;
    } else {
      KdInterior *interior = (KdInterior *)node;
			// random decision
			if (mask&1)
				nodeStack.push(interior->mBack);
			else
				nodeStack.push(interior->mFront);
			mask = mask>>1;
		}
	}
  return NULL;
}