#include "context.h"
#include "coverage_export.h"
#include "features_index.h"
#include "geobase.h"
#include "graph_coverage.h"
#include "graph_length.h"
#include "strings.h"
#include "yt_tables.h"

#include <maps/wikimap/mapspro/services/mrc/libs/common/include/algorithm/parallel_for_each.h>
#include <maps/wikimap/mapspro/services/mrc/libs/object/include/revision_loader.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/common.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/libs/common/include/math.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/introspection/include/comparison.h>
#include <maps/libs/log8/include/log8.h>
#include <maps/libs/xrange_view/include/xrange_view.h>
#include <maps/libs/introspection/include/stream_output.h>

#include <mapreduce/yt/library/parallel_io/parallel_writer.h>
#include <mapreduce/yt/interface/client.h>
#include <sys/types.h>

#include <library/cpp/geobase/lookup.hpp>
#include <util/generic/size_literals.h>

#include <pqxx/transaction_base>

#include <map>
#include <unordered_set>
#include <vector>

using namespace std::literals::chrono_literals;

namespace maps::mrc::graph_coverage_export {

namespace {

class GenerateTimelineReducer: public yt::Reducer
{
    struct DatedLengthDiff {
        int32_t date;
        int32_t actualizationDate;
        double lengthDiff;
    };

    struct CoverageChangesLog {
        CoverageEdgeData coverageEdgeData;
        std::vector<DatedLengthDiff> sortedLengthDiffs;
    };

public:
    Y_SAVELOAD_JOB(timelineEndDate_);

    GenerateTimelineReducer() = default;

    GenerateTimelineReducer(int32_t timelineEndDate)
        : timelineEndDate_(timelineEndDate)
    {}

    void Do(yt::Reader* reader, yt::Writer* writer) override
    {
        auto coverageChangesLog = readCoverageChangesLog(reader);

        auto it = std::begin(coverageChangesLog.sortedLengthDiffs);
        ASSERT(timelineEndDate_ < maxDate().count());

        std::map<int32_t, double> actDateToLengthDiffMap;

        auto end = std::end(coverageChangesLog.sortedLengthDiffs);
        while(it != end && it->date <= timelineEndDate_)
        {
            auto intervalIt = it;
            for (; intervalIt != end &&
                    intervalIt->date == it->date; ++intervalIt)
            {
                actDateToLengthDiffMap[intervalIt->actualizationDate] +=
                    intervalIt->lengthDiff;
            }

            DatedCoverageEdgeData datedCoverageEdgeData{
                .coverageEdgeData = coverageChangesLog.coverageEdgeData,
                .date = it->date
            };

            do {
                auto classifiedLengths = calculateAgeCategoriesLengths(
                    actDateToLengthDiffMap, datedCoverageEdgeData.date);
                double totalLength = 0;
                for (auto [ageCategory, length]: classifiedLengths) {
                    totalLength += length;
                    datedCoverageEdgeData.coverageEdgeData.length = length;
                    datedCoverageEdgeData.ageCategory = static_cast<int32_t>(ageCategory);
                    writer->AddRow(yt::serialize(datedCoverageEdgeData));
                }

                datedCoverageEdgeData.date += 1;
            } while(datedCoverageEdgeData.date < std::min(intervalIt->date, timelineEndDate_));
            it = intervalIt;
        }
    }

private:

    CoverageChangesLog readCoverageChangesLog(yt::Reader* reader) const
    {
        std::vector<DatedLengthDiff> dataArray;
        std::optional<CoverageEdgeData> coverageEdgeData;
        for (; reader->IsValid(); reader->Next()) {
            auto row = reader->GetRow();
            auto intervalCoverageEdgeData = yt::deserialize<IntervalCoverageEdgeData>(row);
            if (! coverageEdgeData.has_value()) {
                coverageEdgeData = intervalCoverageEdgeData.coverageEdgeData;
            }
            dataArray.push_back(
                {
                    intervalCoverageEdgeData.fromDate,
                    intervalCoverageEdgeData.actualizationDate,
                    intervalCoverageEdgeData.coverageEdgeData.length
                }
            );
            dataArray.push_back(
                {
                    intervalCoverageEdgeData.toDate,
                    intervalCoverageEdgeData.actualizationDate,
                    -intervalCoverageEdgeData.coverageEdgeData.length
                }
            );
        }

        std::sort(
            std::begin(dataArray),
            std::end(dataArray),
            [](const auto& one, const auto& other) { return one.date < other.date; });

        ASSERT(coverageEdgeData.has_value());
        return {.coverageEdgeData = std::move(coverageEdgeData.value()),
                .sortedLengthDiffs = std::move(dataArray)};
    }

