#include <maps/wikimap/mapspro/services/mrc/tools/experiment_sign_position_accuracy/lib/include/dataset.h>

#include <maps/libs/log8/include/log8.h>
#include <yandex/maps/proto/offline-mrc/results.sproto.h>
#include <maps/wikimap/mapspro/services/mrc/tools/sign-positioning-dataset/common/constants.h>

#include <vector>

namespace spresults = yandex::maps::sproto::offline::mrc::results;

namespace maps::mrc::tracks_with_sensors {

db::TrackPoints loadTrack(const std::string& datasetPrefix)
{
    std::string resultsPath = datasetPrefix + dataset::RESULTS_FILENAME_POSTFIX;

    spresults::Results mrcResults;
    std::ifstream{resultsPath, std::ios::binary} >> mrcResults;

    db::TrackPoints resultTrackPoints(mrcResults.track().size());

    for (size_t i = 0; i < mrcResults.track().size(); i++) {
        const auto& trackPoint = mrcResults.track()[i];
        resultTrackPoints[i].setTimestamp(
            maps::chrono::TimePoint(
                std::chrono::nanoseconds(trackPoint.time())));
        resultTrackPoints[i].setGeodeticPos({
                trackPoint.location().point().lon(),
                trackPoint.location().point().lat()});

        if (trackPoint.location().speed()) {
            resultTrackPoints[i].setSpeedMetersPerSec(
                *trackPoint.location().speed());
        }
        if (trackPoint.location().heading()) {
            resultTrackPoints[i].setHeading(
                geolib3::Heading(*trackPoint.location().heading()));
        }
        if (trackPoint.location().accuracy()) {
            resultTrackPoints[i].setAccuracyMeters(
                *trackPoint.location().accuracy());
        }
    }

    return resultTrackPoints;
}

SensorEvents loadSensors(const std::string& datasetPrefix)
{
    std::string resultsPath = datasetPrefix + dataset::RESULTS_FILENAME_POSTFIX;

    spresults::Results mrcResults;
    std::ifstream{resultsPath, std::ios::binary} >> mrcResults;

    std::vector<std::stringstream> reports;

    INFO() << "loaded " << mrcResults.reports().size() << " reports";
    for (size_t i = 0; i < mrcResults.reports().size(); i++) {
        reports.push_back(std::stringstream(mrcResults.reports()[i]));
    }
    return sensors_feature_positioner::extractSensors(
        reports, std::nullopt, std::nullopt);
}

Photos loadPhotos(const std::string& datasetPrefix,
                  std::optional<chrono::TimePoint> minTime,
                  std::optional<chrono::TimePoint> maxTime)
{
    std::string metaPath = datasetPrefix + dataset::IMAGES_META_FILENAME_POSTFIX;
    std::string resultsPath = datasetPrefix + dataset::RESULTS_FILENAME_POSTFIX;

    spresults::Results mrcResults;
    std::ifstream{resultsPath, std::ios::binary} >> mrcResults;
    const maps::json::Value imagesMeta
        = maps::json::Value::fromFile(metaPath);

    REQUIRE(imagesMeta.size() == mrcResults.images().size(),
            "Images metadata must be aligned with MRC results!");

    Photos photos;
    for (size_t i = 0; i < imagesMeta.size(); i++) {
        auto orientation = common::ImageOrientation::fromExif(
            imagesMeta[i]["exif_orientation"].as<int>());
        cv::Mat decodedImage = common::decodeImage(mrcResults.images()[i].image());
        common::Size size{(std::size_t)decodedImage.cols,
                          (std::size_t)decodedImage.rows};
        decodedImage = common::transformByImageOrientation(decodedImage, orientation);
        geolib3::Point2 mercatorPos{mrcResults.images()[i].estimatedPosition()->point().lon(),
                                    mrcResults.images()[i].estimatedPosition()->point().lat()};
        chrono::TimePoint timestamp = maps::chrono::TimePoint(
            chrono::TimePoint::duration(imagesMeta[i]["feature_ts"].as<std::uint64_t>()));
        if ((maxTime && timestamp > maxTime)
            || (minTime && timestamp < minTime))
        {
            continue;
        }

        photos.push_back(Photo{
                imagesMeta[i]["feature_id"].as<std::int64_t>(),
                timestamp,
                mercatorPos,
                geolib3::Heading(*mrcResults.images()[i].estimatedPosition()->heading()),
                size,
                orientation});
    }

    INFO() << "loaded " << photos.size() << "photos";
    return photos;
}

SignPhotos loadSigns(const std::string& datasetPrefix,
                     const std::optional<std::unordered_set<db::TId>>& featureIds)
{
    std::string metaPath = datasetPrefix + dataset::IMAGES_META_FILENAME_POSTFIX;
    std::string resultsPath = datasetPrefix + dataset::RESULTS_FILENAME_POSTFIX;

    spresults::Results mrcResults;
    std::ifstream{resultsPath, std::ios::binary} >> mrcResults;
    const maps::json::Value imagesMeta
        = maps::json::Value::fromFile(metaPath);

    REQUIRE(imagesMeta.size() == mrcResults.images().size(),
            "Images metadata must be aligned with MRC results!");
    SignPhotos signPhotos;

    for (std::size_t idx = 0; idx < imagesMeta.size(); ++idx) {
        REQUIRE(mrcResults.images()[idx].created()
                    == imagesMeta[idx]["feature_ts"].as<std::uint64_t>(),
                "The data must be aligned!");

        const auto featureId
            = imagesMeta[idx]["feature_id"].as<std::int64_t>();
        if (featureIds && !featureIds->count(featureId)) {
            continue;
        }
        chrono::TimePoint featureTs = maps::chrono::TimePoint(
            chrono::TimePoint::duration(
                imagesMeta[idx]["feature_ts"].as<std::uint64_t>()));

        auto orientation = common::ImageOrientation::fromExif(
            imagesMeta[idx]["exif_orientation"].as<int>());
        cv::Mat decodedImage
            = common::decodeImage(mrcResults.images()[idx].image());

        for (const auto& sign : imagesMeta[idx]["signs"]) {
            const auto signType = sign["sign_type"].as<std::string>();
            const auto minX = sign["bbox"][0].as<std::size_t>();
            const auto minY = sign["bbox"][1].as<std::size_t>();
            const auto maxX = sign["bbox"][2].as<std::size_t>();
            const auto maxY = sign["bbox"][3].as<std::size_t>();

            common::Size photoSize = {(std::size_t)decodedImage.cols,
                                      (std::size_t)decodedImage.rows};

            const auto imageBox = revertByImageOrientation(
                {minX, minY, maxX, maxY},
                photoSize,
                orientation);

            signPhotos.push_back(SignPhoto{
                featureId,
                featureTs,
                photoSize,
                traffic_signs::stringToTrafficSign(signType),
                imageBox,
                sign["sign_id"].as<int>()});
        }
    }

    INFO() << "loaded " << signPhotos.size() << "signs";
    return signPhotos;
}


std::vector<Sign> loadGtSigns(const std::string& datasetPrefix,
                              const std::optional<std::unordered_set<SignGtId>>& signGtIds)
{
    std::string signsPath = datasetPrefix + dataset::SIGNS_FILENAME_POSTFIX;

    const maps::json::Value signsJson
        = maps::json::Value::fromFile(signsPath);

    std::vector<Sign> signs;
    signs.reserve(signsJson.size());

    for (std::size_t i = 0; i < signsJson.size(); i++) {
        const int signGtId = signsJson[i]["sign_id"].as<int>();
        if (signGtIds && !signGtIds->count(signGtId)) {
            continue;
        }

        geolib3::Point2 mercatorPos = geolib3::geoPoint2Mercator({
                signsJson[i]["wgs84_pos"][1].as<double>(),
                signsJson[i]["wgs84_pos"][0].as<double>()});
        signs.push_back(Sign{
            mercatorPos,
            traffic_signs::stringToTrafficSign(
                signsJson[i]["sign_type"].as<std::string>()),
                geolib3::Heading(signsJson[i]["azimuth"].as<double>()),
                signGtId});
    }
    return signs;
}

} // namespace maps::mrc::tracks_with_sensors
