#include <maps/libs/log8/include/log8.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/common/include/file_utils.h>
#include <maps/libs/cmdline/include/cmdline.h>

#include <maps/wikimap/mapspro/services/autocart/libs/geometry/include/hex_wkb.h>
#include <maps/wikimap/mapspro/services/autocart/libs/utils/include/multithreading.h>
#include <maps/wikimap/mapspro/services/autocart/libs/satellite/include/load_sat_image.h>
#include <maps/wikimap/mapspro/services/autocart/libs/satellite/include/load_sat_params.h>

#include <maps/libs/csv/include/rows.h>
#include <maps/libs/csv/include/iterator.h>
#include <maps/libs/csv/include/input_stream.h>
#include <maps/libs/csv/include/output_stream.h>

#include <maps/libs/geolib/include/polygon.h>
#include <maps/libs/geolib/include/conversion.h>
#include <maps/libs/geolib/include/bounding_box.h>

#include <maps/libs/tile/include/utils.h>

#include <util/system/fs.h>
#include <util/random/random.h>
#include <util/random/entropy.h>
#include <util/random/mersenne.h>

#include <opencv2/opencv.hpp>

#include <string>
#include <fstream>

namespace csv = maps::csv;
namespace tile = maps::tile;
namespace common = maps::common;
namespace geolib3 = maps::geolib3;
namespace autocart = maps::wiki::autocart;

