#include <library/cpp/testing/common/env.h>
#include <library/cpp/testing/gtest/gtest.h>
#include <maps/wikimap/mapspro/services/mrc/libs/fb_rtree/include/packer.h>
#include <maps/libs/geolib/include/spatial_relation.h>

namespace maps::mrc::fb_rtree::tests {
namespace {

using namespace geolib3;

using IdToPointMap = std::unordered_map<Id, Point2>;
using IdToBoxMap = std::unordered_map<Id, BoundingBox>;

LeafNodes grid(int width, int height, float probability)
{
    auto result = LeafNodes{};
    while (result.empty()) {
        for (int i = 0; i < width; ++i) {
            for (int j = 0; j < height; ++j) {
                if (rand() < RAND_MAX * probability) {
                    result.push_back({Id(i * height + j),
                                      Point2(i * 10, j * 10).boundingBox()});
                }
            }
        }
    }
    std::random_shuffle(result.begin(), result.end());
    return result;
}

IdToBoxMap makeIdToBoxMap(const LeafNodes& leaves)
{
    auto result = IdToBoxMap{};
    for (const auto& leaf : leaves) {
        result.insert({leaf.id, leaf.bbox});
    }
    return result;
}

void verifyNumberFound(const Rtree& rtree,
                       const BoundingBox& bbox,
                       IntersectsWithId intersects,
                       size_t expected)
{
    auto rng = rtree.allIdsInWindow(bbox, intersects);
    size_t size = std::distance(rng.begin(), rng.end());
    EXPECT_EQ(size, expected);
}

void verifyClosest(const IdToBoxMap& idToBoxMap,
                   const Rtree& rtree,
                   const Point2& point,
                   uint32_t count)
{
    auto distance = SquaredDistance{};
    auto distanceToId = [&](const geolib3::Point2& point, Id id) {
        return distance(point, idToBoxMap.at(id));
    };
    auto expectedDistances = std::vector<double>{};
    auto idToDistanceMap = std::unordered_map<Id, double>{};
    for (auto& [id, box] : idToBoxMap) {
        double dist = distance(point, box);
        expectedDistances.push_back(dist);
        idToDistanceMap.insert({id, dist});
    }
    std::sort(expectedDistances.begin(), expectedDistances.end());
    if (expectedDistances.size() > count) {
        expectedDistances.resize(count);
    }

    auto answer = nearestIds(rtree, point, distance, distanceToId, count);
    EXPECT_EQ(answer.size(), expectedDistances.size());
    for (size_t i = 0; i < answer.size(); ++i) {
        EXPECT_DOUBLE_EQ(idToDistanceMap.at(answer[i]), expectedDistances[i]);
    }
}

}  // anonymous namespace

TEST(rtree, base)
{
    // create the rtree
    auto leaves = LeafNodes{};
    for (int i = 0; i < 10; ++i) {
        leaves.push_back(
            {Id(i), BoundingBox(Point2(i, i), Point2(i + 0.5, i + 0.5))});
    }
    auto rtree = buildRtree("version", leaves);
    auto idToBoxMap = makeIdToBoxMap(leaves);

    // find values intersecting some area defined by a box
    struct {
        BoundingBox queryBox;
        Ids expectedInWindow;
    } testsInWindow[] = {
        {BoundingBox(Point2(0, 0), Point2(1, 1)), Ids{0, 1}},
        {BoundingBox(Point2(3, 3), Point2(7, 7)), Ids{3, 4, 5, 6, 7}},
        {BoundingBox(Point2(9, 0), Point2(9, 9)), Ids{9}},
    };
    auto intersects = [&](const auto& box, Id id) {
        return spatialRelation(box, idToBoxMap.at(id), geolib3::Intersects);
    };
    for (auto& [queryBox, expectedInWindow] : testsInWindow) {
        auto resultInWindow = Ids{};
        for (Id id : rtree->allIdsInWindow(queryBox, intersects)) {
            resultInWindow.push_back(id);
        }
        EXPECT_TRUE(std::is_permutation(resultInWindow.begin(),
                                        resultInWindow.end(),
                                        expectedInWindow.begin(),
                                        expectedInWindow.end()));
    }

    // find nearest values to a point
    struct {
        Point2 point;
        size_t count;
        Ids expectedNearest;
    } testsNearest[] = {
        {Point2(1, 0), 2, Ids{0, 1}},
        {Point2(5, 5), 5, Ids{5, 4, 6, 3, 7}},
        {Point2(9, 8), 3, Ids{8, 9, 7}},
    };
    auto distance = SquaredDistance{};
    auto distanceToId = [&](const geolib3::Point2& point, Id id) {
        return distance(point, idToBoxMap.at(id));
    };
    for (auto& [point, count, expectedNearest] : testsNearest) {
        auto resultNearest =
            nearestIds(*rtree, point, distance, distanceToId, count);
        EXPECT_EQ(resultNearest, expectedNearest);
    }
}

TEST(rtree, in_window)
{
    for (int size = 2; size < 50; ++size) {
        for (int branching = 2; branching < 10; ++branching) {
            for (float probability : {0.1, 0.9}) {
                auto leaves = grid(size, size, probability);
                auto rtree = buildRtree("version", leaves, branching);
                auto idToBoxMap = makeIdToBoxMap(leaves);
                auto intersects = [&](const auto& box, Id id) {
                    return spatialRelation(
                        box, idToBoxMap.at(id), geolib3::Intersects);
                };

                int span = size * 10;
                verifyNumberFound(
                    *rtree,
                    BoundingBox(
                        Point2(-span / 2, -span / 2), span / 4, span / 4),
                    intersects,
                    0);

                verifyNumberFound(
                    *rtree,
                    BoundingBox(
                        Point2(span / 2, span / 2), span * 1.1, span * 1.1),
                    intersects,
                    leaves.size());

                verifyNumberFound(*rtree,
                                  BoundingBox(Point2(0, 0), span * 3, span * 3),
                                  intersects,
                                  leaves.size());

                for (const auto& leaf : leaves) {
                    auto rng = rtree->allIdsInWindow(leaf.bbox, intersects);
                    EXPECT_EQ(std::distance(rng.begin(), rng.end()), 1);
                    EXPECT_EQ(*rng.begin(), leaf.id);
                }
            }
        }
    }
}

TEST(rtree, nearest)
{
    int size = 10;
    float probability = 0.5;
    auto leaves = grid(size, size, probability);
    auto rtree = buildRtree("version", leaves);
    auto idToBoxMap = makeIdToBoxMap(leaves);
    for (int i = 0; i != std::min(size, 5); ++i) {
        auto point = Point2(rand() % size, rand() % size);
        verifyClosest(idToBoxMap, *rtree, point, std::max(rand(), 1));
    }
}

TEST(rtree, empty)
{
    auto rtree = buildRtree("version", LeafNodes{});
    auto largeBbox = BoundingBox(Point2(0, 0),
                                 std::numeric_limits<double>::max(),
                                 std::numeric_limits<double>::max());
    auto intersects = [](auto&&...) { return true; };
    verifyNumberFound(*rtree, largeBbox, intersects, 0);
}

TEST(rtree, multiple_leaves_in_one_cell)
{
    // create the rtree
    auto idToPointMap = IdToPointMap{
        {Id(1), Point2(44.0149, 56.3126)},
        {Id(2), Point2(44.0147, 56.3127)},
        {Id(3), Point2(44.0146, 56.3128)},
        {Id(4), Point2(44.0145, 56.3130)},
        {Id(5), Point2(44.0144, 56.3131)},
        {Id(6), Point2(44.0141, 56.3134)},
        {Id(7), Point2(44.0140, 56.3136)},
        {Id(8), Point2(44.0142, 56.3133)},
        {Id(9), Point2(37.5994, 55.6293)},
        {Id(10), Point2(37.5996, 55.6295)},
        {Id(11), Point2(37.5997, 55.6296)},
        {Id(12), Point2(37.5999, 55.6294)},
    };
    auto leaves = LeafNodes{};
    for (auto& [id, point] : idToPointMap) {
        leaves.push_back({id, point.boundingBox()});
    }
    auto rtree = buildRtree("version", leaves);

    // self intersection
    auto intersects = [&](const auto& box, Id id) {
        return spatialRelation(box, idToPointMap.at(id), geolib3::Intersects);
    };
    for (auto [id, point] : idToPointMap) {
        auto rng = rtree->allIdsInWindow(point.boundingBox(), intersects);
        EXPECT_EQ(std::distance(rng.begin(), rng.end()), 1u);
        EXPECT_EQ(*rng.begin(), id);
    }

    // self nearest
    auto distance = FastGeoDistance{};
    auto distanceToId = [&](const geolib3::Point2& point, Id id) {
        return distance(point, idToPointMap.at(id));
    };
    for (auto [id, point] : idToPointMap) {
        auto ids = nearestIds(*rtree, point, distance, distanceToId, 1);
        EXPECT_EQ(ids.size(), 1u);
        EXPECT_EQ(ids.front(), id);
    }
}

TEST(rtree, all_ids)
{
    auto idToPointMap = IdToPointMap{
        {Id(1), Point2(44.0149, 56.3126)},
        {Id(2), Point2(44.0147, 56.3127)},
        {Id(3), Point2(44.0146, 56.3128)},
        {Id(4), Point2(44.0145, 56.3130)},
        {Id(5), Point2(44.0144, 56.3131)},
        {Id(6), Point2(44.0141, 56.3134)},
        {Id(7), Point2(44.0140, 56.3136)},
        {Id(8), Point2(44.0142, 56.3133)},
        {Id(9), Point2(37.5994, 55.6293)},
        {Id(10), Point2(37.5996, 55.6295)},
        {Id(11), Point2(37.5997, 55.6296)},
        {Id(12), Point2(37.5999, 55.6294)},
    };
    auto leaves = LeafNodes{};
    Ids expectedIds;
    for (auto& [id, point] : idToPointMap) {
        leaves.push_back({id, point.boundingBox()});
        expectedIds.push_back(id);
    }
    auto rtree = buildRtree("version", leaves);
    auto idsRange  = rtree->allIds();
    EXPECT_THAT(Ids(idsRange.begin(), idsRange.end()),
        ::testing::UnorderedPointwise(::testing::Eq(), expectedIds));

}

}  // namespace maps::mrc::fb_rtree::tests
