/** @file 
 *  @brief Representation of the optimized k-d tree.
 *  @author Remus Dumitru. Jacobs University Bremen, Germany
 *  @author Corneliu-Claudiu Prodescu. Jacobs University Bremen, Germany
 *  @author Andreas Nuechter. Jacobs University Bremen, Germany
 *  @author Kai Lingemann. Inst. of CS, University of Osnabrueck, Germany
 *  @author Thomas Escher. Inst. of CS, University of Osnabrueck, Germany
 */

#ifndef __KD_TREE_IMPL_H__
#define __KD_TREE_IMPL_H__

#include "slam6d/kdparams.h"
#include "globals.icc"

#ifdef _MSC_VER
#if !defined _OPENMP && defined OPENMP 
#define _OPENMP
#endif
#endif

#ifdef _OPENMP
#include <omp.h>
#endif

class PointCompare {
public:
    bool operator() (const std::pair<Point, double>& left,
				 const std::pair<Point, double>& right)
  {
    return left.second > right.second;
  }
};

/**
 * @brief The optimized k-d tree. 
 * 
 * A kD tree for points, with limited
 * capabilities (find nearest point to
 * a given point, or to a ray).
 **/
template<class PointData, class AccessorData, class AccessorFunc>
class KDTreeImpl {
public:
  inline KDTreeImpl() { }
  
  virtual inline ~KDTreeImpl() {
    if (!npts) {
#ifdef WITH_OPENMP_KD
      omp_set_num_threads(OPENMP_NUM_THREADS);
#pragma omp parallel for schedule(dynamic)
#endif
      for (int i = 0; i < 2; i++) {
        if (i == 0 && node.child1) delete node.child1;
        if (i == 1 && node.child2) delete node.child2;
      }
    } else {
      if (leaf.p) delete [] leaf.p;
    }
  }

  virtual void create(PointData pts, AccessorData *indices, size_t n) {
    AccessorFunc point;
    
    // Find bbox
    double xmin = point(pts, indices[0])[0], xmax = point(pts, indices[0])[0];
    double ymin = point(pts, indices[0])[1], ymax = point(pts, indices[0])[1];
    double zmin = point(pts, indices[0])[2], zmax = point(pts, indices[0])[2];
    for(unsigned int i = 1; i < n; i++) {
      xmin = min(xmin, point(pts, indices[i])[0]);
      xmax = max(xmax, point(pts, indices[i])[0]);
      ymin = min(ymin, point(pts, indices[i])[1]);
      ymax = max(ymax, point(pts, indices[i])[1]);
      zmin = min(zmin, point(pts, indices[i])[2]);
      zmax = max(zmax, point(pts, indices[i])[2]);
    }

    // Leaf nodes
    if ((n > 0) && (n <= 10)) {
      npts = n;
      leaf.p = new AccessorData[n];
      // fill leaf index array with indices
      for(unsigned int i = 0; i < n; ++i) {
        leaf.p[i] = indices[i];
      }
      return;
    }

    // Else, interior nodes
    npts = 0;

    node.center[0] = 0.5 * (xmin+xmax);
    node.center[1] = 0.5 * (ymin+ymax);
    node.center[2] = 0.5 * (zmin+zmax);
    node.dx = 0.5 * (xmax-xmin);
    node.dy = 0.5 * (ymax-ymin);
    node.dz = 0.5 * (zmax-zmin);
    node.r2 = sqr(node.dx) + sqr(node.dy) + sqr(node.dz);

    // Find longest axis
    if (node.dx > node.dy) {
      if (node.dx > node.dz) {
        node.splitaxis = 0;
      } else {
        node.splitaxis = 2;
      }
    } else {
      if (node.dy > node.dz) {
        node.splitaxis = 1;
      } else {
        node.splitaxis = 2;
      }
    }

    // Partition
    double splitval = node.center[node.splitaxis];

    if ( fabs(max(max(node.dx,node.dy),node.dz)) < 0.01 ) {
      npts = n;
      leaf.p = new AccessorData[n];
      // fill leaf index array with indices
      for(unsigned int i = 0; i < n; ++i) {
        leaf.p[i] = indices[i];
      }
      return;
    }

    AccessorData* left = indices, * right = indices + n - 1;
    while(true) {
      while(point(pts, *left)[node.splitaxis] < splitval)
        left++;
      while(point(pts, *right)[node.splitaxis] >= splitval)
        right--;
      if(right < left)
        break;
      std::swap(*left, *right);
    }

    // Build subtrees
    int i;
#ifdef WITH_OPENMP_KD                   // does anybody know the reason why this is slower ?? --Andreas
    omp_set_num_threads(OPENMP_NUM_THREADS);
#pragma omp parallel for schedule(dynamic) 
#endif
    for (i = 0; i < 2; i++) {
      if (i == 0) {
        node.child1 = new KDTreeImpl();
        node.child1->create(pts, indices, left - indices);
      }
      if (i == 1) {
        node.child2 = new KDTreeImpl();
        node.child2->create(pts, left, n - (left - indices));
      }
    }
  }

protected:
  /**
   * storing the parameters of the k-d tree, i.e., the current closest point,
   * the distance to the current closest point and the point itself.
   * These global variable are needed in this search.
   *
   * Padded in the parallel case.
   */
#ifdef _OPENMP
#ifdef __INTEL_COMPILER
  __declspec (align(16)) static KDParams params[MAX_OPENMP_NUM_THREADS];
#else
  static KDParams params[MAX_OPENMP_NUM_THREADS];
#endif //__INTEL_COMPILER
#else
  static KDParams params[MAX_OPENMP_NUM_THREADS];
#endif	