namespace {

using IdToHeight = std::unordered_map<size_t, size_t>;
using HeightAndRequirement = std::pair<size_t, size_t>;
using HeightToRequirement = std::unordered_map<size_t, size_t>;
using HeightToLoadedCount = std::unordered_map<size_t, TAtomic>;

struct Building {
    size_t id;
    geolib3::Polygon2 geom;
    size_t height;
};

using HeightToBuildingVec
    = std::unordered_map<size_t, std::vector<Building>>;

struct BuildingWithData : public Building {
    cv::Mat image;
    cv::Mat mask;
    autocart::SatelliteParams satParams;
};

HeightToRequirement loadRequirements(const std::string& path) {
    static const std::string HEIGHT_COLUMN_NAME = "height";
    static const std::string REQUIREMENT_COLUMN_NAME = "requirement";
    std::ifstream ifs(path);
    REQUIRE(ifs.is_open(), "Failed to open file: " + path);
    csv::MapRowRange rows(ifs);
    HeightToRequirement heightToRequirement;
    for (const auto& row : rows) {
        size_t height = row[HEIGHT_COLUMN_NAME].as<size_t>();
        size_t requirement = row[REQUIREMENT_COLUMN_NAME].as<size_t>();
        heightToRequirement[height] = requirement;
    }
    return heightToRequirement;
}

IdToHeight loadBldsWithSuitableHeight(const std::string& ymapsdfPath,
                                      const std::set<size_t> heights) {
    static const std::string BLD_TABLE = "bld.csv";
    static const std::string BLD_ID_COLUMN_NAME = "bld_id";
    static const std::string HEIGHT_COLUMN_NAME = "height";

    std::string bldTablePath = common::joinPath(ymapsdfPath, BLD_TABLE);
    std::ifstream ifs(bldTablePath);
    REQUIRE(ifs.is_open(), "Failed to open file: " + bldTablePath);
    csv::MapRowRange rows(ifs);

    IdToHeight idToHeight;
    for (const auto& row : rows) {
        if (row[HEIGHT_COLUMN_NAME].isEmpty()) {
            continue; // building without height
        }
        size_t id = row[BLD_ID_COLUMN_NAME].as<size_t>();
        size_t height = row[HEIGHT_COLUMN_NAME].as<size_t>();
        if (heights.find(height) != heights.end()) {
            idToHeight[id] = height;
        }
    }
    return idToHeight;
}

HeightToBuildingVec
loadBldsGeometry(const std::string& ymapsdfPath,
                 const IdToHeight& idToHeight) {
    static const std::string BLD_GEOM_TABLE = "bld_geom.csv";
    static const std::string BLD_ID_COLUMN_NAME = "bld_id";
    static const std::string SHAPE_COLUMN_NAME = "shape";

    std::string bldGeomTablePath
        = common::joinPath(ymapsdfPath, BLD_GEOM_TABLE);
    std::ifstream ifs(bldGeomTablePath);
    REQUIRE(ifs.is_open(), "Failed to open file: " + bldGeomTablePath);
    csv::MapRowRange rows(ifs);

    HeightToBuildingVec heightToBuildingVec;

    for (const auto& row : rows) {
        size_t id = row[BLD_ID_COLUMN_NAME].as<size_t>();
        auto it = idToHeight.find(id);
        if (it != idToHeight.end()) {
            Building bld;
            bld.id = id;
            bld.height = it->second;
            std::string_view hexWKB = row[SHAPE_COLUMN_NAME];
            bld.geom = autocart::hexWKBToPolygon(hexWKB);
            heightToBuildingVec[bld.height].push_back(bld);
        }
    }

    for (auto& [height, blds] : heightToBuildingVec) {
        std::shuffle(
            blds.begin(), blds.end(),
            TMersenne<ui64>(Seed())
        );
    }

    return heightToBuildingVec;
}

cv::Mat drawBuilding(const geolib3::BoundingBox& bbox,
                     const geolib3::Polygon2& mercatorBld,
                     size_t zoom) {
    auto firstCorner = tile::mercatorToDisplaySigned(bbox.lowerCorner(), zoom);
    auto secondCorner = tile::mercatorToDisplaySigned(bbox.upperCorner(), zoom);
    tile::SignedDisplayCoord leftTop(std::min(firstCorner.x(), secondCorner.x()),
                                     std::min(firstCorner.y(), secondCorner.y()));
    size_t maskWidth = abs(firstCorner.x() - secondCorner.x()) + 1;
    size_t maskHeight = abs(firstCorner.y() - secondCorner.y()) + 1;

    std::vector<cv::Point2i> polygon;
    for (size_t i = 0; i < mercatorBld.pointsNumber(); i++) {
        geolib3::Point2 point = mercatorBld.pointAt(i);
        tile::SignedDisplayCoord dispPoint = tile::mercatorToDisplaySigned(point, zoom);
        polygon.emplace_back(dispPoint.x() - leftTop.x(), dispPoint.y() - leftTop.y());
    }

    cv::Mat mask(maskHeight, maskWidth, CV_8UC1, cv::Scalar::all(0));
    std::vector<std::vector<cv::Point2i>> polygons(1, polygon);

    cv::fillPoly(mask, polygons, cv::Scalar::all(255));

    return mask;
}

std::optional<BuildingWithData>
downloadData(const Building& bld, double gapSize, size_t zoom) {
    geolib3::Polygon2 mercatorBld
        = geolib3::convertGeodeticToMercator(bld.geom);
    geolib3::BoundingBox mercatorBbox = mercatorBld.boundingBox();
    mercatorBbox = geolib3::resizeByRatio(mercatorBbox, 1 + gapSize);
    geolib3::BoundingBox geodeticBbox
        = geolib3::convertMercatorToGeodetic(mercatorBbox);

    BuildingWithData bldWithData;
    bldWithData.id = bld.id;
    bldWithData.geom = bld.geom;
    bldWithData.height = bld.height;

    try {
        bldWithData.image = autocart::loadSatImage(mercatorBbox, zoom);
    } catch (std::exception&) {
        INFO() << "Failed to download image for "
               << bld.id << " image";
        return std::nullopt;
    }

    bldWithData.mask = drawBuilding(mercatorBbox, mercatorBld, zoom);

    bldWithData.satParams
        = autocart::getSatelliteParams(bldWithData.image.size(), geodeticBbox);

    return bldWithData;
}

using BldQueue = TLockFreeQueue<Building, autocart::QueueCounter<Building>>;
using BldWithDataQueue
    = TLockFreeQueue<BuildingWithData, autocart::QueueCounter<BuildingWithData>>;

void wait() {
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
}

bool isRequirementStatisfied(size_t requirement, const size_t& loadedCount) {
    return loadedCount >= requirement;
}

class DataDownloader : public autocart::ObjectInQueueWithData {
public:
    DataDownloader(BldQueue* inpQueue, BldWithDataQueue* outQueue,
                   double gapSize, size_t zoom,
                   const HeightToRequirement& heightToRequirement,
                   HeightToLoadedCount* heightToLoadedCount)
        : inpQueue_(inpQueue),
          outQueue_(outQueue),
          gapSize_(gapSize),
          zoom_(zoom),
          heightToRequirement_(heightToRequirement),
          heightToLoadedCount_(heightToLoadedCount),
          running_(1) {}

    void Process(void* /*threadSpecificResource*/) override {
        for (;;) {
            Building bld;
            if (inpQueue_->Dequeue(&bld)) {
                Do(bld);
            } else if (!isWaitData()) {
                break;
            } else {
                wait();
            }
        }
        AtomicSet(running_, 0);
    }

