//
//  BBoxTree.hpp
//  Hexer
//
//  Created by Andreas Bærentzen on 22/02/2021.
//

#ifndef BBoxTree_hpp
#define BBoxTree_hpp

#include <GEL/CGLA/Vec3d.h>
#include <GEL/Geometry/GridAlgorithm.h>
#include <algorithm>
#include <limits>
#include <set>
#include <unordered_map>

/** A simple bounding box class.
 This class represents an axis aligned bounding box. For many applications it
 might be most useful as a parent
 class of a bounding box containing a payload. Meanwhile, instances of BBox can
 also be used in queries. */
struct BBox {
    CGLA::Vec3d min_corner = CGLA::Vec3d(std::numeric_limits<double>::max());
    CGLA::Vec3d max_corner = CGLA::Vec3d(std::numeric_limits<double>::lowest());

    /// Expanding the bounding box to contain point, p, passed as argument.
    void extend(const CGLA::Vec3d &p) {
        min_corner = CGLA::v_min(min_corner, p);
        max_corner = CGLA::v_max(max_corner, p);
    }

    /// Expanding bounding box to contain the other_box passed as argument.
    void extend(const BBox &other_box) {
        min_corner = CGLA::v_min(min_corner, other_box.min_corner);
        max_corner = CGLA::v_max(max_corner, other_box.max_corner);
    }

    /// Returns true if the other_box has an overlap with the present box.
    bool overlap(const BBox &other_box) const {
        if (other_box.max_corner.all_ge(min_corner) &&
            other_box.min_corner.all_le(max_corner))
            return true;
        return false;
    }

    /// Returns true if the point has an overlap with the present box.
    bool contains(const CGLA::Vec3d &point) const {
        if (point.all_ge(min_corner) && point.all_le(max_corner)) return true;
        return false;
    }

    /// Returns centroid of bounding box.
    CGLA::Vec3d centroid() const { return 0.5 * (min_corner + max_corner); }

    /// Returns diagonal vector of bounding box
    CGLA::Vec3d diagonal() const { return max_corner - min_corner; }

    /** Comparison operator. Returns true if the minimum corners x value is
       smaller than other_box.
        This is useful if you are making a sorted list of boxes. */
    bool operator<(const BBox &other_box) const {
        if (min_corner[0] < other_box.min_corner[0]) return true;
        return false;
    }
};

/** A spatial data structure for bounding boxes.
 This template is used to efficiently find bounding boxes that overlap. The
 template argument is meant to be derived from BBox. Constructed with a vector
 of such bounding boxes, an instantiated class provides facilities for quickly
 locating the overlapping bounding boxes. It also provides access to the
 elements via begin and end functions and an index operator. Internally this
 class relies on the unordered_multimap class from the C++ standard library
 which would usually be implemented as a hash map that allows several values for
 each key.
 */
