#include <mapreduce/yt/interface/client.h>

#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/wikimap/mapspro/services/mrc/libs/signdetect/include/signdetect_faster_rcnn.h>
#include <maps/wikimap/mapspro/services/mrc/libs/traffic_signs/include/yandex/maps/mrc/traffic_signs/signs.h>

#include <opencv2/opencv.hpp>
#include <opencv2/imgcodecs/imgcodecs_c.h>

#include <util/generic/size_literals.h>
#include <util/stream/output.h>
#include <util/system/user.h>

#include <vector>

using namespace NYT;

class TSignDetectMapper
        : public IMapper<TTableReader<TNode>, TTableWriter<TNode>>
{
public:
    maps::mrc::signdetect::FasterRCNNDetector detector;

    void Do(TReader* reader, TWriter* writer) override
    {
        for (; reader->IsValid(); reader->Next()) {
            const auto& row = reader->GetRow();

            TNode outRow;
            outRow["feature_id"] = row["feature_id"];

            int nSize = row["image"].AsString().size();
            cv::Mat image = cv::imdecode(cv::Mat(1, nSize, CV_8UC1, (void*)row["image"].AsString().data()), CV_LOAD_IMAGE_COLOR);
            auto detections = detector.detect(image);

            for (auto detection = detections.begin(); detection != detections.end(); ++detection) {
                outRow["x"] = detection->box.x;
                outRow["y"] = detection->box.y;
                outRow["width"] = detection->box.width;
                outRow["height"] = detection->box.height;
                outRow["confidence"] = detection->confidence;
                outRow["sign_type"] = TString(maps::mrc::traffic_signs::toString(detection->sign));
                outRow["sign_category"] = TString(maps::mrc::traffic_signs::toString(
                        maps::mrc::traffic_signs::toTrafficSignCategory(detection->sign)
                ));
                writer->AddRow(outRow);
            }
        }
    }
};
REGISTER_MAPPER(TSignDetectMapper);

int main(int argc, char** argv) {
    Initialize(argc, argv);

    maps::cmdline::Parser parser("Detects signs on images from input-table, "
                                 " exports detections as bounding boxes to output-table");
    auto inputTable = parser.string("input-table")
        .help("Input table with road images")
        .required();

    auto outputTable = parser.string("output-table")
        .help("Output table with bounding boxes, sign categories and types")
        .required();

    parser.parse(argc, argv);

    auto client = CreateClient("hahn");

    client->Map(
            TMapOperationSpec()
                    .MapperSpec(TUserJobSpec().MemoryLimit(8_GB))
                    .DataSizePerJob(16_MB)
                    .AddInput<TNode>(TString(inputTable))
                    .AddOutput<TNode>(TString(outputTable)),
            new TSignDetectMapper);

    return 0;
}
