// based on Eigen unsupported bvh KDTree

#pragma once

/** \file */

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/StdVector>
#include <algorithm>
#include <queue>

#include "hgeom/bvh/bvh_algo.hpp"
#include "hgeom/geom/primitive.hpp"
#include "hgeom/util/types.hpp"

/**
\namespace hgeom
\brief top level namespace
*/
namespace hgeom {
/**
\namespace hgeom::bvh
\brief Bounding volume hierarchy namespace
*/
namespace bvh {

using namespace geom;
using namespace Eigen;

// internal pair class for the BVH--used instead of std::pair because of
// alignment
template <class F, int DIM> struct VintPair {
    using first_type = Matrix<F, DIM, 1>;
    using secont_type = int;
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW_IF_VECTORIZABLE_FIXED_SIZE(F, DIM)
    VintPair(const first_type &v, int i) : first(v), second(i) {}
    first_type first;
    int second;
};

// these templates help the tree initializer get the bounding spheres either in
// from a provided iterator range or using bounding_vol in a unified way
template <typename Objs, typename Vols, typename Iter> struct get_bvols_helper {
    void operator()(const Objs &objs, Iter sphbeg, Iter sphend, Vols &out) {
        out.insert(out.end(), sphbeg, sphend);
        eigen_assert(out.size() == objs.size());
    }
};

template <typename Objs, typename Vols>
struct get_bvols_helper<Objs, Vols, int> {
    void operator()(const Objs &objs, int, int, Vols &out) {
        out.reserve(objs.size());
        for (int i = 0; i < (int)objs.size(); ++i)
            out.push_back(bounding_vol(objs[i]));
    }
};

template <class RAiter> struct P1Range {
    using value_type = typename RAiter::value_type::first_type;
    RAiter a, b;
    auto const &operator[](size_t i) const { return (a + i)->first; }
    auto &operator[](size_t i) { return (a + i)->first; }
    size_t size() const { return b - a; }
    auto get_index(size_t i) const { return (a + i)->second; }
};
template <class RAiter> auto p1range(RAiter a, RAiter b) {
    P1Range<RAiter> r;
    r.a = a;
    r.b = b;
    return r;
}

template <typename F, bool idrange = false> struct WelzlBoundingSphere {
    template <typename SubtreeObjs>
    static Sphere<F> bound(SubtreeObjs subtree_objs) {
        return welzl_bounding_sphere<true>(subtree_objs);
    }
};

template <typename _Scalar, typename _Object, int _DIM = 3,
          typename _Volume = Sphere<_Scalar>,
          typename BoundingSphere = WelzlBoundingSphere<_Scalar, true>>
class SphereBVH {
  public:
    static int const DIM = _DIM;
    typedef _Object Object;
    typedef std::vector<Object, Eigen::aligned_allocator<Object>> Objs;
    typedef _Scalar F;
    // typedef Eigen::AlignedBox<F, DIM> Volume;
    typedef _Volume Volume;
    typedef std::vector<Volume, Eigen::aligned_allocator<Volume>> Vols;
    typedef int Index;
    typedef const int *VolumeIterator; // the iterators are just pointers into
                                       // the tree's vectors
    typedef const Object *ObjectIterator;

    std::vector<int> child; // child of x are child[2x] and
                            // child[2x+1], indices bigger than
                            // vols.size() index into objs.
    Vols vols;
    Objs objs;

    SphereBVH() {}

    template <typename Iter> SphereBVH(Iter begin, Iter end) {
        init(begin, end, 0, 0);
    } // int is recognized by init as not being an iterator type

    template <typename OIter, typename BIter>
    SphereBVH(OIter begin, OIter end, BIter sphbeg, BIter sphend) {
        init(begin, end, sphbeg, sphend);
    }

    size_t size() const { return objs.size(); }

    /** Given an iterator range over \a Object references, constructs the BVH,
     * overwriting whatever is in there currently.
     * Requires that bounding_vol(Object) return a Volume. */
    template <typename Iter> void init(Iter begin, Iter end) {
        init(begin, end, 0, 0);
    }

    /** Given an iterator range over \a Object references and an iterator range
     * over their bounding vols,
     * constructs the BVH, overwriting whatever is in there currently. */
    template <typename OIter, typename BIter>
    void init(OIter begin, OIter end, BIter sphbeg, BIter sphend) {
        objs.clear();
        vols.clear();
        child.clear();

        objs.insert(objs.end(), begin, end);
        int n = static_cast<int>(objs.size());

        // if we have at most one object, we don't need any internal nodes
        if (n < 2) return;

        Vols ovol;
        VIPairs ocen;

        // compute the bounding vols depending on BIter type
        get_bvols_helper<Objs, Vols, BIter>()(objs, sphbeg, sphend, ovol);

        ocen.reserve(n);
        vols.reserve(n - 1);
        child.reserve(2 * n - 2);

        for (int i = 0; i < n; ++i) ocen.push_back(VIPair(ovol[i].cen, i));

        // the recursive part of the algorithm
        build(ocen, 0, n, ovol, 0);

        Objs tmp(n);
        tmp.swap(objs);
        for (int i = 0; i < n; ++i) objs[i] = tmp[ocen[i].second];
    }

    /** \returns the index of the root of the hierarchy */
    inline Index getRootIndex() const { return (int)vols.size() - 1; }