template <typename BBoxT> class BBoxMap {
  public:
    using IterType      = typename std::vector<BBoxT>::iterator;
    using ConstIterType = typename std::vector<BBoxT>::const_iterator;

  private:
    /** Hash function for vectors. This function is designed to combine hashes
     for each dimension and returns a single combined hash. The function works
     in  principle for any number of dimensions although the intent is that is
     should be used for d=2,3, or 4. */
    struct HashVec3i {
        size_t operator()(const CGLA::Vec3i &v) const {
            const size_t k = 83492791;
            size_t       h = v[0];
            for (unsigned int i = 1; i < v.get_dim(); ++i) {
                h *= k;
                h += v[i];
            }
            return h;
        }
    };
    using IndexMapType =
        std::unordered_multimap<CGLA::Vec3i, size_t, HashVec3i>;

    std::vector<BBoxT> boxes; // vector containing the actual data
    double             scale; // scaling used when going from float to integer
    IndexMapType index_map;   // mapping from the (hashed) spatial indices to
                              // linear indices used for boxes.

    /** Convert 3D floating point vector to 3D integer vector useful for
     * indexing a point in space. */
    CGLA::Vec3i disc_vec(const CGLA::Vec3d &x) const {
        return CGLA::Vec3i(scale * x);
    }

    /// Given a bounding box, this function returns a range of nD indices
    Geometry::Range3D bbox_range(const BBox &b) const {
        const CGLA::Vec3i pl(disc_vec(b.min_corner));
        const CGLA::Vec3i ph =
            CGLA::Vec3i(1) + CGLA::Vec3i(disc_vec(b.max_corner));
        return Geometry::Range3D(pl, ph);
    }
    /// Given a bounding box, this function returns a range of nD indices
    Geometry::Range3D bbox_range(const CGLA::Vec3d &p) const {
        const CGLA::Vec3i pl(disc_vec(p - CGLA::Vec3d(1e-3)));
        const CGLA::Vec3i ph =
            CGLA::Vec3i(1) + CGLA::Vec3i(disc_vec(p + CGLA::Vec3d(1e-3)));
        return Geometry::Range3D(pl, ph);
    }

    /// Given the vector index of a bbox, this function inserts spatial indices
    /// in the index_map.
    void create_indices(size_t idx) {
        for (const CGLA::Vec3i &p : bbox_range(boxes[idx]))
            index_map.insert({p, idx});
    }

  public:
    /** Constructor which is provided with a vector of bounding boxes. This
       constructor copies
        the data to an internal representation and builds the spatial indices.
     */
    BBoxMap(const std::vector<BBoxT> &_boxes) : boxes(_boxes) {
        double sum_side_lengths = 0;
        for (const auto &b : boxes)
            sum_side_lengths += dot(CGLA::Vec3d(1), b.diagonal());
        scale = 0.7 * (boxes.size() * 3) / sum_side_lengths;
        index_map.clear();
        for (size_t i = 0; i < boxes.size(); ++i) create_indices(i);
    }

    /** Insert a single box in the data structure. Having called this function,
     the data structure is ready for new queries where box could potentially be
     found. Just be aware that there is no check whether box fits in the
     bounding box provided to the constructor or computed in build. Thus, the
     responsibility is on the user for ensuring this. ---> this behaviour might
     have to change to something safer ... in principle we could index outside
     the overall bounding box since the grid is hash based. It is just that in
     this case the indices are not unique. Perhaps the soundest way of dealing
     with this is by using the vector type as key in the
     unordered_multimap and then providing a hash function... <----  */
    void insert(const BBoxT &box) {
        size_t idx = boxes.size();
        boxes.push_back(box);
        create_indices(idx);
    }

    /// Index operator providing access to elements.
    BBoxT &operator[](size_t idx) { return boxes[idx]; }

    /// Const index operator providing const access to elements
    const BBoxT &operator[](size_t idx) const { return boxes[idx]; }

    /// Begin returns iterator to first element in array of boxes
    IterType begin() { return boxes.begin(); }

    /// End returns iterator to one past last elemet in array of boxes
    IterType end() { return boxes.end(); }

    /** This function returns a vector of indices to the boxes that overlap with
     the box provided as argument. Access to the relevant boxes can then be
     obtained using the above index operators. */
    std::vector<size_t> overlap(const BBox &box) const {
        std::set<size_t> indices;
        for (const CGLA::Vec3i &p : bbox_range(box)) {
            auto range = index_map.equal_range(p);
            for (auto it = range.first; it != range.second; ++it)
                indices.insert(it->second);
        }
        std::vector<size_t> idx_out;
        for (auto idx : indices) {
            if (boxes[idx].overlap(box)) idx_out.push_back(idx);
        }
        return idx_out;
    }

    /** This function returns a vector of indices to the boxes that overlap with
     the point provided as argument. Access to the relevant boxes can then be
     obtained using the above index operators. */
    std::vector<size_t> contains(const CGLA::Vec3d &point) const {
        std::set<size_t> indices;
        for (const CGLA::Vec3i &p : bbox_range(point)) {
            auto range = index_map.equal_range(p);
            for (auto it = range.first; it != range.second; ++it)
                indices.insert(it->second);
        }
        std::vector<size_t> idx_out;
        for (auto idx : indices) {
            if (boxes[idx].contains(point)) idx_out.push_back(idx);
        }
        return idx_out;
    }

    /** Returns a set of index pairs containing all pairs of overlapping boxes.
     The reason for a set is that boxes are likely to overlap in a number of
     grid cells, so to make sure that each pair is reported only once, we use a
     set rather than a vector. */
    std::set<std::pair<size_t, size_t>> overlapping_pairs() const {
        std::set<std::pair<size_t, size_t>> pair_set;
        for (size_t i = 0; i < index_map.bucket_count(); ++i)
            for (auto iter0 = index_map.begin(i); iter0 != index_map.end(i);
                 ++iter0)
                for (auto iter1 = index_map.begin(i); iter1 != index_map.end(i);
                     ++iter1) {
                    const auto bidx0 = iter0->second;
                    const auto bidx1 = iter1->second;
                    if (bidx0 < bidx1 && boxes[bidx0].overlap(boxes[bidx1]))
                        pair_set.insert(
                            std::make_pair(iter0->second, iter1->second));
                }
        return pair_set;
    }
};

#endif /* BBoxTree_hpp */
