#include <cstring>
#include <algorithm>
#include <limits>
#include <iostream>

namespace CosmoTool {

  template<int N, typename ValType>
  class CellCompare
  {
  public:
    CellCompare(int k)
    {
      rank = k;
    }

    bool operator()(const KDCell<N,ValType> *a, const KDCell<N,ValType> *b) const
    {
      return (a->coord[rank] < b->coord[rank]);
    }
  protected:
    int rank;
  };

  template<int N, typename ValType>
  KDTree<N,ValType>::~KDTree()
  {
  }

  template<int N, typename ValType>
  KDTree<N,ValType>::KDTree(KDCell<N,ValType> *cells, uint32_t Ncells)
  {
    
    numNodes = Ncells;
    nodes = new KDTreeNode<N,ValType>[numNodes];

    sortingHelper = new KDCell<N,ValType> *[Ncells];
    for (uint32_t i = 0; i < Ncells; i++)
	sortingHelper[i] = &cells[i];    

    optimize();
  }

  template<int N, typename ValType>
  void KDTree<N,ValType>::optimize()
  {
    coords absoluteMin, absoluteMax;

    std::cout << "Optimizing the tree..." << std::endl;
    uint32_t activeCells = gatherActiveCells(sortingHelper, numNodes);
    std::cout << "  number of active cells = " << activeCells << std::endl;

    lastNode = 0;
    for (int i = 0; i < N; i++)
      {
	absoluteMin[i] = -std::numeric_limits<typeof (absoluteMin[0])>::max();
	absoluteMax[i] = std::numeric_limits<typeof (absoluteMax[0])>::max();
      }
    
    std::cout << "   rebuilding the tree..." << std::endl;
    root = buildTree(sortingHelper, activeCells, 0, absoluteMin, absoluteMax);    
    std::cout << "   done." << std::endl;
  }

  template<int N, typename ValType>
  uint32_t KDTree<N,ValType>::getIntersection(const coords& x, CoordType r, 
					       KDCell<N, ValType> **cells,
					       uint32_t numCells)
    throw (NotEnoughCells)
  {
    RecursionInfoCells<N,ValType> info;

    memcpy(info.x, x, sizeof(x));
    info.r = r;
    info.r2 = r*r;
    info.cells = cells;
    info.currentRank = 0;
    info.numCells = numCells;
    info.distances = 0;

    recursiveIntersectionCells(info, root, 0);
    return info.currentRank;
  }

  template<int N, typename ValType>
  uint32_t KDTree<N,ValType>::getIntersection(const coords& x, CoordType r, 
					      KDCell<N, ValType> **cells,
					      CoordType *distances,
					      uint32_t numCells)
    throw (NotEnoughCells)
  {
    RecursionInfoCells<N,ValType> info;

    memcpy(info.x, x, sizeof(x));
    info.r = r;
    info.r2 = r*r;
    info.cells = cells;
    info.currentRank = 0;
    info.numCells = numCells;
    info.distances = distances;

    recursiveIntersectionCells(info, root, 0);
    return info.currentRank;
  }

  template<int N, typename ValType>
  void KDTree<N,ValType>::recursiveIntersectionCells(RecursionInfoCells<N,ValType>& info, 
						      KDTreeNode<N,ValType> *node,
						      int level)
    throw (NotEnoughCells)
  {
    int axis = level % N;
    CoordType d2 = 0;

    if (node->value->active) 
      {
	for (int j = 0; j < 3; j++)
	  {
	    CoordType delta = info.x[j]-node->value->coord[j];
	    d2 += delta*delta;
	  }
	if (d2 < info.r2)
	  {
	    if (info.currentRank == info.numCells)
	      throw NotEnoughCells();
	    info.cells[info.currentRank] = node->value;
	    if (info.distances)
	      info.distances[info.currentRank] = d2;
	    info.currentRank++;
	  }
      }

    // The hypersphere intersects the left child node
    if (((info.x[axis]+info.r) > node->minBound[axis]) &&
	((info.x[axis]-info.r) < node->value->coord[axis]))
      {
	if (node->children[0] != 0)
	  recursiveIntersectionCells(info, node->children[0],
				     level+1);
      }
    if (((info.x[axis]+info.r) > node->value->coord[axis]) &&
	((info.x[axis]-info.r) < node->maxBound[axis]))
      {
	if (node->children[1] != 0)
	  recursiveIntersectionCells(info, node->children[1],
				     level+1);
      }
  }

  template<int N, typename ValType>
  uint32_t gatherActiveCells(KDCell<N,ValType> **cells,
			 uint32_t Ncells)
  {
    uint32_t swapId = Ncells-1;
    uint32_t i = 0;

    while (!cells[swapId]->active && swapId > 0)
      swapId--;

    while (i < swapId)
      {
	if (!cells[i]->active)
	  {
	    std::swap(cells[i], cells[swapId]);
	    while (!cells[swapId]->active && swapId > i)
	      {
		swapId--;
	      }
	  }
	i++;
      }
    return swapId+1;
  }

  template<int N, typename ValType>
  KDTreeNode<N, ValType> *KDTree<N,ValType>::buildTree(KDCell<N,ValType> **cell0,
							 uint32_t Ncells,
							 uint32_t depth,
							 coords minBound,
							 coords maxBound)
  {
    if (Ncells == 0)
      return 0;

    int axis = depth % N;
    KDTreeNode<N,ValType> *node = &nodes[lastNode++];
    uint32_t mid = Ncells/2;
    coords tmpBound;
    
    // Isolate the environment
    {
      CellCompare<N,ValType> compare(axis);
      std::sort(cell0, cell0+Ncells, compare);
    }

    node->value = *(cell0+mid);
    memcpy(&node->minBound[0], &minBound[0], sizeof(coords));
    memcpy(&node->maxBound[0], &maxBound[0], sizeof(coords));

    memcpy(tmpBound, maxBound, sizeof(coords));
    tmpBound[axis] = node->value->coord[axis];

    depth++;
    node->children[0] = buildTree(cell0, mid, depth, minBound, tmpBound);
    
    memcpy(tmpBound, minBound, sizeof(coords));
    tmpBound[axis] = node->value->coord[axis];
    node->children[1] = buildTree(cell0+mid+1, Ncells-mid-1, depth, 
				  tmpBound, maxBound);

    return node;
  }

