#include "ground_truth_feature.h"

#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/io.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/feature_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/types.h>

#include <maps/libs/geolib/include/heading.h>
#include <maps/libs/geolib/include/point.h>

#include <mapreduce/yt/interface/client.h>
#include <opencv2/opencv.hpp>


namespace maps::mrc::opensfm_experiment {

namespace {

cv::Mat deserializeIntrinsics(const NYT::TNode& node) {
    cv::Mat intrinsics = cv::Mat::zeros(3, 3, CV_64F);

    for (size_t i = 0; i != node.Size(); ++i) {
        for (size_t j = 0; j != node[0].Size(); ++j) {
            intrinsics.at<double>(i, j) = node[i][j].AsDouble();
        }
    }

    return intrinsics;
}

NYT::TNode serializeIntrinsics(const cv::Mat& intrinsics) {
    auto result = NYT::TNode::CreateList();

    for (int i = 0; i != intrinsics.size().height; ++i) {
        result.Add(NYT::TNode::CreateList());
        for (int j = 0; j != intrinsics.size().width; ++j) {
            result[i].Add(yt::serialize(intrinsics.at<double>(i, j)));
        }
    }
    return result;
}

} // namespace

NYT::TNode serialize(const GroundTruthFeature& value) {
    auto result = yt::serialize(static_cast<db::Feature>(value.feature));
    auto image = NYT::TNode::CreateEntity();

    if (!value.image.empty()) {
        yt::Bytes bytes;
        REQUIRE(cv::imencode(".jpg", value.image, bytes), "Impossible serialize image!");
        image = yt::serialize(bytes);
    }

    auto mask = NYT::TNode::CreateEntity();
    if (!value.mask.empty()) {
        yt::Bytes bytes;
        REQUIRE(cv::imencode(".png", value.mask, bytes), "Impossible serialize mask!");
        mask = yt::serialize(bytes);
    }

    auto truePos = yt::serialize(value.truePos);
    auto trueHeading = yt::serialize(value.trueHeading.value());
    auto intrinsics = serializeIntrinsics(value.intrinsics);

    result
        ("image", image)
        ("mask", mask)
        ("true_pos", truePos)
        ("true_heading", trueHeading)
        ("intrinsics", intrinsics);

    return result;
}

GroundTruthFeature deserialize(const NYT::TNode& node) {
    auto decode = [&node](const std::string& name, int flags) {
        if (node[name].IsNull()) {
            return cv::Mat();
        }

        return cv::imdecode(
             cv::_InputArray(reinterpret_cast<const common::Byte*>(node[name].AsString().data()), node[name].AsString().size()),
             flags
        );
    };

    GroundTruthFeature result(yt::deserialize<db::Feature>(node));

    result.truePos = yt::deserialize<geolib3::Point2>(node["true_pos"]);
    result.trueHeading = geolib3::Heading(node["true_heading"].AsDouble());
    result.image = decode("image", cv::IMREAD_COLOR | cv::IMREAD_IGNORE_ORIENTATION);
    result.mask = decode("mask", cv::IMREAD_GRAYSCALE | cv::IMREAD_IGNORE_ORIENTATION);
    result.intrinsics = deserializeIntrinsics(node["intrinsics"]);

    return result;
}

} // namespace maps::mrc::opensfm_experiment