    CoverageAgeCategory ageCategory(int32_t days) const {
        if (days < 7) {
            return CoverageAgeCategory::Age_0_7;
        } else if (days < 30) {
            return CoverageAgeCategory::Age_7_30;
        } else if (days < 90) {
            return CoverageAgeCategory::Age_30_90;
        } else if (days < 180) {
            return CoverageAgeCategory::Age_90_180;
        } else {
            return CoverageAgeCategory::Age_180;
        }
    }

    std::map<CoverageAgeCategory, double>
    calculateAgeCategoriesLengths(
        const std::map<int32_t, double>& actDateToLengthDiffMap,
        int32_t date) const
    {
        std::map<CoverageAgeCategory, double> result;
        for (auto [actDate, lengthDiff] : actDateToLengthDiffMap) {
            REQUIRE(actDate <= date, "actDate must be less then date, but "
                << actDate << " > " << date);
            result[ageCategory(date - actDate)] += lengthDiff;
        }
        return result;
    }

    int32_t timelineEndDate_;
};

REGISTER_REDUCER(GenerateTimelineReducer);

using GeoIdsSet = std::unordered_set<int32_t>;

void writeRegionsToYt(NYT::IClientBase& ytClient,
                      const std::string& ytPath,
                      const GeoIdsSet& geoIds,
                      geobase::GeobasePtr geobasePtr)
{
    INFO() << "Writing regions to " << ytPath;
    const auto regionsWriter =
        ytClient.CreateTableWriter<NYT::TNode>(
            NYT::TRichYPath(TString(ytPath))
                    .Schema(yt::getSchemaOf<Region>())
                    .OptimizeFor(NYT::OF_SCAN_ATTR));

    for (auto geoId: geoIds) {
        regionsWriter->AddRow(yt::serialize(
            Region{geoId, geobasePtr->getRegionById(geoId).name})
        );
    }

    regionsWriter->Finish();
}

} // namespace

void coverageExport(const common::Config& config,
                    const wiki::common::ExtendedXmlDoc& wikiConfig,
                    const std::string& mrcDatasetPath,
                    const std::string& mrcFeaturesSecretDatasetPath,
                    const std::string& geobasePath,
                    const std::string& geodataPatchYtDir,
                    const std::string& mrcRoadGraphPath,
                    const std::string& mrcPedestrianGraphPath,
                    const TString& ytDir,
                    const std::optional<geolib3::BoundingBox>& geoBbox)
{
    INFO() << "Output YT dir: " << ytDir;
    NYT::IClientPtr ytClient = config.externals().yt().makeClient();
    wiki::common::PoolHolder mrcPoolHolder =
        config.makePoolHolder(maps::mrc::common::LONG_READ_DB_ID,
                              maps::mrc::common::LONG_READ_POOL_ID);

    geobase::GeobasePtr geobasePtr =
        loadPatchedGeobasePtr(ytClient,
                              geodataPatchYtDir,
                              std::make_unique<geobase::Geobase6Adapter>(geobasePath));

    wiki::common::PoolHolder wikiPoolHolder(
        wikiConfig, "long-read", "long-read");
    auto objectLoader =
        object::makeRevisionLoader(wikiPoolHolder.pool().slaveTransaction());

    auto ytTxn = ytClient->StartTransaction();

    ytTxn->Create(
        ytDir, NYT::NT_MAP, NYT::TCreateOptions().Recursive(true).Force(true));


    INFO() << "Loading features index";
    FeaturesIndex featuresIndex(mrcDatasetPath, mrcFeaturesSecretDatasetPath);
    INFO() << "Loaded features index";

    TString edgesTableName = ytDir + "/" + EdgeData::YT_TABLE_NAME;
    GeoIdsSet referencedGeoIds;

    const struct {
        std::string path;
        db::GraphType type;
    } graphs[] = {{mrcRoadGraphPath, db::GraphType::Road},
                  {mrcPedestrianGraphPath, db::GraphType::Pedestrian}};

    auto graphLengthTableWriter = ytTxn->CreateTableWriter<NYT::TNode>(
            NYT::TRichYPath(ytDir + "/" + GraphLength::YT_TABLE_NAME)
                .Schema(yt::getSchemaOf<GraphLength>())
                .OptimizeFor(NYT::OF_SCAN_ATTR)
        );

    for (const auto& [graphPath, graphType] : graphs) {
        INFO() << "Processing graph " << graphType;
        Context ctx(graphPath, graphType, geobasePtr, *objectLoader);
        fb::CoverageRtreeReader coverageRtreeReader(ctx.graph(), graphPath);

        INFO() << "Computing graph length";
        saveGraphLengthToYt(ctx, *graphLengthTableWriter, geoBbox);

        const auto edgesWriter =
            NYT::CreateParallelUnorderedTableWriter<NYT::TNode>(
                ytTxn, NYT::TRichYPath(edgesTableName).Append(true));

        INFO() << "Writing data to " << edgesTableName;

        auto coveredEdgeIdsRange =
            geoBbox.has_value() ? coverageRtreeReader.getCoveredEdgeIdsByBbox(geoBbox.value())
                : coverageRtreeReader.getCoveredEdgeIds();

        common::parallelForEach<THREADS_NUMBER>(
            coveredEdgeIdsRange.begin(),
            coveredEdgeIdsRange.end(),
            [&](auto& mutex, road_graph::EdgeId edgeId)
            {
                road_graph::EdgeData edgeData = ctx.graph().edgeData(edgeId);
                if (!areCompatible(edgeData.accessIdMask(), ctx.graphType())) {
                    return;
                }
                geolib3::Polyline2 edgeGeometry = edgeData.geometry();
                db::Features features = searchFeaturesForEdge(edgeGeometry, featuresIndex);

                std::lock_guard<decltype(mutex)> lock(mutex);
                ctx.processEdge(edgeData, edgeGeometry, features,
                                referencedGeoIds, *edgesWriter);
            }
        );

        edgesWriter->Finish();
        INFO() << "Finished writing data to " << edgesTableName;
    }
    graphLengthTableWriter->Finish();

    writeRegionsToYt(*ytTxn, ytDir + "/" + Region::YT_TABLE_NAME, referencedGeoIds, geobasePtr);
    convertToCoverage(*ytTxn, edgesTableName, ytDir + "/" + IntervalCoverageEdgeData::YT_TABLE_NAME);

    calculateGraphCoverageTimeline(
        *ytTxn,
        chrono::sinceEpoch<std::chrono::days>(),
        ytDir + "/" + IntervalCoverageEdgeData::YT_TABLE_NAME,
        ytDir + "/" + DatedCoverageEdgeData::YT_TABLE_NAME,
        yt::PoolType::Processing);

    ytTxn->Commit();

    INFO() << "Done writing YT tables";
}

void calculateGraphCoverageTimeline(
    NYT::IClientBase& ytClient,
    int32_t timelineEnd,
    const TString& ytGraphCoveragePath,
    const TString& ytOutputPath,
    std::optional<yt::PoolType> poolType)
{
    INFO() << "Calculating coverage timeline";
    auto columnsToReduceBy = NYT::TSortColumns(
        {col::CAMERA_DEVIATION,
         col::DATASET,
         col::GRAPH_TYPE,
         col::GEO_ID,
         col::FC,
         col::IS_TOLL,
         col::PRIVATE_AREA});

    ytClient.Sort(
        NYT::TSortOperationSpec()
            .AddInput(ytGraphCoveragePath)
            .Output(ytGraphCoveragePath)
            .SortBy(columnsToReduceBy),
        NYT::TOperationOptions().Spec(
            yt::baseCpuOperationSpec("sort_coverage_timeline", poolType)
        ));

    ytClient.Reduce(
        NYT::TReduceOperationSpec()
            .AddInput<NYT::TNode>(ytGraphCoveragePath)
            .AddOutput<NYT::TNode>(
                NYT::TRichYPath(ytOutputPath)
                    .Schema(yt::getSchemaOf<DatedCoverageEdgeData>())
                    .OptimizeFor(NYT::OF_SCAN_ATTR)
            )
            .ReduceBy(columnsToReduceBy),
        new GenerateTimelineReducer(timelineEnd ),
        NYT::TOperationOptions().Spec(
            yt::baseCpuOperationSpec("calculate_coverage_timeline", poolType)
                ("reducer", NYT::TNode::CreateMap()
                    ("memory_limit", 4_GB)
                )
        )
    );

    ytClient.Sort(
        NYT::TSortOperationSpec()
            .AddInput(ytOutputPath)
            .Output(ytOutputPath)
            .SortBy({col::DATE}),
        NYT::TOperationOptions().Spec(
            yt::baseCpuOperationSpec("sort_coverage_timeline", poolType)
        ));


}

}  // namespace maps::mrc::graph_coverage_export
