#include <Eigen/Dense>
#include <algorithm>
#include <array>
#include <bitset>
#include <deque>
#include <execution>
#include <forward_list>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <memory>
#include <mutex>
#include <queue>
#include <random>
#include <shared_mutex>
#include <tuple>
#include <unordered_set>

template <unsigned int dim, typename number, typename T> class Orthtree {
public:
  typedef Eigen::Vector<number, dim> Point;
  typedef std::function<bool(std::unique_ptr<T> &, const Point &)>
      data_access_f;
  typedef std::function<bool(const std::unique_ptr<T> &, const Point &)>
      cdata_access_f;

  class BBox {
  public:
    enum Relation {
      NO_OVERLAP,
      OVERLAP,
      CONTAINS,
      CONTAINED,
    };
    Point min, max;
    Point center() { return (min + max) / 2; }
    Relation relation(const BBox &b) {
      if ((b.min.array() > max.array()).any() ||
          (b.max.array() < min.array()).any())
        return NO_OVERLAP;
      if ((b.min.array() > min.array()).all() &&
          (b.max.array() < max.array()).all())
        return CONTAINS;
      if ((b.min.array() < min.array()).all() &&
          (b.max.array() > max.array()).all())
        return CONTAINED;
      return OVERLAP;
    };
    BBox(const Point &min, const Point &max) : min(min), max(max) {
      assert((min.array() < max.array()).all());
    }
  };
  typedef std::bitset<dim> coord;
  class Node;
  struct NodeInfo {
    const Node *node;
    int depth;
    number width;
    coord crd;
  };
  typedef std::forward_list<NodeInfo> branch;
  class Node {
  public:
    Point center;
    std::unique_ptr<std::array<Node, (1 << dim)>> children;
    std::unique_ptr<T> data;

    mutable std::shared_mutex mtx;

    void create_children(const Point &parent_center, number parent_width) {
      children = std::make_unique<std::array<Node, (1 << dim)>>();
      for (unsigned int i = 0; i < (1 << dim); ++i) {
        Point offset = Point::Ones() * (parent_width / 4);
        for (unsigned int j = 0; j < dim; ++j)
          if (!coord(i)[j])
            offset(j) *= -1;
        (*children)[i].center = parent_center + offset;
      }
    };

    coord get_child_coord(const Point &p) const {
      coord child_coord;
      for (unsigned int i = 0; i < dim; ++i)
        child_coord[i] = center[i] < p[i];
      return child_coord;
    };

    void depth(const Point &point, int depth, data_access_f f, number width) {
      if (depth == 0) {
        std::unique_lock guard(mtx);
        f(data, center);
      } else {
        mtx.lock_shared();
        if (!children) {
          mtx.unlock_shared();
          {
            std::unique_lock uguard(mtx);
            if (!children)
              create_children(center, width);
          }
          mtx.lock_shared();
        }
        (*children)[get_child_coord(point).to_ulong()].depth(point, depth - 1,
                                                             f, width / 2);

        mtx.unlock_shared();
      }
    };

    BBox bb(number width) const {
      Point offset = Point::Ones() * (width / 2);
      return BBox(center - offset, center + offset);
    };

    void access_data(cdata_access_f f) const {
      std::unique_lock guard(mtx);
      if (data) {
        f(data, center);
      }
    }
    bool access_data_ret(cdata_access_f f, bool &ret) const {
      std::unique_lock guard(mtx);
      if (data) {
        ret = f(data, center);
        return true;
      }
      return false;
    }
    void all_data(cdata_access_f f) const {
      std::unique_lock guard(mtx);
      if (data) {
        f(data, center);
      }
      if (children) {
        std::for_each(std::execution::par_unseq, std::begin(*children),
                      std::end(*children),
                      [f](const auto &c) { c.all_data(f); });
      }
    }
    bool is_leaf() const {
      std::shared_lock guard(mtx);
      return !bool(children);
    }
    void nodes_to_leaf(const Point &p, const number width, const int depth,
                       branch &l) const {
      std::shared_lock guard(mtx);
      if (!children)
        return;
      auto c = get_child_coord(p);
      auto child = &(*children)[c.to_ullong()];
      l.push_front({child, depth - 1, width / 2, c});
      child->nodes_to_leaf(p, width / 2, depth - 1, l);
    }
  };
  std::unique_ptr<Node> root;
  number root_width;

  mutable std::mutex mtx;
  Orthtree(const Point &root_center, number root_width)
      : root(new Node()), root_width(root_width) {
    root->center = root_center;
  };

  void all_data(cdata_access_f f) const { root->all_data(f); };

  void depth(const Point &point, int depth, data_access_f f) {
    root->depth(point, depth, f, root_width);
  }

  // std::vector<NodeInfo> const std::forward_list<std::pair<Node *, coord>>&
  // branch)

  branch get_branch_to_leaf(const Point &p) const {
    branch b({root.get, 0, root_width, coord{0}});
    root->node_to_leaf(p, root_width, b);
    return b;
  }

  // std::vector<branch> get_leaf_neighours(branch b) {
  //   auto node = b.front();
  //   b.pop_front();
  //   if (b.empty())
  //     return {};
  //   std::vector<branch> ret;
  //   ret.reserve(2<<dim);
  //   for (int i = 0; i < (1<<dim); i += 1) {
  //     auto& c = *b.front()->children;
  //     if (coord{i} != node->coord) {
  //       ret.push_back(b);
  //       ret.back().push_front(&c);
  //     }
  //   }
  //   return ret;
  // }
  template<typename Container>
  void get_nodes_in_dir(const branch b, int dir,
                        std::queue<branch,Container> &nodes) const {
    nodes.push_back(b);
    const unsigned int dim_dir = std::abs(dir) - 1;
    if (!b.front().node->children)
      return;
    for (int i = 0; i < 1 << dim; i++) {
      if (coord{i}[dim_dir] == dir > 0)
        continue;
      branch child_branch(b.begin(), b.end());
      child_branch.push_front({b.front().node->children->at(i),
                               b.front().depth - 1, b.front().width / 2,
                               coord{i}});
      get_nodes_in_dir(child_branch, dir, nodes);
    }
  }

  template<typename Container>
  void get_adjacent_nodes(branch b, int dir,
                          std::queue<branch,Container> &ad_nodes) const {
    auto n = b.first();
    const unsigned int dim_dir = std::abs(dir) - 1;
    int up = 1;
    auto b_it = b.begin();
    if (std::next(b_it) == b.end())
      return;
    for (; b_it.crd[dim_dir] == dir > 0; std::advance(b_it, 1)) {
      if (std::next(b_it) == b.end())
        return;
      up++;
    }
    branch c(b.begin(), std::next(b_it));
    branch p(std::next(b_it), b.end()); // TODO remove from b not cpy
    c.reverse();
    for (int i = 0; i < up; i += 1) {
      coord pc = coord(c.front().crd).flip(dim_dir);
      p.push_front({c.front().node->children->at(pc), c.front().depth - i,
                    c.front().width / (1 << i), pc});
      c.pop_front();
    }
    get_nodes_in_dir(p, -dir, ad_nodes);
  }

  void datas_by_distance(const Point &p, cdata_access_f f) const {
    branch b = get_branch_to_leaf(p);
    std::deque<branch> to_visit{b};
    std::unordered_set<const Node *> visited;
    while (!to_visit.empty()) {
      const branch& curr_visit = to_visit.front();
      for (const auto &n : curr_visit) {
        if (visited.contains(n.node))
          break;
        bool ret;
        if (n.node->data_acces_ret(f, ret) && !ret)
          return;
        visited.insert(n.node);
      }
      to_visit.pop_front();
      for (int i = 0; i < dim; i += 1) {
        get_adjacent_nodes(curr_visit,dim+1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
        get_adjacent_nodes(curr_visit,-dim-1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
      }
    }
  }
  // void datas_by_distance(const Point &p, cdata_access_f f) const {
  //   branch b = get_branch_to_leaf(p);
  //   std::deque<std::pair<Point, branch>> to_visit{{p, b}};
  //   std::unordered_set<Node *> visited;
  //   while (!to_visit.empty()) {
  //     for (auto &n : to_visit.second.front()) {
  //       if (visited.contains(n))
  //         break;
  //       bool ret;
  //       if (n.node->data_acces_ret(f, ret) && !ret)
  //         return;
  //       visited.insert(n.node);
  //     }
  //     auto search_p = to_visit.front().first;
  //     to_visit.pop_front();
  //     for (int i = 0; i < dim; i += 1) {
  //       Point offset = Point::Zero();
  //       offset(i) = b.width;
  //       auto p1 = search_p + offset;
  //       auto p2 = search_p + offset;
  //       to_visit.push_back({p1, get_branch_to_leaf(p1)});
  //       to_visit.push_back({p2, get_branch_to_leaf(p2)});
  //     }
  //   }
  // }

  void data_in_bb(const BBox &bb, cdata_access_f f) const {
    number curr_width = root_width;
    std::vector<const Node *> to_visit{root.get()};
    std::vector<const Node *> to_visit_children;
    bool stop = false;
    while (!to_visit.empty() && !stop) {
      auto cbb = to_visit.back()->bb(curr_width);
      auto bbrel = cbb.relation(bb);
      // std::cout << to_visit.back()->center.transpose() << " - "
      // << cbb.min.transpose() << " - " << bbrel << std::endl;
      bool ret;
      switch (bbrel) {
      case (Orthtree::BBox::CONTAINED):
        to_visit.back()->all_data(f);
        break;
      case (Orthtree::BBox::OVERLAP):
      case Orthtree::BBox::CONTAINS:
        to_visit.back()->access_data_ret(f, ret);
        if (to_visit.back()->access_data_ret(f, ret) && !ret)
          stop = true;
        if (!to_visit.back()->is_leaf()) {
          // std::transform(std::execution::par_unseq,to_visit.back()->children->cbegin(),
          //                to_visit.back()->children->cend(),
          //                to_visit_children.begin(), std::addressof<const
          //                Node>);
          for (const auto &n : *to_visit.back()->children) {
            to_visit_children.push_back(&n);
          }
          // std::cout << "leaf "<<to_visit_children.size() << std::endl;
        }
        break;
      default:
        break;
      }
      to_visit.pop_back();
      if (to_visit.empty()) {
        to_visit.swap(to_visit_children);
        to_visit_children.clear();
        curr_width /= 2;
      }
    }
  }
};