#include <maps/wikimap/mapspro/services/tasks_social/src/assessment_sampler/lib/limits.h>
#include <maps/wikimap/mapspro/services/tasks_social/src/assessment_sampler/lib/limits_impl.h>

#include "helpers.h"

#include <maps/wikimap/mapspro/libs/assessment/tests/helpers/unit_creator.h>

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


namespace maps::wiki::assessment::sampler::tests {

using assessment::tests::UnitCreator;

std::set<Mode> getKeys(auto map)
{
    std::set<Mode> result;
    for (const auto& [key, _]: map) {
        result.insert(key);
    }
    return result;
}


Y_UNIT_TEST_SUITE(limits_tests)
{
    Y_UNIT_TEST(all_mode_to_smth_maps_should_have_same_keys)
    {
        const auto modeToUnitsRatioKeys      = getKeys(impl::MODE_TO_UNITS_RATIO);
        const auto modeToSampleSizeLimitKeys = getKeys(impl::MODE_TO_MAX_UNITS);
        const auto modeToTasksPerUnitKeys    = getKeys(impl::MODE_TO_TASKS_PER_UNIT);

        UNIT_ASSERT_EQUAL(modeToUnitsRatioKeys, modeToSampleSizeLimitKeys);
        UNIT_ASSERT_EQUAL(modeToUnitsRatioKeys, modeToTasksPerUnitKeys);
    }

    Y_UNIT_TEST(should_decrease_groups_in_ratio_times)
    {
        const double RATIO = 0.5;
        const size_t MAX_UNITS = 100;

        const auto result = cropGroups(
            GroupNameToUnits{
                {"group 1", {*UnitCreator(11), *UnitCreator(12), *UnitCreator(13), *UnitCreator(14)}},
                {"group 2", {*UnitCreator(21), *UnitCreator(22), *UnitCreator(23)}},
                {"group 3", {*UnitCreator(31), *UnitCreator(32)}},
                {"group 4", {*UnitCreator(41)}},
            },
            RATIO,
            MAX_UNITS
        );

        UNIT_ASSERT_EQUAL(result.size(), 4);
        UNIT_ASSERT_EQUAL(result.at("group 1").size(), std::floor(4 * RATIO));
        UNIT_ASSERT_EQUAL(result.at("group 2").size(), std::floor(3 * RATIO));
        UNIT_ASSERT_EQUAL(result.at("group 3").size(), std::floor(2 * RATIO));
        UNIT_ASSERT_EQUAL(result.at("group 4").size(), 1); // min 1 unit in a group is allowed.
    }

    Y_UNIT_TEST(should_crop_sample_down_to_max_sample_size)
    {
        const double RATIO = 1;
        const size_t MAX_UNITS = 5;

        const auto result = cropGroups(
            GroupNameToUnits{
                {"group 1", {*UnitCreator(11), *UnitCreator(12), *UnitCreator(13), *UnitCreator(14), *UnitCreator(15)}},
                {"group 2", {*UnitCreator(21), *UnitCreator(22), *UnitCreator(23), *UnitCreator(24)}},
            },
            RATIO,
            MAX_UNITS
        );

        UNIT_ASSERT_EQUAL(result.size(), 2);
        UNIT_ASSERT_EQUAL(
            result.at("group 1").size() + result.at("group 2").size(),
            MAX_UNITS
        );
    }

    Y_UNIT_TEST(should_not_crop_groups_to_length_less_than_one)
    {
        const double RATIO = 1;
        const size_t MAX_UNITS = 2;

        const auto result = cropGroups(
            GroupNameToUnits{
                {"group 1", {*UnitCreator(11)}},
                {"group 2", {*UnitCreator(21), *UnitCreator(22)}},
                {"group 3", {*UnitCreator(31), *UnitCreator(32), *UnitCreator(33)}},
            },
            RATIO,
            MAX_UNITS
        );

        UNIT_ASSERT_EQUAL(result.size(), 3);
        UNIT_ASSERT_EQUAL(result.at("group 1").size(), 1);
        UNIT_ASSERT_EQUAL(result.at("group 2").size(), 1);
        UNIT_ASSERT_EQUAL(result.at("group 3").size(), 1);
    }
}

} // namespace maps::wiki::assessment::sampler::tests
