#include <maps/wikimap/mapspro/services/mrc/eye/lib/detection/tests/fixture.h>

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/id.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detection/include/greedy_clusterizer.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/location/include/rotation.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/unit_test/include/frame.h>

#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/recognition_gateway.h>

#include <library/cpp/testing/gmock_in_unittest/gmock.h>
#include <library/cpp/testing/unittest/registar.h>

namespace maps::mrc::eye::tests {

struct GreedyClusteringFixture: public Fixture {
    GreedyClusteringFixture();

    db::eye::Devices devices;
    db::eye::Frames frames;
    db::eye::FrameLocations locations;
    db::eye::FramePrivacies privacies;
    db::eye::DetectionGroups groups;
    db::eye::Detections detections;

    std::unique_ptr<DetectionStore> store;
};

GreedyClusteringFixture::GreedyClusteringFixture()
{
    devices = insertx<db::eye::DeviceGateway>(
        db::eye::Devices {
            {db::eye::MrcDeviceAttrs{""}}
        }
    );

    frames = insertx<db::eye::FrameGateway>(
        db::eye::Frames {
            {devices[0].id(), identical, makeUrlContext(1, "1"), {1200, 800}, time()},
            {devices[0].id(), identical, makeUrlContext(2, "2"), {1200, 800}, time()},
            {devices[0].id(), identical, makeUrlContext(3, "3"), {1200, 800}, time()},
            {devices[0].id(), identical, makeUrlContext(4, "4"), {1200, 800}, time()},
            {devices[0].id(), identical, makeUrlContext(5, "5"), {1200, 800}, time()},
            {devices[0].id(), identical, makeUrlContext(5, "5"), {1200, 800}, time()},
        }
    );

    locations = insertx<db::eye::FrameLocationGateway>(
        db::eye::FrameLocations {
            {frames[0].id(), geolib3::Point2{0, 0}, toRotation(geolib3::Heading(90), identical)},
            {frames[1].id(), geolib3::Point2{1, 0}, toRotation(geolib3::Heading(90), identical)},
            {frames[2].id(), geolib3::Point2{0, 2}, toRotation(geolib3::Heading(90), identical)},
            {frames[3].id(), geolib3::Point2{1, 1}, toRotation(geolib3::Heading(90), identical)},
            {frames[4].id(), geolib3::Point2{1, 1}, toRotation(geolib3::Heading(90), identical)},
            {frames[5].id(), geolib3::Point2{1, 1}, toRotation(geolib3::Heading(90), identical)},
        }
    );

    for (const auto& frame : frames) {
        privacies.emplace_back(frame.id(), db::FeaturePrivacy::Public);
    }
    privacies = insertx<db::eye::FramePrivacyGateway>(std::move(privacies));

    groups = insertx<db::eye::DetectionGroupGateway>(
        db::eye::DetectionGroups {
            {frames[0].id(), db::eye::DetectionType::HouseNumber},
            {frames[1].id(), db::eye::DetectionType::HouseNumber},
            {frames[2].id(), db::eye::DetectionType::HouseNumber},
            {frames[3].id(), db::eye::DetectionType::HouseNumber},
            {frames[4].id(), db::eye::DetectionType::HouseNumber},
        }
    );

    detections = insertx<db::eye::DetectionGateway>(
        db::eye::Detections {
            {groups[0].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
            {groups[1].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
            {groups[2].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
            {groups[3].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
            {groups[4].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
            {groups[1].id(), db::eye::DetectedHouseNumber{{100, 100, 200, 200}, 1.0, "12"}},
        }
    );

    store = std::make_unique<DetectionStore>(
        byId(groups),
        byId(detections),
        byId(frames),
        byId(locations),
        byId(privacies),
        byId(devices)
    );
};

Y_UNIT_TEST_SUITE_F(greedy_clustering, GreedyClusteringFixture)
{

Y_UNIT_TEST(empty)
{
    db::TIdSet detectionIds{};
    MatchedFrameDetections matches{};

    GreedyDetectionClusterizer clusterizer;
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{};

    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(no_edges)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id()
    };
    MatchedFrameDetections matches{};

    GreedyDetectionClusterizer clusterizer;
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {detections[0].id()},
        {detections[1].id()},
        {detections[2].id()}
    };

    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(one_cluster)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id(),
        detections[3].id(),
        detections[4].id()
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[1].id(), detections[1].id()},
            {frames[2].id(), detections[2].id()},
            1.0
        },
        {
            {frames[2].id(), detections[2].id()},
            {frames[3].id(), detections[3].id()},
            1.0
        },
        {
            {frames[3].id(), detections[3].id()},
            {frames[4].id(), detections[4].id()},
            1.0
        },
    };

    GreedyDetectionClusterizer clusterizer;
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {
            detections[0].id(),
            detections[1].id(),
            detections[2].id(),
            detections[3].id(),
            detections[4].id()
        }
    };

    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(simple_merge)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id(),
        detections[3].id(),
        detections[4].id()
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[2].id(), detections[2].id()},
            {frames[0].id(), detections[0].id()},
            1.0
        },
        {
            {frames[3].id(), detections[3].id()},
            {frames[4].id(), detections[4].id()},
            1.0
        },
    };

    GreedyDetectionClusterizer clusterizer;
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {
            detections[0].id(),
            detections[1].id(),
            detections[2].id()
        },
        {
            detections[3].id(),
            detections[4].id()
        }
    };

    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(greedy_criteria)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        // detections[1], detections[5] in one frame
        detections[1].id(),
        detections[5].id(),
        detections[3].id(),
        detections[4].id()
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[1].id(), detections[5].id()},
            {frames[0].id(), detections[0].id()},
            0.01
        },
        {
            {frames[3].id(), detections[3].id()},
            {frames[4].id(), detections[4].id()},
            0.7
        },
    };

    GreedyDetectionClusterizer clusterizer;
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {detections[0].id(), detections[1].id()},
        {detections[5].id()},
        {detections[3].id(), detections[4].id()}
    };
    UNIT_ASSERT_VALUES_EQUAL(result.size(), expected.size());
    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(min_conf)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id(),
        detections[3].id(),
        detections[4].id()
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[2].id(), detections[2].id()},
            {frames[0].id(), detections[0].id()},
            0.01
        },
        {
            {frames[3].id(), detections[3].id()},
            {frames[4].id(), detections[4].id()},
            0.7
        },
    };

    GreedyDetectionClusterizer clusterizer(0.5);
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {
            detections[0].id(),
            detections[1].id()
        },
        {
            detections[2].id()
        },
        {
            detections[3].id(),
            detections[4].id()
        }
    };

    UNIT_ASSERT(std::is_permutation(result.begin(), result.end(), expected.begin()));
}