    bool IsRunning() const {
        return (0 != AtomicGet(running_));
    }

private:
    void Do(const Building& bld) {
        const size_t& height = bld.height;
        const size_t& requirement = heightToRequirement_.at(height);
        size_t loadedCount = AtomicGet(heightToLoadedCount_->at(height));
        if (isRequirementStatisfied(requirement, loadedCount)) {
            return;
        }
        std::optional<BuildingWithData> data = downloadData(bld, gapSize_, zoom_);
        if (data) {
            size_t prevLoadedCount
                = AtomicGetAndIncrement(heightToLoadedCount_->at(height));
            if (!isRequirementStatisfied(requirement, prevLoadedCount)) {
                outQueue_->Enqueue(*data);
            }
        }
    }

    BldQueue* inpQueue_;
    BldWithDataQueue* outQueue_;
    double gapSize_;
    size_t zoom_;
    const HeightToRequirement& heightToRequirement_;
    HeightToLoadedCount* heightToLoadedCount_;
    TAtomic running_;
};


class DataWriter : public autocart::ObjectInQueueWithData {
public:
    DataWriter(BldWithDataQueue* inpQueue,
               const std::string& datasetPath,
               const HeightToRequirement& heightToRequirement)
        : inpQueue_(inpQueue),
          datasetPath_(datasetPath),
          heightToRequirement_(heightToRequirement),
          heightToRemainingRequirement_(heightToRequirement),
          imagePath_(common::joinPath(datasetPath_, "image")),
          maskPath_(common::joinPath(datasetPath_, "mask")),
          csvWriter_(dataCSVOfs_, csv::COMMA),
          index_(0),
          running_(1) {
        NFs::MakeDirectoryRecursive(TString(imagePath_));
        NFs::MakeDirectoryRecursive(TString(maskPath_));
        std::string csvPath = common::joinPath(datasetPath_, "data.csv");
        dataCSVOfs_.open(csvPath);
        REQUIRE(dataCSVOfs_.is_open(), "Faild to create file: " + csvPath);
        // CSV header
        csvWriter_ << "id"
                   << "height"
                   << "elev" << "azim"
                   << "pixel_width" << "pixel_height"
                   << csv::endl;
    }

    void Process(void* /*threadSpecificResource*/) override {
        for (;;) {
            BuildingWithData bld;
            if (inpQueue_->Dequeue(&bld)) {
                Do(bld);
            } else if (!isWaitData()) {
                break;
            } else {
                wait();
            }
        }
        dataCSVOfs_.close();
        for (const auto& [height, requirement] : heightToRequirement_) {
            const size_t& remainingRequirement
                = heightToRemainingRequirement_[height];
            INFO() << "Height: " << height << ", "
                   << "saved count: " << requirement - remainingRequirement;
        }
        AtomicSet(running_, 0);
    }

    bool IsRunning() const {
        return (0 != AtomicGet(running_));
    }

private:
    void Do(const BuildingWithData& bld) {
        if (heightToRemainingRequirement_[bld.height] > 0) {
            heightToRemainingRequirement_[bld.height] -= 1;
        } else {
            return;
        }

        csvWriter_ << bld.id << bld.height;

        if (bld.satParams.angles) {
            csvWriter_ << bld.satParams.angles->elev;
            csvWriter_ << bld.satParams.angles->azimAngle;
        } else {
            csvWriter_ << "" << "";
        }

        csvWriter_ << bld.satParams.pixelSize.widthInMeters;
        csvWriter_ << bld.satParams.pixelSize.heightInMeters;
        csvWriter_ << csv::endl;

        cv::imwrite(common::joinPath(imagePath_, std::to_string(bld.id) + ".jpg"), bld.image);
        cv::imwrite(common::joinPath(maskPath_, std::to_string(bld.id) + ".png"), bld.mask);

        index_++;
        if (index_ % 100 == 0) {
            INFO() << "Processed " << index_ << " items";
        }
    }