    /** Given an \a index of a node, on exit, \a vbeg and \a vend range
     * over the indices of the volume children of the node
     * and \a obeg and \a oend range over the object children of the node
     */
    EIGEN_STRONG_INLINE
    void getChildren(Index index, VolumeIterator &vbeg, VolumeIterator &vend,
                     ObjectIterator &obeg, ObjectIterator &oend) const {
        // inlining this function should open lots of optimization opportunities
        // to the compiler
        if (index < 0) {
            vbeg = vend;
            if (!objs.empty()) obeg = &(objs[0]);
            oend = obeg + objs.size(); // output all objs--necessary
                                       // when the tree has only one
                                       // object
            return;
        }

        int nvol = static_cast<int>(vols.size());

        int idx = index * 2;
        if (child[idx + 1] < nvol) { // second index is always bigger
            vbeg = &(child[idx]);
            vend = vbeg + 2;
            obeg = oend;
        } else if (child[idx] >= nvol) { // if both child are objs
            vbeg = vend;
            obeg = &(objs[child[idx] - nvol]);
            oend = obeg + 2;
        } else { // if the first child is a volume and the second is an object
            vbeg = &(child[idx]);
            vend = vbeg + 1;
            obeg = &(objs[child[idx + 1] - nvol]);
            oend = obeg + 1;
        }
        // std::cout << " gc" << vend - vbeg;
    }

    inline const Volume &getVolume(Index index) const { return vols[index]; }

  private:
    typedef VintPair<F, DIM> VIPair;
    typedef std::vector<VIPair, Eigen::aligned_allocator<VIPair>> VIPairs;
    typedef Eigen::Matrix<F, DIM, 1> VectorType;

    struct AxisComparator {
        int dim;
        AxisComparator(int inDim) : dim(inDim) {}
        inline bool operator()(const VIPair &v1, const VIPair &v2) const {
            return v1.first[dim] < v2.first[dim];
        }
    };
    struct DotComparator {
        Eigen::Matrix<F, DIM, 1> normal;
        DotComparator(const Eigen::Matrix<F, DIM, 1> &n) : normal(n) {}
        inline bool operator()(const VIPair &v1, const VIPair &v2) const {
            return v1.first.dot(normal) < v2.first.dot(normal);
        }
    };
    // struct DotComparator {
    //     Matrix<F, DIM, 1> normal;
    //     DotComparator(Matrix<F, DIM, 1> n) : normal(n) {}
    //     template <class Pair>
    //     DotComparator(Pair p)
    //         : normal((p.second.eval() - p.first.eval()).eval()) {}
    //     // : normal((Eigen::Matrix<F, DIM, 1>(p.second) -
    //     // Eigen::Matrix<F, DIM, 1>(p.first))
    //     // .eval()) {}
    //     inline bool operator()(const VIPair &v1, const VIPair &v2) const {
    //         return v1.first.dot(normal) < v2.first.dot(normal);
    //     }
    // };

    // Build the part of the tree between objs[from] and objs[to] (not
    // including objs[to]). This routine partitions the ocen in [from, to) along
    // the dimension dim, recursively constructs the two halves, and adds their
    // parent node.  TODO: a cache-friendlier layout
    void build(VIPairs &ocen, int from, int to, Vols const &ovol,
               int dim) noexcept {
        eigen_assert(to - from > 1);
        if (to - from == 2) {
            auto merge =
                ovol[ocen[from].second].merged(ovol[ocen[from + 1].second]);
            vols.push_back(merge);
            child.push_back(from + (int)objs.size() - 1);
            child.push_back(from + (int)objs.size());
        } else if (to - from == 3) {
            int mid = from + 2;
            auto subtree_objs = p1range(ocen.begin() + from, ocen.begin() + to);
            auto sep = most_separated_points_on_AABB(subtree_objs);
            auto normal = (sep.second - sep.first).eval();
            nth_element(ocen.begin() + from, ocen.begin() + mid,
                        ocen.begin() + to, DotComparator(normal));
            // nth_element(
            // ocen.begin() + from, ocen.begin() + mid, ocen.begin() + to,
            // DotComparator(most_separated_points_on_AABB(subtree_objs)));
            // AxisComparator(dim));
            build(ocen, from, mid, ovol, (dim + 1) % DIM);
            int idx1 = (int)vols.size() - 1;
            Volume bound = BoundingSphere::bound(subtree_objs);
            vols.push_back(bound);
            // Volume merge = vols[idx1].merged(ovol[ocen[mid].second]);
            // if (merge.rad + 0.0001 < bound.rad)
            // std::cout << "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << std::endl;
            // vols.push_back(bound.rad < merge.rad ? bound : merge);
            child.push_back(idx1);
            child.push_back(mid + (int)objs.size() - 1);
        } else {
            int mid = from + (to - from) / 2;
            auto subtree_objs = p1range(ocen.begin() + from, ocen.begin() + to);
            auto sep = most_separated_points_on_AABB(subtree_objs);
            auto normal = (sep.second - sep.first).eval();
            nth_element(ocen.begin() + from, ocen.begin() + mid,
                        ocen.begin() + to, DotComparator(normal));
            // AxisComparator(dim));
            build(ocen, from, mid, ovol, (dim + 1) % DIM);
            int idx1 = (int)vols.size() - 1;
            build(ocen, mid, to, ovol, (dim + 1) % DIM);
            int idx2 = (int)vols.size() - 1;
            Volume bound = BoundingSphere::bound(subtree_objs);
            vols.push_back(bound);
            // Volume merge = vols[idx1].merged(vols[idx2]);
            // if (merge.rad + 0.0001 < bound.rad)
            // std::cout << "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" << std::endl;
            // vols.push_back(bound.rad < merge.rad ? bound : merge);
            child.push_back(idx1);
            child.push_back(idx2);
        }
    }
};
} // namespace bvh
} // namespace hgeom
