#include <maps/wikimap/mapspro/services/mrc/eye/lib/detection/include/greedy_clusterizer.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detection/impl/cluster_store.h>

namespace maps::mrc::eye {

namespace {

MatchedFrameDetections sortByRelevanceWithVerdict(MatchedFrameDetections matches) {
    auto cmpKey = [](const auto& match) -> auto {
        return RelevanceWithVerdict{
            .relevance = match.relevance(),
            .verdict = match.verdict()
        };
    };
    std::sort(matches.begin(), matches.end(),
        [&](const auto& lhs, const auto& rhs) {
            using namespace std::rel_ops;
            return cmpKey(lhs) > cmpKey(rhs);
        }
    );
    return matches;
}


std::map<FrameDetectionId, FrameDetectionIdSet>
makeDetectionToForbiddenDetections(const MatchedFrameDetections& sortedMatches)
{
    std::map<FrameDetectionId, FrameDetectionIdSet> result;
    for (auto it = sortedMatches.rbegin();
        it != sortedMatches.rend() && it->verdict() == MatchedFrameDetection::Verdict::No;
        ++it)
    {
        result[it->id0()].insert(it->id1());
        result[it->id1()].insert(it->id0());
    }
    return result;
}

} // namespace

GreedyDetectionClusterizer::GreedyDetectionClusterizer(float minRelevance)
    : minRelevance_(minRelevance)
{}

std::vector<db::TIdSet> GreedyDetectionClusterizer::clusterize(
    const DetectionStore& /*store*/,
    const db::TIdSet& detectionIds,
    const MatchedFrameDetections& matches) const
{
    ClusterStore clusters;
    const auto sortedMatches = sortByRelevanceWithVerdict(matches);

    for (const auto& [fdId, forbiddenDetections] :
            makeDetectionToForbiddenDetections(sortedMatches))
    {
        clusters.createNewCluster({fdId}, std::move(forbiddenDetections));
    }

    // Greedy merge
    for (const MatchedFrameDetection& match : sortedMatches) {
        if (match.verdict() == MatchedFrameDetection::Verdict::No ||
            !match.verdict().has_value() && match.relevance() < minRelevance_)
        {
            break;
        }
        const FrameDetectionId& fdId0 = match.id0();
        const FrameDetectionId& fdId1 = match.id1();

        ASSERT(fdId0.frameId != fdId1.frameId);

        auto clusterIt0 = clusters.find(fdId0.detectionId);
        auto clusterIt1 = clusters.find(fdId1.detectionId);

        if (clusters.end() == clusterIt0 && clusters.end() == clusterIt1)
        {// no clusters
            clusters.createNewCluster({fdId0, fdId1});
        } else if (clusters.end() == clusterIt0) {
            clusters.tryAddNewElementInCluster(clusterIt1->second, fdId0);
        } else if (clusters.end() == clusterIt1) {
            clusters.tryAddNewElementInCluster(clusterIt0->second, fdId1);
        } else {
            clusters.tryMergeTwoClusters(clusterIt0->second, clusterIt1->second);
        }
    }
    std::vector<db::TIdSet> result = clusters.makeDetectionClusters();
    expandClustersByUnusedDetectionIds(detectionIds, result);
    return result;
}

} // namespace maps::mrc::eye