Y_UNIT_TEST(negative_verdict_is_taken_into_account)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id(),
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[0].id(), detections[0].id()},
            {frames[2].id(), detections[2].id()},
            1.0
        },
        {
            {frames[2].id(), detections[2].id()},
            {frames[1].id(), detections[1].id()},
            1.0,
            MatchedFrameDetection::Verdict::No
        }
    };

    GreedyDetectionClusterizer clusterizer(0.5);
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {
            detections[0].id(),
            detections[1].id()
        },
        {
            detections[2].id()
        }
    };

    EXPECT_THAT(result, ::testing::UnorderedElementsAreArray(expected));
}

Y_UNIT_TEST(positive_verdict_is_taken_into_account)
{
    db::TIdSet detectionIds{
        detections[0].id(),
        detections[1].id(),
        detections[2].id(),
    };
    MatchedFrameDetections matches{
        {
            {frames[0].id(), detections[0].id()},
            {frames[1].id(), detections[1].id()},
            1.0
        },
        {
            {frames[0].id(), detections[0].id()},
            {frames[2].id(), detections[2].id()},
            1.0
        },
        {
            {frames[2].id(), detections[2].id()},
            {frames[1].id(), detections[1].id()},
            0.0,
            MatchedFrameDetection::Verdict::Yes
        }
    };

    GreedyDetectionClusterizer clusterizer(0.5);
    std::vector<db::TIdSet> result = clusterizer.clusterize(*store, detectionIds, matches);

    std::vector<db::TIdSet> expected{
        {
            detections[0].id(),
            detections[1].id(),
            detections[2].id()
        }
    };

    EXPECT_THAT(result, ::testing::UnorderedElementsAreArray(expected));
}


} // Y_UNIT_TEST_SUITE

} // namespace maps::mrc::eye::tests
