#include "eytzinger_aabb_winding_number_tree.h"
#include "matlab_format.h"

#include <Eigen/Sparse>
#include <cassert>
#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>

template <
  typename DerivedE,
  typename Derivedleaf,
  typename DerivedI,
  typename DerivedC>
IGL_INLINE void igl::eytzinger_aabb_winding_number_tree(
  const Eigen::MatrixBase<DerivedE> & E,
  const Eigen::MatrixBase<Derivedleaf> & leaf,
  Eigen::PlainObjectBase<DerivedI> & I,
  Eigen::PlainObjectBase<DerivedC> & C)
{
  const int num_vertices = E.maxCoeff() + 1;

  // helper to increment or decrement count
  const auto add_into = [](
      std::unordered_map<int,int> & count_map,
      const int v, const int delta)
  {
    auto it = count_map.find(v);
    if(it == count_map.end())
    {
      count_map[v] = delta;
    }else
    {
      it->second += delta;
    }
  };

  // I claim this is still O(n log n) even though it seems like it's 
  // recomputing each subtree O(log n) times.
  std::vector<std::vector<int>> vI(leaf.size());
  for(int r = 0;r<leaf.size();r++)
  {
    if(leaf(r) == -2) { continue; }
    if(leaf(r) >= 0){ continue; }

    // Use unordered map to count contributions
    std::unordered_map<int,int> count_map;

    std::vector<int> stack;
    stack.push_back(r);
    while(!stack.empty())
    {
      const int i = stack.back();
      stack.pop_back();
      if(leaf(i) >= 0)
      {
        // leaf node
        add_into(count_map,E(leaf(i),0),1);
        add_into(count_map,E(leaf(i),1),-1);
      }else
      {
        const int left_i = 2*i + 1;
        const int right_i = 2*i + 2;
        stack.push_back(right_i);
        stack.push_back(left_i);
      }
    }

    std::vector<int> positive;
    std::vector<int> negative;
    for(const auto & it : count_map)
    {
      if(it.second != 0)
      {
        for(int c = 0;c<std::abs(it.second);c++)
        {
          if(it.second > 0)
          {
            positive.push_back(it.first);
          }else
          {
            negative.push_back(it.first);
          }
        }
      }
    }
    assert(positive.size() == negative.size());
    std::vector<int> alternating;
    for(int k = 0;k<positive.size();k++)
    {
      alternating.push_back(positive[k]);
      alternating.push_back(negative[k]);
    }
    vI[r] = alternating;
  }
 
  // first collect cummulative sizes
  C.resize(leaf.size()+1);
  C(0) = 0;
  for(int r = 0;r<leaf.size();r++)
  {
    assert(vI[r].size() % 2 == 0);
    C(r+1) = C(r) + vI[r].size();
  }
  I.resize(C(leaf.size()));
  for(int r = 0;r<leaf.size();r++)
  {
    for(int k = 0;k<vI[r].size();k++)
    {
      I(C(r)+k) = vI[r][k];
    }
  }

}

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