#include "inspect_quality.h"
#include "tools.h"

#include <maps/libs/log8/include/log8.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/opencv.h>

#include <boost/optional.hpp>
#include <boost/range/algorithm_ext/erase.hpp>

#include <algorithm>

namespace maps {
namespace mrc {
namespace img_qa {
namespace {

boost::optional<db::ugc::TolokaStatus>
getStatus(const toloka::ImageQualityC12nTask& task)
{
    switch (task.status()) {
    case toloka::TaskStatus::New:
    case toloka::TaskStatus::InProgress:
        return db::ugc::TolokaStatus::InProgress;
    case toloka::TaskStatus::Finished: {
        auto output = task.outputValues();
        REQUIRE(output, "no output: " << task.id());
        switch (output->answer()) {
        case toloka::ImageQualityC12nAnswer::Ok:
            return db::ugc::TolokaStatus::Accepted;
        case toloka::ImageQualityC12nAnswer::HasDefect:
            return db::ugc::TolokaStatus::Rejected;
        case toloka::ImageQualityC12nAnswer::NotLoaded:
            break;
        }
    } break;
    case toloka::TaskStatus::Failed:
    case toloka::TaskStatus::Free:
    case toloka::TaskStatus::Cancelling:
    case toloka::TaskStatus::Cancelled:
        break;
    }
    return boost::none;
}

/**
 * - send image classification tasks to Toloka queue
 * - save mapping between photos and Toloka tasks
 */
void createImageQualityC12nTasks(Context& ctx,
                                 const db::Features& photos)
{
    static constexpr size_t BATCH_SIZE = 100;
    for (const auto& batch : maps::common::makeBatches(photos, BATCH_SIZE)) {
        auto tasks = ctx.createImageQualityC12nTasks(batch);
        auto photoTasks = zip(batch, tasks);
        ctx.insert(photoTasks);
    }
}

void filterPhotosForToloka(
    db::Features& photos,
    const PhotoIdToTolokaStatusMap& tolokaPhotoToStatusMap,
    const Segments& uncovered)
{
    static constexpr double MIN_ADDITIONAL_CAPTURED_LENGTH_PER_IMAGE
        = 0.1; // meters
    auto uncoveredLength = geoLength(uncovered);
    boost::remove_erase_if(
        photos, [&](const db::Feature& feature) {
            if (tolokaPhotoToStatusMap.count(feature.id())
                || feature.quality() >= GOOD_QUALITY_THRESHOLD
                || feature.quality() <= BAD_QUALITY_THRESHOLD
                || !feature.hasPos() || !feature.hasHeading()) {
                return true;
            }
            if (feature.cameraDeviation() != db::CameraDeviation::Front ||
                feature.privacy() != db::FeaturePrivacy::Public) {
                return true;
            }

            auto cuttings = db::getUncovered(feature, uncovered);
            return geoLength(cuttings)
                       + MIN_ADDITIONAL_CAPTURED_LENGTH_PER_IMAGE
                   >= uncoveredLength;
        });
}

} // anonymous namespace

PhotoIdToTolokaStatusMap loadPhotoToTolokaStatusMap(
    Context& ctx,
    const FeatureBatch& photos)
{
    constexpr size_t BATCH_SIZE = 1000;
    PhotoIdToTolokaStatusMap result;
    for (const auto& batch : maps::common::makeBatches(photos, BATCH_SIZE)) {
        auto photoTasks = ctx.loadFeatureQaTasks({batch.begin(), batch.end()});
        auto tasks = ctx.getImageQualityC12nTasks(photoTasks);
        std::unordered_map<db::TId, db::TId> tolokaToPhotoMap;
        tolokaToPhotoMap.reserve(photoTasks.size());
        for (const auto& photoTask : photoTasks) {
            tolokaToPhotoMap.emplace(photoTask.tolokaId(),
                                     photoTask.featureId());
        }

        for (const auto& task : tasks) {
            auto photoId = tolokaToPhotoMap.at(task.id());
            auto status = getStatus(task);
            if (status) {
                result.emplace(photoId, *status);
            }
        }
    }
    return result;
}

TolokaStatistics inspectQuality(Context& ctx,
                                db::TId assignmentId,
                                db::CameraDeviation cameraDeviation,
                                const Segments& uncovered,
                                bool createTolokaTasks)
{
    auto photos = ctx.loadProcessedPhotos(assignmentId, cameraDeviation);
    auto photoToStatusMap
        = loadPhotoToTolokaStatusMap(ctx, {photos.begin(), photos.end()});
    auto stat
        = makeTolokaStatistics(assignmentId, photos.size(), photoToStatusMap);
    if (stat.tolokaComing > 0) {
        sortByQualityDesc(photos);
        filterPhotosForToloka(photos, photoToStatusMap, uncovered);
        stat.tolokaComing = photos.size();
        if (createTolokaTasks) {
            createImageQualityC12nTasks(ctx, photos);
        }
    }
    return stat;
}

} // img_qa
} // mrc
} // maps
