#include "eytzinger_aabb.h"
#include "PlainMatrix.h"
#include "median.h"
#include <vector>
#include <algorithm>

template <
  typename DerivedPB,
  typename DerivedB,
  typename Derivedleaf
>
IGL_INLINE void igl::eytzinger_aabb(
  const Eigen::MatrixBase<DerivedPB> & PB1,
  const Eigen::MatrixBase<DerivedPB> & PB2,
  Eigen::PlainObjectBase<DerivedB> & B1,
  Eigen::PlainObjectBase<DerivedB> & B2,
  Eigen::PlainObjectBase<Derivedleaf> & leaf)
{
  using Scalar = typename DerivedPB::Scalar;
  using VectorXS = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
  // Number of primitives
  const int m = PB1.rows();
  assert(PB1.rows() == PB2.rows());
  if(m==0)
  {
    B1.resize(0,PB1.cols());
    B2.resize(0,PB2.cols());
    leaf.resize(0);
    return;
  }
  const PlainMatrix<DerivedPB> PBC = (PB1 + PB2) / 2;
  const int complete_m = 1 << ((int)std::ceil(std::log2(m))+1);
  B1.resize(complete_m, PB1.cols());
  B2.resize(complete_m, PB2.cols());
  leaf.resize(complete_m);
  leaf.setConstant(-2);
  std::vector<int> I(m); for(int i = 0; i < m; i++) { I[i] = i; }

  std::function<void(const int, const std::vector<int> &)> recursive_helper;
  recursive_helper = [&](const int i, const std::vector<int> & I)
  {
    if(I.size() == 0) { return; }

    B1.row(i) = PB1(I, Eigen::all).colwise().minCoeff();
    B2.row(i) = PB2(I, Eigen::all).colwise().maxCoeff();

    if(I.size() == 1) { leaf(i) = I[0]; return; }

    leaf(i) = -1;
    int dir;
    (B2.row(i) - B1.row(i)).maxCoeff(&dir);

    Scalar split_value;
    VectorXS Z = PBC(I, dir);
    igl::median(Z,split_value); 
    std::vector<int> left; left.reserve(I.size()/2+1);
    std::vector<int> right; right.reserve(I.size()/2+1);
    for(int j = 0; j < (int)I.size(); j++)
    {
      if(Z(j) < split_value)
      {
        left.push_back(I[j]);
      }else if(Z(j) > split_value)
      {
        right.push_back(I[j]);
      }
    }
    for(int j = 0; j < (int)I.size(); j++)
    {
      if(Z(j) == split_value)
      {
        (left.size()<right.size() ? left : right).push_back(I[j]);
      }
    }
    assert(std::abs(int(left.size()-right.size())) <= 1);

    const int left_i = 2*i + 1;
    const int right_i = 2*i + 2;
    recursive_helper(left_i, left);
    recursive_helper(right_i, right);
  };
  recursive_helper(0,I);
}

#ifdef IGL_STATIC_LIBRARY
// Explicit template instantiation
// generated by autoexplicit.sh
template void igl::eytzinger_aabb<Eigen::Matrix<double, -1, 3, 1, -1, 3>, Eigen::Matrix<double, -1, 3, 1, -1, 3>, Eigen::Matrix<int, -1, 1, 0, -1, 1>>(Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>> const&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>>&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>>&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, 1, 0, -1, 1>>&);
template void igl::eytzinger_aabb<Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<int, -1, 1, 0, -1, 1>>(Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>>&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>>&, Eigen::PlainObjectBase<Eigen::Matrix<int, -1, 1, 0, -1, 1>>&);
#endif