#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_road_marking/impl/yt_worker.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_road_marking/impl/yt_spec.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_road_marking/impl/strings.h>

#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/recognition_gateway.h>

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/id.h>

#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/schema.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/serialization.h>

#include <maps/wikimap/mapspro/services/mrc/libs/signdetect/include/detected_sign.h>
#include <maps/wikimap/mapspro/services/mrc/libs/roadmark_detector/include/roadmarkdetector.h>
#include <maps/wikimap/mapspro/services/mrc/libs/birdview/include/vanishing_point.h>
#include <maps/wikimap/mapspro/services/mrc/libs/birdview/include/birdview.h>
#include <maps/wikimap/mapspro/services/mrc/libs/birdview/include/utility.h>

#include <maps/libs/log8/include/log8.h>

namespace maps::mrc::eye {

namespace {

std::vector<signdetect::DetectedSign> findRoadMarking(
    const roadmarkdetector::RoadMarkDetector& roadMarkDetector,
    const cv::Mat& image,
    const std::optional<cv::Point>& vanishingPoint,
    int bottomEdge)
{
    if (!vanishingPoint) {
        return {};
    }

    static constexpr double MIN_CONFIDENCE = 0.7;
    static constexpr int MIN_SIDE_PX = 15;
    static const cv::Size BIRDVIEW_SIZE {800, 800};

    std::vector<signdetect::DetectedSign> result;
    auto homography = birdview::findBirdviewTransformationMatrix(
        *vanishingPoint, bottomEdge, BIRDVIEW_SIZE);
    cv::Mat birdview;
    cv::warpPerspective(image, birdview, homography, BIRDVIEW_SIZE);
    auto inv = homography.inv();
    auto imageBoundary = cv::Rect(0, 0, image.cols, image.rows);
    for (const auto& roadMark: roadMarkDetector.detect(birdview)) {
        if (roadMark.confidence > MIN_CONFIDENCE) {
            cv::Rect rect =
                birdview::findBoxOfPerspectiveTransform(roadMark.box, inv) &
                imageBoundary;
            if (rect.width >= MIN_SIDE_PX && rect.height >= MIN_SIDE_PX &&
                rect.y > vanishingPoint->y) {
                result.push_back({rect, roadMark.type, roadMark.confidence});
            }
        }
    }
    return result;
}

db::eye::Recognition makeRecognition(
    const db::eye::Frame& frame,
    const std::vector<signdetect::DetectedSign>& detections)
{
    static const int16_t version = 0;

    db::eye::DetectedRoadMarkings roadMarkings;
    for (const signdetect::DetectedSign& detection : detections) {
        roadMarkings.push_back({
            common::revertByImageOrientation(
                detection.box, frame.originalSize(), frame.orientation()
            ),
            detection.sign,
            detection.confidence
        });
    }

    return {
        frame.id(),
        frame.orientation(),
        db::eye::RecognitionType::DetectRoadMarking,
        db::eye::RecognitionSource::Model,
        version,
        roadMarkings
    };
}

} // namespace

DetectRoadMarkingWorker::DetectRoadMarkingWorker(const FrameLoader& frameLoader)
    : frameLoader_(frameLoader)
{}

void DetectRoadMarkingWorker::Do(
    NYT::TTableReader<NYT::TNode>* reader,
    NYT::TTableWriter<NYT::TNode>* writer)
{
    roadmarkdetector::RoadMarkDetector roadMarkDetector;

    for (; reader->IsValid(); reader->Next()) {
        const NYT::TNode& node = reader->GetRow();

        const db::eye::Frame frame
            = yt::deserialize<db::eye::Frame>(node[FRAME_COLUMN]);
        INFO() << "Process frame: " << frame.id();

        const cv::Mat image = frameLoader_.load(frame);
        const auto vanishingPoint = birdview::findVanishingPoint(image);
        if (vanishingPoint.has_value()) {
            INFO() << "Vanishing point: " << vanishingPoint.value();
        } else {
            INFO() << "Failed to find vanishing point";
        }

        const db::eye::Recognition panel
            = yt::deserialize<db::eye::Recognition>(node[PANEL_COLUMN]);
        INFO() << "Recognition id: " << panel.id();
        const common::ImageBox panelBox
            = common::transformByImageOrientation(
                panel.value<db::eye::DetectedPanel>().restBox,
                frame.originalSize(),
                panel.orientation()
            );
        INFO() << "Panel box: " << panel.value<db::eye::DetectedPanel>().restBox;
        INFO() << "Oriented panel box: " << panelBox;
        const int panelRow = panelBox.maxY();
        INFO() << "Panel row: " << panelRow;

        std::vector<signdetect::DetectedSign> roadMarkings{};

        if (vanishingPoint.has_value() && panelRow > vanishingPoint->y) {
            roadMarkings = findRoadMarking(
                roadMarkDetector, image,
                vanishingPoint, panelRow
            );

            INFO() << "Detected " << roadMarkings.size() << " road markings";
        } else {
            INFO() << "Detection has not been executed for frame " << frame.id();
        }

        writer->AddRow(yt::serialize(makeRecognition(frame, roadMarkings)));
    }
}

REGISTER_MAPPER(DetectRoadMarkingWorker);

} // namespace maps::mrc::eye