  /**
   * number of points. If this is 0: intermediate node. If nonzero: leaf.
   */
  int npts;
  
  /**
   * Cue the standard rant about anon unions but not structs in C++
   */
  union {
    /** 
     * in case of internal node... 
     */
    struct {	 
      double center[3]; ///< storing the center of the voxel (R^3)
      double dx,  ///< defining the voxel itself
	        dy,  ///< defining the voxel itself
	        dz,  ///< defining the voxel itself
             r2;  ///< defining the voxel itself
      int splitaxis;   ///< defining the kind of splitaxis
      KDTreeImpl *child1;  ///< pointers to the childs
      KDTreeImpl *child2;  ///< pointers to the childs
    } node;
    /** 
     * in case of leaf node ... 
     */
    struct {
      /** 
       * store the value itself
       * Here we store just a pointer to the data
       */
      AccessorData* p;
    } leaf;
  };

  void _FindClosest(const PointData& pts, int threadNum) const {
    AccessorFunc point;

    // Leaf nodes
    if (npts) {
      for (int i = 0; i < npts; i++) {
        double myd2 = Dist2(params[threadNum].p, point(pts, leaf.p[i]));
        if (myd2 < params[threadNum].closest_d2) {
          params[threadNum].closest_d2 = myd2;
          params[threadNum].closest = point(pts, leaf.p[i]);
        }
      }
      return;
    }

    // Quick check of whether to abort  
    double approx_dist_bbox =
	 max(max(fabs(params[threadNum].p[0]-node.center[0])-node.dx,
		    fabs(params[threadNum].p[1]-node.center[1])-node.dy),
		fabs(params[threadNum].p[2]-node.center[2])-node.dz);
    if (approx_dist_bbox >= 0 &&
	   sqr(approx_dist_bbox) >= params[threadNum].closest_d2)
      return;

    // Recursive case
    double myd = node.center[node.splitaxis] - params[threadNum].p[node.splitaxis];
    if (myd >= 0.0) {
      node.child1->_FindClosest(pts, threadNum);
      if (sqr(myd) < params[threadNum].closest_d2) {
        node.child2->_FindClosest(pts, threadNum);
      }
    } else {
      node.child2->_FindClosest(pts, threadNum);
      if (sqr(myd) < params[threadNum].closest_d2) {
        node.child1->_FindClosest(pts, threadNum);
      }
    }
  }

  void _FindClosestAlongDir(const PointData& pts, int threadNum) const {
    AccessorFunc point;

    // Leaf nodes
    if (npts) {
      for (int i=0; i < npts; i++) {
        double p2p[] =  { params[threadNum].p[0] - point(pts, leaf.p[i])[0],
                          params[threadNum].p[1] - point(pts, leaf.p[i])[1],
                          params[threadNum].p[2] - point(pts, leaf.p[i])[2] };
        double myd2 = Len2(p2p) - sqr(Dot(p2p, params[threadNum].dir));
        if ((myd2 < params[threadNum].closest_d2)) {
          params[threadNum].closest_d2 = myd2;
          params[threadNum].closest = point(pts, leaf.p[i]);
        }
      }
      return;
    }


    // Quick check of whether to abort
    double p2c[] = { params[threadNum].p[0] - node.center[0],
                     params[threadNum].p[1] - node.center[1],
                     params[threadNum].p[2] - node.center[2] };
    double myd2center = Len2(p2c) - sqr(Dot(p2c, params[threadNum].dir));
    if (myd2center > node.r2 + params[threadNum].closest_d2 + 2.0f * max(node.r2, params[threadNum].closest_d2))
      return;


    // Recursive case
    if (params[threadNum].p[node.splitaxis] < node.center[node.splitaxis] ) {
      node.child1->_FindClosestAlongDir(pts, threadNum);
      node.child2->_FindClosestAlongDir(pts, threadNum);
    } else {
      node.child2->_FindClosestAlongDir(pts, threadNum);
      node.child1->_FindClosestAlongDir(pts, threadNum);
    }
  }

