#include "matcher.h"
#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/graph_matcher_adapter/include/feature_positioner.h>
#include <maps/libs/geolib/include/distance.h>
#include <maps/libs/log8/include/log8.h>

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

#include <unordered_map>
#include <algorithm>
#include <optional>

namespace maps::mrc::opensfm_experiment {

namespace {

db::Features getFilteredFeatures(
    NYT::IClientBase& client,
    const TString& table,
    const std::chrono::seconds& minInterval,
    double minDistance
) {
    auto features = yt::loadFromTable<db::Features>(client, table);

    std::sort(
        features.begin(),
        features.end(),
        [](const auto& lhs, const auto& rhs) {return lhs.timestamp() < rhs.timestamp();}
    );

    const db::Feature* prev = nullptr;
    const auto it = std::remove_if(
        features.begin(), features.end(),
        [&prev, &minInterval, &minDistance](const auto& current) -> bool {
            if (not prev) {
                prev = &current;
                return false;
            }

            bool tooCloseTime = current.timestamp() - prev->timestamp() < minInterval;
            bool tooCloseDistance = maps::geolib3::fastGeoDistance(current.geodeticPos(), prev->geodeticPos()) < minDistance;

            if (not tooCloseTime and not tooCloseDistance) {
                prev = &current;
            }

            return tooCloseTime or tooCloseDistance;
        }
    );

    features.erase(it, features.end());

    INFO() << features.size() << " features after filtering";
    return features;
}


std::unordered_map<std::string, db::TrackPoints> toTrackPoints(const db::Features& features) {
    std::unordered_map<std::string, db::TrackPoints> points;

    for (const auto& feature : features) {
        db::TrackPoint point;

        point.setSourceId(feature.sourceId())
             .setTimestamp(feature.timestamp())
             .setMercatorPos(feature.mercatorPos());

        points[feature.sourceId()].push_back(point);
    }

    return points;
}


db::Features matchFeaturesToGraph(const db::Features& rawFeatures) {
    auto trackPoints = toTrackPoints(rawFeatures);

    auto tracksProvider = [&trackPoints] (
        const std::string& sourceId,
        chrono::TimePoint startTime,
        chrono::TimePoint endTime
    ) {
        db::TrackPoints result;
        db::TrackPoint startTrackPoint = db::TrackPoint().setTimestamp(startTime);

        auto begin = std::lower_bound(
            trackPoints[sourceId].begin(),
            trackPoints[sourceId].end(),
            startTrackPoint,
            [](const auto& lhs, const auto& rhs) {return lhs.timestamp() < rhs.timestamp();}
        );

        db::TrackPoint endTrackPoint = db::TrackPoint().setTimestamp(endTime);
        auto end = std::upper_bound(
            trackPoints[sourceId].begin(),
            trackPoints[sourceId].end(),
            endTrackPoint,
            [](const auto& lhs, const auto& rhs) {return lhs.timestamp() < rhs.timestamp();}
        );

        std::copy(begin, end, std::back_inserter(result));
        return result;
    };

    db::Features result(rawFeatures);

    adapters::CompactGraphMatcherAdapter matcher("/var/spool/yandex/maps/graph/19.08.29-0/");
    adapters::FeaturePositioner positioner({{db::GraphType::Road, &matcher}}, tracksProvider);

    positioner(result);
    return result;
}

} // namespace

void refinePositions(
    NYT::IClientBase& client,
    const TString& input,
    const TString& output,
    const std::chrono::seconds& minInterval,
    double minDistance,
    bool matchToGraph
) {
    INFO() << "Loading features from " << input;
    auto features = getFilteredFeatures(client, input, minInterval, minDistance);
    if (matchToGraph) {
        INFO() << "Matching features";
        features = matchFeaturesToGraph(features);
    }

    INFO() << "Writing updated positions and headings to " << output;
    auto newFeatures = std::unordered_map<db::TId, const db::Feature*>();

    for (const auto& feature : features) {
        newFeatures[feature.id()] = &feature;
    }

    auto reader = client.CreateTableReader<NYT::TNode>(input);
    auto writer = client.CreateTableWriter<NYT::TNode>(output);

    for (; reader->IsValid(); reader->Next()) {
        GroundTruthFeature row = deserialize(reader->GetRow());

        if (newFeatures.find(row.feature.id()) == newFeatures.end() or !newFeatures[row.feature.id()]->hasHeading()) {
            continue;  // input may contain more features (before filtering by time and distance)
        }

        auto newFeature = newFeatures[row.feature.id()];

        row.feature.setGeodeticPos(newFeature->geodeticPos());
        row.feature.setHeading(newFeature->heading());

        writer->AddRow(serialize(row));
    }
}

}; // namespace maps::mrc::opensfm_experiment
