#include <maps/wikimap/mapspro/services/tasks_social/src/assessment_sampler/lib/make_expert_sample.h>
#include <maps/wikimap/mapspro/services/tasks_social/src/assessment_sampler/lib/make_expert_sample_impl.h>

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

#include <maps/wikimap/mapspro/libs/unittest/include/yandex/maps/wiki/unittest/arcadia.h>
#include <maps/wikimap/mapspro/services/tasks_social/src/assessment_sampler/tests/helpers.h>

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

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

namespace feedback = maps::wiki::social::feedback;

namespace {

TIds createUnits(assessment::Gateway& assessmentGw, size_t count)
{
    TIds unitIds;
    for (const auto& unit: generateUnits(count)) {
        unitIds.insert(assessmentGw.getOrCreateUnit(unit.entity, unit.action));
    }
    return unitIds;
}

void createBasicSample(
    assessment::Gateway& assessmentGw,
    const std::string& name,
    size_t unitsNumber,
    size_t gradedUnitsNumber)
{
    ASSERT(unitsNumber >= gradedUnitsNumber);

    const size_t TASKS_PER_UNIT = 2;
    const TUid UID = 1;

    const auto unitIds = createUnits(assessmentGw, unitsNumber);
    const TId sampleId = assessmentGw.createSample(Entity::Domain::Feedback, Qualification::Basic, name, unitIds, TASKS_PER_UNIT);

    auto console = assessmentGw.console(UID);
    while (gradedUnitsNumber-- > 0) {
        const auto unit = console.acquireUnit(sampleId, SkipAcquired::No);
        UNIT_ASSERT(unit);
        console.gradeUnit(unit->id, Grade::Value::Correct, "comment", Qualification::Basic);
    }
}

std::optional<SampleInfo> getSample(pqxx::transaction_base& txn, const std::string& name)
{
    const auto rows = txn.exec(
        "SELECT sample_id FROM assessment.sample WHERE name = '" + name + "'"
    );

    if (rows.empty()) {
        return std::nullopt;
    }

    return assessment::Gateway(txn).findSample(rows[0][0].as<TId>());
}

std::vector<TId> setIntersection(const std::vector<TId>& lhs, const std::vector<TId>& rhs)
{
    std::vector<TId> result;
    std::set_intersection(lhs.cbegin(), lhs.cend(), rhs.cbegin(), rhs.cend(), std::back_inserter(result));
    return result;
}

} // namespace

Y_UNIT_TEST_SUITE(expert_sample_small_tests) {
    Y_UNIT_TEST(print_seed)
    {
        std::cerr << "Seed: " << SEED << "\n";
    }

    Y_UNIT_TEST(should_pick_expert_units)
    {
        const double UNITS_RATIO = 0.5;
        const double INCONSISTENT_UNITS_RATIO = 2.0 / 3.0;
        std::mt19937 rnd{SEED};

        const std::vector<TId> inconsistent = {11, 12, 13, 14};
        const std::vector<TId> consistent   = {21, 22, 23, 24};
        const std::vector<TId> notGraded    = {31, 32, 33, 34};
        const auto totalUnitsCount = inconsistent.size() + consistent.size() + notGraded.size();

        auto result = impl::pickExpertUnits(
            rnd,
            {inconsistent, consistent, notGraded},
            UNITS_RATIO,
            INCONSISTENT_UNITS_RATIO
        );
        UNIT_ASSERT_EQUAL(result.size(), UNITS_RATIO * totalUnitsCount);

        std::sort(result.begin(), result.end());
        UNIT_ASSERT_EQUAL(
            setIntersection(result, inconsistent).size(),
            totalUnitsCount * UNITS_RATIO * INCONSISTENT_UNITS_RATIO
        );
        UNIT_ASSERT_EQUAL(
            setIntersection(result, consistent).size(),
            totalUnitsCount * UNITS_RATIO * (1 - INCONSISTENT_UNITS_RATIO)
        );
        UNIT_ASSERT(setIntersection(result, notGraded).empty());
    }
}

Y_UNIT_TEST_SUITE_F(expert_sample_tests, unittest::ArcadiaDbFixture) {
    Y_UNIT_TEST(should_make_expert_sample) {
        auto txn = pool().masterWriteableTransaction();
        const auto [timepointMin, timepointMax] = txnNowSpan(*txn);
        std::mt19937 rnd{SEED};

        assessment::Gateway assessmentGw{*txn};

        const size_t UNITS_NUMBER = 20;
        const size_t GRADED_UNITS_NUMBER = 15;
        createBasicSample(assessmentGw, "basic-sample", UNITS_NUMBER, GRADED_UNITS_NUMBER);

        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Staff, HypothesesMode::No, "expert-sample");

        const auto sampleInfo = getSample(*txn, "expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->entityDomain, Entity::Domain::Feedback);
        UNIT_ASSERT_EQUAL(sampleInfo->qualification, Qualification::Expert);
        UNIT_ASSERT_EQUAL(sampleInfo->gradedCount, 0);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, UNITS_NUMBER * EXPERT_UNITS_RATIO);
    }

    Y_UNIT_TEST(should_make_expert_sample_from_appropriate_basic_sample) {
        auto txn = pool().masterWriteableTransaction();
        const auto [timepointMin, timepointMax] = txnNowSpan(*txn);
        std::mt19937 rnd{SEED};

        assessment::Gateway assessmentGw{*txn};

        createBasicSample(assessmentGw, "basic-sample",                      30, 30);
        createBasicSample(assessmentGw, "hypotheses-basic-sample",           25, 25);
        createBasicSample(assessmentGw, "piecework-basic-sample",            20, 20);
        createBasicSample(assessmentGw, "piecework-hypotheses-basic-sample", 15, 15);
        createBasicSample(assessmentGw, "outsource-basic-sample",            10, 10);
        createBasicSample(assessmentGw, "outsource-hypotheses-basic-sample",  5,  5);

        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Staff,     HypothesesMode::No,  "expert-sample");
        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Staff,     HypothesesMode::Yes, "hypotheses-expert-sample");
        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Piecework, HypothesesMode::No,  "piecework-expert-sample");
        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Piecework, HypothesesMode::Yes, "piecework-hypotheses-expert-sample");
        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Outsource, HypothesesMode::No,  "outsource-expert-sample");
        makeExpertSample(*txn, Entity::Domain::Feedback, timepointMin, timepointMax, rnd, StaffType::Outsource, HypothesesMode::Yes, "outsource-hypotheses-expert-sample");

        auto sampleInfo = getSample(*txn, "expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 30 * EXPERT_UNITS_RATIO);

        sampleInfo = getSample(*txn, "hypotheses-expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 25 * EXPERT_UNITS_RATIO);

        sampleInfo = getSample(*txn, "piecework-expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 20 * EXPERT_UNITS_RATIO);

        sampleInfo = getSample(*txn, "piecework-hypotheses-expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 15 * EXPERT_UNITS_RATIO);

        sampleInfo = getSample(*txn, "outsource-expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 10 * EXPERT_UNITS_RATIO);

        sampleInfo = getSample(*txn, "outsource-hypotheses-expert-sample");
        UNIT_ASSERT(sampleInfo);
        UNIT_ASSERT_EQUAL(sampleInfo->totalCount, 5 * EXPERT_UNITS_RATIO);
    }
}

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