  void _FixedRangeSearch(const PointData& pts, int threadNum) const {
    AccessorFunc point;

    // Leaf nodes
    if (npts) {
	 for (int i = 0; i < npts; i++) {
	   double myd2 = Dist2(params[threadNum].p, point(pts, leaf.p[i]));
	   if (myd2 < params[threadNum].closest_d2) {
		params[threadNum].closest = point(pts, leaf.p[i]);

		Point newPt;
		double* currPt = point(pts, leaf.p[i]);
		newPt.x = currPt[0];
		newPt.y = currPt[1];
		newPt.z = currPt[2];
		params[threadNum].heap.push_back(std::make_pair(newPt, myd2));
		std::push_heap(params[threadNum].heap.begin(),
					params[threadNum].heap.end(),
					PointCompare());
	   }
	 }
	 return;
    }

    // Quick check of whether to abort
    double approx_dist_bbox =
	 max(max(fabs(params[threadNum].p[0]-node.center[0])-node.dx,
		    fabs(params[threadNum].p[1]-node.center[1])-node.dy),
		fabs(params[threadNum].p[2]-node.center[2])-node.dz);
    if (approx_dist_bbox >= 0 &&
	   sqr(approx_dist_bbox) >= params[threadNum].closest_d2)
	 return;

    // Recursive case
    double myd = node.center[node.splitaxis] - params[threadNum].p[node.splitaxis];
    if (myd >= 0.0) {
	 node.child1->_FixedRangeSearch(pts, threadNum);
	 if (sqr(myd) < params[threadNum].closest_d2) {
	   node.child2->_FixedRangeSearch(pts, threadNum);
	 }
    } else {
	 node.child2->_FixedRangeSearch(pts, threadNum);
	 if (sqr(myd) < params[threadNum].closest_d2) {
	   node.child1->_FixedRangeSearch(pts, threadNum);
	 }
    }
  }


  void _KNNSearch(const PointData& pts, int threadNum) const {
    AccessorFunc point;

    // Leaf nodes
    if (npts) {
	 for (int i = 0; i < npts; i++) {
	   double myd2 = Dist2(params[threadNum].p, point(pts, leaf.p[i]));

	   if (myd2 < params[threadNum].closest_d2) {
		Point newPt;
		double* currPt = point(pts, leaf.p[i]);
		newPt.x = currPt[0];
		newPt.y = currPt[1];
		newPt.z = currPt[2];
		params[threadNum].heap.push_back(std::make_pair(newPt, myd2));
		std::push_heap(params[threadNum].heap.begin(),
					params[threadNum].heap.end(),
					PointCompare());

		params[threadNum].closest = point(pts, leaf.p[i]);
	   }
	 }
	 return;
    }

    // Quick check of whether to abort
    double approx_dist_bbox =
	 max(max(fabs(params[threadNum].p[0]-node.center[0])-node.dx,
		    fabs(params[threadNum].p[1]-node.center[1])-node.dy),
		fabs(params[threadNum].p[2]-node.center[2])-node.dz);
    if (approx_dist_bbox >= 0 &&
	   sqr(approx_dist_bbox) >= params[threadNum].closest_d2)
	 return;

    // Recursive case
    double myd = node.center[node.splitaxis] - params[threadNum].p[node.splitaxis];
    if (myd >= 0.0) {
	 node.child1->_KNNSearch(pts, threadNum);
	 if (sqr(myd) < params[threadNum].closest_d2) {
	   node.child2->_KNNSearch(pts, threadNum);
	 }
    } else {
	 node.child2->_KNNSearch(pts, threadNum);
	 if (sqr(myd) < params[threadNum].closest_d2) {
	   node.child1->_KNNSearch(pts, threadNum);
	 }
    }
  }
};


#endif