    BldWithDataQueue* inpQueue_;
    std::string datasetPath_;
    const HeightToRequirement& heightToRequirement_;
    HeightToRequirement heightToRemainingRequirement_;
    std::string imagePath_;
    std::string maskPath_;
    std::ofstream dataCSVOfs_;
    csv::OutputStream csvWriter_;
    size_t index_;
    TAtomic running_;
};

void readBldsToQueue(const HeightToBuildingVec& heightToBuildingVec,
                     BldQueue* bldQueue,
                     const HeightToRequirement& heightToRequirement,
                     const HeightToLoadedCount& heightToLoadedCount) {
    constexpr size_t PREFETCH_SIZE = 100;
    for (const auto& [height, buildingVec] : heightToBuildingVec) {
        const TAtomic& loadedCount = heightToLoadedCount.at(height);
        const size_t& requirement = heightToRequirement.at(height);
        for (size_t i = 0; i < buildingVec.size(); i++) {
            if (isRequirementStatisfied(requirement, AtomicGet(loadedCount))) {
                break;
            }
            bldQueue->Enqueue(buildingVec[i]);
            for (; !isRequirementStatisfied(requirement, AtomicGet(loadedCount));) {
                auto counter = bldQueue->GetCounter();
                if (PREFETCH_SIZE > static_cast<size_t>(AtomicGet(counter.Counter))) {
                    break;
                }
                wait();
            }
        }
    }
}

void dump(const std::string& datasetPath,
          const HeightToBuildingVec& heightToBuildingVec,
          const HeightToRequirement& heightToRequirement,
          double gapSize, size_t zoom,
          size_t threadsNum) {
    HeightToLoadedCount heightToLoadedCount;
    for (const auto& [height, requirement] : heightToRequirement) {
        heightToLoadedCount[height] = 0;
    }

    TAutoPtr<IThreadPool> mtpQueue = CreateThreadPool(threadsNum + 1);

    BldQueue inpQueue;
    BldWithDataQueue outQueue;

    std::vector<TAutoPtr<DataDownloader>> downloaders;
    for (size_t i = 0; i < threadsNum; i++) {
        downloaders.emplace_back(
            new DataDownloader(
                &inpQueue, &outQueue,
                gapSize, zoom,
                heightToRequirement, &heightToLoadedCount)
        );
    }
    TAutoPtr<DataWriter> writer = new DataWriter(
        &outQueue,
        datasetPath,
        heightToRequirement
    );

    mtpQueue->SafeAdd(writer.Get());
    for (const auto& downloader : downloaders) {
        mtpQueue->SafeAdd(downloader.Get());
    }

    readBldsToQueue(
        heightToBuildingVec,
        &inpQueue,
        heightToRequirement,
        heightToLoadedCount
    );

    for (const auto& downloader : downloaders) {
        downloader->DataEnded();
    }

    while (!inpQueue.IsEmpty()) {
        wait();
    }

    bool waitDownloader = true;
    while (waitDownloader) {
        wait();
        waitDownloader = false;
        for (size_t i = 0; i < downloaders.size(); i++) {
            waitDownloader |= downloaders[i]->IsRunning();
        }
    }

    writer->DataEnded();

    while (writer->IsRunning()) {
        wait();
    }

    mtpQueue->Stop();
}

} // namespace

int main(int argc, const char** argv)
try {
    maps::cmdline::Parser parser;

    auto ymapsdfPath = parser.string("ymapsdf_path")
        .required()
        .help("Path to folder with YMapsDF tables");

    auto requirementsPath = parser.string("requirements_path")
        .required()
        .help("Path to file with heights and requirements");

    auto datasetPath = parser.string("dataset_path")
        .required()
        .help("Path to dataset");

    auto gapSize = parser.real("gap")
        .required()
        .help("Gap size");

    auto zoom = parser.num("zoom")
        .required()
        .help("Satellite image zoom");

    auto threadsNum = parser.size_t("threads_num")
        .required()
        .help("Threads number");

    parser.parse(argc, const_cast<char**>(argv));

    INFO() << "Loading dataset requirements: " << requirementsPath;
    HeightToRequirement heightToRequirement
        = loadRequirements(requirementsPath);
    std::set<size_t> heights;
    INFO() << "Requirements:";
    for (const auto& [height, requirement] : heightToRequirement) {
        heights.emplace(height);
        INFO() << "  Height: " << height << ", requirement: " << requirement;
    }

    INFO() << "Loading all buildings with suitable height";
    IdToHeight idToHeight = loadBldsWithSuitableHeight(ymapsdfPath, heights);

    INFO() << "Loading buildings polygon";
    HeightToBuildingVec heightToBuildingVec
        = loadBldsGeometry(ymapsdfPath, idToHeight);

    INFO() << "Saving dataset: " << datasetPath;
    dump(datasetPath,
        heightToBuildingVec,
        heightToRequirement,
        gapSize,
        zoom,
        threadsNum);

    INFO() << "Success!";

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    INFO() << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    INFO() << e.what();
    return EXIT_FAILURE;
}
catch (...) {
    INFO() << "Caught unknown exception";
    return EXIT_FAILURE;
}