  template<int N, typename ValType>
  uint32_t KDTree<N,ValType>::countActives() const
  {
    uint32_t numActive = 0;
    for (uint32_t i = 0; i < lastNode; i++)
      {
	if (nodes[i].value->active)
	  numActive++;
      }
    return numActive;
  }

  template<int N, typename ValType>
  typename KDDef<N>::CoordType
  KDTree<N,ValType>::computeDistance(KDCell<N, ValType> *cell, const coords& x)
  {
    CoordType d2 = 0;

    for (int i = 0; i < N; i++)
      {
	CoordType delta = cell->coord[i] - x[i];
	d2 += delta*delta;
      }
    return d2;
  }

  template<int N, typename ValType>
  void
  KDTree<N,ValType>::recursiveNearest(
				       KDTreeNode<N,ValType> *node,
				       int level,
				       const coords& x,
				       CoordType& R2,
				       KDCell<N,ValType> *& best)
  {
    CoordType d2 = 0;
    int axis = level % N;
    KDTreeNode<N,ValType> *other, *go;

    if (x[axis] < node->value->coord[axis])
      {
	// The best is potentially in 0.
	go = node->children[0];
	other = node->children[1];
      }
    else
      {
	// If not it is in 1.
	go = node->children[1];
	other = node->children[0];
	if (go == 0)
	  {
	    go = other;
	    other = 0;
	  }
      }

    if (go != 0)
      {
	recursiveNearest(go, level+1,
			 x, R2,best);
      }
    else
      {
	CoordType thisR2 = computeDistance(node->value, x);
	if (thisR2 < R2)
	  {
	    R2 = thisR2;
	    best = node->value;
	  }
	return;
      }
    
    // Check if current node is not the nearest
    CoordType thisR2 = 
      computeDistance(node->value, x);
    
    if (thisR2 < R2)
      {
	R2 = thisR2;
	best = node->value;
      }    
    
    // Now we found the best. We check whether the hypersphere
    // intersect the hyperplane of the other branch

    CoordType delta1;

    delta1 = x[axis]-node->value->coord[axis];
    if (delta1*delta1 < R2)
      {
	// The hypersphere intersects the hyperplane. Try the
	// other branch
	if (other != 0)
	  {
	    recursiveNearest(other, level+1, x, R2, best);
	  }
      }
  }

 template<int N, typename ValType>
 KDCell<N, ValType> *
 KDTree<N,ValType>::getNearestNeighbour(const coords& x)
 {
   CoordType R2 = INFINITY;
   KDCell<N,ValType> *best = 0;

   recursiveNearest(root, 0, x, R2, best);   
   
   return best;
 }

  template<int N, typename ValType>
  void
  KDTree<N,ValType>::recursiveMultipleNearest(RecursionMultipleInfo<N,ValType>& info, KDTreeNode<N,ValType> *node,
					       int level)
  {
    CoordType d2 = 0;
    int axis = level % N;
    KDTreeNode<N,ValType> *other, *go;

    if (info.x[axis] < node->value->coord[axis])
      {
	// The best is potentially in 0.
	go = node->children[0];
	other = node->children[1];
      }
    else
      {
	// If not it is in 1.
	go = node->children[1];
	other = node->children[0];
	if (go == 0)
	  {
	    go = other;
	    other = 0;
	  }
      }

    if (go != 0)
      {
	recursiveMultipleNearest(info, go, level+1);
      }
    
    // Check if current node is not the nearest
    CoordType thisR2 = 
      computeDistance(node->value, info.x);
    info.queue.push(node->value, thisR2);
    info.traversed++;
    if (go == 0)
      return;
    
    // Now we found the best. We check whether the hypersphere
    // intersect the hyperplane of the other branch

    CoordType delta1;

    delta1 = info.x[axis]-node->value->coord[axis];
    if (delta1*delta1 < info.queue.getMaxPriority())
      {
	// The hypersphere intersects the hyperplane. Try the
	// other branch
	if (other != 0)
	  {
	    recursiveMultipleNearest(info, other, level+1);
	  }
      }    
  }

 template<int N, typename ValType>
 void KDTree<N,ValType>::getNearestNeighbours(const coords& x, uint32_t N2,
					       KDCell<N, ValType> **cells)
 {
   RecursionMultipleInfo<N,ValType> info(x, cells, N2);
   
   for (int i = 0; i < N2; i++)
     cells[i] = 0;

   recursiveMultipleNearest(info, root, 0);

   std::cout << "Traversed = " << info.traversed << std::endl;
 }

 template<int N, typename ValType>
 void KDTree<N,ValType>::getNearestNeighbours(const coords& x, uint32_t N2,
					      KDCell<N, ValType> **cells,
					      CoordType *distances)
 {
   RecursionMultipleInfo<N,ValType> info(x, cells, N2);
   
   for (int i = 0; i < N2; i++)
     cells[i] = 0;

   recursiveMultipleNearest(info, root, 0);
   memcpy(distances, info.getPriorities(), sizeof(CoordType)*N2);

   std::cout << "Traversed = " << info.traversed << std::endl;
 }

};