#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/secure_config.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_panel/include/worker.h>

#include <maps/wikimap/mapspro/services/mrc/libs/birdview/include/panel_detector.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/opencv.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/load.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/io.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/exception.h>
#include <maps/libs/log8/include/log8.h>

#include <util/generic/size_literals.h>

#include <iterator>
#include <vector>

namespace maps::mrc::eye {

Partitions makePartitions(
        const std::vector<db::eye::Frames>& sequences,
        size_t size,
        size_t margin)
{
    const size_t minSize = margin + 1;
    const size_t maxSize = std::max(size_t(1.5 * size), size + minSize);

    ASSERT(size >= minSize);

    std::vector<PartedFrame> result;

    db::TId partId = 0;
    for (const auto& sequence: sequences) {
        if (sequence.size() < minSize) {
            continue;
        }

        auto begin = sequence.begin();
        auto end = begin;

        for (; begin != sequence.end(); begin = end) {
            end = size_t(sequence.end() - begin) >= maxSize
                ? begin + size
                : sequence.end();

            std::transform(
                begin, end,
                std::back_inserter(result),
                [&](const auto& frame) {
                    return PartedFrame {partId, frame};
                }
            );

            ++partId;
        }
    }

    return {result, size_t(partId)};
}

void savePartitions(NYT::IClientBase& client, const TString& path, const Partitions& partitions)
{
    const auto pathWithSchema = NYT::TRichYPath(path).Schema(
        yt::getSchemaOf<PartedFrame>()
            .SortBy({"part_id"})
    );

    yt::saveToTable(client, pathWithSchema, partitions.frames);

    constexpr size_t chunkSize = 2_KB;

    const auto spec = NYT::TNode::CreateMap()
        ("job_io", NYT::TNode::CreateMap()
             ("table_writer", NYT::TNode::CreateMap()
                 ("desired_chunk_size", chunkSize)
                 ("block_size", chunkSize)
             )
        );

    client.Merge(
        NYT::TMergeOperationSpec()
            .ForceTransform(true)
            .Mode(NYT::EMergeMode::MM_ORDERED)
            .AddInput(pathWithSchema)
            .Output(pathWithSchema),
        NYT::TOperationOptions()
            .Spec(spec)
    );
}

class DetectPanelWorker: public yt::Reducer {

public:
    DetectPanelWorker() = default;

    DetectPanelWorker(const FrameLoader& frameLoader, size_t margin)
        : frameLoader_(frameLoader),
          margin_(margin)
    {}

    void Do(yt::Reader* reader, yt::Writer* writer) override;

    Y_SAVELOAD_JOB(frameLoader_, margin_);

private:
    FrameLoader frameLoader_;
    size_t margin_;
};

void DetectPanelWorker::Do(yt::Reader* reader, yt::Writer* writer)
{
    const size_t coldStartN = margin_ + 1;

    birdview::PanelDetector detector;

    std::optional<common::Size> globalSize;
    bool coldStartReady = false;
    db::eye::Frames coldStartFrames = {};
    common::ImageBox coldStartPanelBox = {};

    for (; reader->IsValid(); reader->Next()) {
        const auto [partId, frame] = yt::deserialize<PartedFrame>(reader->GetRow());
        const cv::Mat image = frameLoader_.load(frame);

        const auto size = common::sizeOf(image);
        ASSERT(size.width and size.height);

        if (not globalSize) {
            globalSize = frame.size();
        } else if (*globalSize != size) {
            detector = birdview::PanelDetector{};

            globalSize = frame.size();
            coldStartReady = false;
            coldStartFrames = {};
            coldStartPanelBox = {};
        }

        detector.update(image);
        const size_t row = detector.panelRow() > 0 ? detector.panelRow() : 0;
        const common::ImageBox panelBox = {0, 0, size.width, row};

        if (coldStartFrames.size() < coldStartN) {
            coldStartFrames.push_back(frame);
        }

        if (coldStartReady) {
            writer->AddRow(
                yt::serialize(
                    DetectedPanel{
                        frame,
                        common::revertByImageOrientation(
                            panelBox,
                            frame.originalSize(),
                            frame.orientation()
                        )
                    }
                )
            );
        } else if (coldStartFrames.size() == coldStartN) {
            for (const db::eye::Frame& coldStartFrame: coldStartFrames) {
                writer->AddRow(
                    yt::serialize(
                        DetectedPanel{
                            coldStartFrame,
                            common::revertByImageOrientation(
                                panelBox,
                                coldStartFrame.originalSize(),
                                coldStartFrame.orientation()
                            )
                        }
                    )
                );
            }

            coldStartReady = true;
        }
    }
}

REGISTER_REDUCER(DetectPanelWorker);

DetectedPanels detectPanel(const DetectPanelConfig& config, const Partitions& partitions)
{
    {   // check config
        REQUIRE(isValid(config.yt), "Invalid YT config");
    }

    if (not partitions.partitionNumber) {
        return {};
    }

    NYT::ITransactionPtr ytTxn = config.yt.client->StartTransaction();

    const TString root(config.yt.rootPath);
    if (!ytTxn->Exists(root)) {
         ytTxn->Create(root, NYT::NT_MAP, NYT::TCreateOptions().Recursive(true));
        INFO() << "Create node " << root;
    }

    const TString input = root + "/partition";
    const TString output = root + "/recognition";

    INFO() << "Save to table " << input;
    savePartitions(*ytTxn, input, partitions);

    REQUIRE(SecureConfig::isInitialized(), "Secure config was not initialized");

    ytTxn->Reduce(
        NYT::TReduceOperationSpec()
            .AddInput<NYT::TNode>(input)
            .AddOutput<NYT::TNode>(
                NYT::TRichYPath(output)
                    .Schema(yt::getSchemaOf<DetectedPanel>())
            )
            .SortBy({"part_id"})
            .ReduceBy({"part_id"})
            .JobCount(partitions.partitionNumber),
        new DetectPanelWorker(*(config.yt.frameLoader), config.margin),
        NYT::TOperationOptions().Spec(
            yt::baseCpuOperationSpec("eye::detect_panel", config.yt.poolType)
                ("reducer", NYT::TNode::CreateMap()
                    ("memory_limit", 1_GB)
                )
                ("resource_limits", NYT::TNode::CreateMap()
                    ("user_slots", config.yt.concurrency)
                )
                ("max_failed_job_count", 30)
        )
        .SecureVault(SecureConfig::instance())
    );

    INFO() << "Load table " << output;
    DetectedPanels panels = yt::loadFromTable<DetectedPanels>(*ytTxn, output);

    if (config.yt.commit) {
        INFO() << "Yt commited!";
        ytTxn->Commit();
    }

    return panels;
}

} // namespace maps::mrc::eye
