#include <medusa/bits/spatial_search/KDTree_fwd.hpp>

#include "gtest/gtest.h"

#include <medusa/bits/types/Vec.hpp>

namespace mm {

TEST(KDTree, DISABLED_UsageExample) {
    /// [KDTree usage example]
    Range<Vec3d> points = {{5, 4, 0}, {4, 2, 1}, {7, 6, 9}, {2, 2, 1}, {8, 0, 5},
                           {5, 7, 0}, {3, 3, 8}, {9, 7, 3}, {2, 2, 6}, {2, 0, 6}};
    KDTree<Vec3d> tree(points);  // Build kd-tree
    Range<double> distances2;  // squares of distances
    Range<int> closest;
    // Get 3 points closest to {4.1, 3, 0.5}
    std::tie(closest, distances2) = tree.query({4.1, 3, 0.5}, 3);
    // closest = [1, 0, 3];
    Vec3d p = tree.get(closest[0]);  // get coordinates of the closest point
    Range<Vec3d> closest_points = tree.get(closest);
    // distances2 = [1.26, 2.06, 5.66];
    std::cout << tree << std::endl;
    /// [KDTree usage example]
    p = 0.0;  // to not be unused.
    closest_points.clear();
}

TEST(KDTree, 3DQuery) {
    Range<Vec3d> points = {
            {5, 4, 0}, {4, 2, 1}, {7, 6, 9}, {2, 2, 1}, {8, 0, 5}, {5, 7, 0},
            {3, 3, 8}, {9, 7, 3}, {2, 2, 6}, {2, 0, 6}};
    KDTree<Vec3d> kd3(points);
    EXPECT_EQ(points.size(), kd3.size());
    EXPECT_EQ(std::vector<int>({1, 0}), kd3.query({4.1, 3, 0.5}, 2).first);
    EXPECT_EQ(std::vector<int>({3, 1, 8}), kd3.query({2, 2, 3}, 3).first);
}

TEST(KDTree, Get) {
    /// [KDTree get]
    Range<Vec3d> points = {{5, 4, 0}, {4, 2, 1}, {7, 6, 9}, {2, 2, 1}, {8, 0, 5}};
    KDTree<Vec3d> kd3(points);
    ASSERT_EQ(points.size(), kd3.size());
    for (int i = 0; i < kd3.size(); ++i) {
        EXPECT_EQ(points[i], kd3.get(i));
    }

    EXPECT_EQ(points, kd3.get({0, 1, 2, 3, 4}));  // vectorized get

    Range<int> idxs = kd3.query({4.1, 3, 0.5}, 1).first;
    ASSERT_EQ(idxs.size(), 1);
    EXPECT_EQ(Vec3d({4, 2, 1}), kd3.get(idxs[0]));
    /// [KDTree get]
}

TEST(KDTree, 2DQuery) {
    // Test 2D with 50 points
    Range<Vec2d> points = {
            {-0.6441209666616923, 0.5850353695296744}, {-0.6201378980490286, 0.7811830957022021},
            {-0.6784344365993658, 0.9839667035596042}, {0.044496636767509035, 0.4377470158269272},
            {-0.5274564029904492, 0.5696927123044102}, {-0.6259567688687637, -0.7818614711722951},
            {0.9162830614939184, 0.6489213976078669}, {0.7591192942947527, -0.3433314817836943},
            {-0.23782320760874454, 0.1898945196158186}, {-0.5441698103664361, 0.006385132320767317},
            {-0.3654415676489693, 0.06746853328508307}, {0.931189791692042, 0.9181001537912621},
            {-0.3427312891982608, -0.6826010311723985}, {0.891723271821818, -0.8508989951879877},
            {0.11394616056891804, 0.84630479540056}, {-0.5223541537898675, -0.026438581866889965},
            {0.425710227280224, 0.232320814224092}, {-0.12153165386201059, -0.09554948988436873},
            {-0.5360668060424465, 0.8346421613694428}, {-0.48815760323401425, -0.47632741760208863},
            {-0.7355755648085145, 0.021559400675411178}, {0.16684622124842763, -0.5180297967184262},
            {0.2750341401977108, 0.14027191172883025}, {0.6550291085853586, -0.0906298270828052},
            {0.11672354315833089, 0.6864647718719339}, {0.4319012792644372, -0.5056121884571734},
            {-0.3111112525746833, 0.4841999490074793}, {0.9537933710310402, -0.4121810950684375},
            {-0.13430435085483827, -0.7063126623794711}, {0.8509710074195942, 0.4342793339962312},
            {0.8387012530238107, -0.8797447958759539}, {-0.788610725923649, -0.28402465290256296},
            {0.3781644510827291, -0.19185822357265714}, {0.3675280148118516, -0.19959444568922202},
            {0.5441631640348303, 0.9923042351466602}, {0.4379307570126687, -0.7976305308639227},
            {0.6192156716466921, 0.5095818625700927}, {0.4538538344127916, 0.24132983301916466},
            {-0.920957301103486, 0.7946521172232996}, {-0.4620339775822764, -0.6725241136453264},
            {0.6164340529959673, -0.370457261186975}, {0.6353232887100764, -0.4527336246027056},
            {0.5129962829311199, 0.2576340035951956}, {-0.3355508932427227, 0.20820573070479842},
            {-0.03557537269552791, 0.801250103966441}, {-0.6249260602423057, 0.01901038540320088},
            {0.9052024446464844, -0.37370939241714884}, {-0.08417808324593734, -0.4918951358215899},
            {-0.7408521674604813, -0.38351394950328266}, {0.8804020099004968, 0.2246660525735673}};
    int rs;
    KDTree<Vec2d> kd2({{-1, -1}, {-1, 0}, {-4, 4}});
    kd2.reset(points);

    EXPECT_EQ(points.size(), kd2.size());
    EXPECT_EQ(2, kd2.query(points[33], 0.05).first.size());
    EXPECT_EQ(2, kd2.query(points[33], 0.05).first.size());
    auto result = kd2.query(points[33], 3);
    EXPECT_EQ(std::vector<int>({33, 32, 40}), result.first);
    rs = result.first.size();
    for (int i = 0; i < rs; i++) {
        EXPECT_DOUBLE_EQ(
                (points[33] - points[result.first[i]]).squaredNorm(), result.second[i]);
        EXPECT_LE(result.second[i], result.second[rs-1]);
    }

    EXPECT_EQ(2, kd2.query(points[17], 0.09).first.size());
    result = kd2.query(points[17], 3);
    EXPECT_EQ(std::vector<int>({17, 10, 8}), result.first);
    rs = result.first.size();
    for (int i = 0; i < rs; i++) {
        EXPECT_DOUBLE_EQ(
                (points[17] - points[result.first[i]]).squaredNorm(), result.second[i]);
        EXPECT_LE(result.second[i], result.second[rs-1]);
    }

    EXPECT_EQ(4, kd2.query({0., 0}, 0.14).first.size());
    result = kd2.query({0., 0}, 5);
    EXPECT_EQ(std::vector<int>({17, 8, 22, 10, 43}), result.first);
    rs = result.first.size();
    for (int i = 0; i < rs; i++) {
        EXPECT_DOUBLE_EQ(
                (Vec2d({0, 0}) - points[result.first[i]]).squaredNorm(), result.second[i]);
        EXPECT_LE(result.second[i], result.second[rs-1]);
    }

    EXPECT_EQ(3, kd2.query({0.7, 0.7}, 0.101).first.size());
    result = kd2.query({0.7, 0.7}, 5);
    EXPECT_EQ(std::vector<int>({36, 6, 29, 11, 34}), result.first);
    rs = result.first.size();
    for (int i = 0; i < rs; i++) {
        EXPECT_DOUBLE_EQ(
                (Vec2d({0.7, 0.7}) - points[result.first[i]]).squaredNorm(), result.second[i]);
        EXPECT_LE(result.second[i], result.second[rs-1]);
    }
}

TEST(KDTree, DeathTests) {
    Range<Vec2d> points({{1.3, 2.4}});
    KDTree<Vec2d> kd(points);
    EXPECT_DEATH(kd.query(0., 3), "There were not enough points in the tree, you requested 3 "
                                  "points, the tree only contains 1 points.");
}

}  // namespace mm