#include <maps/wikimap/mapspro/libs/tf_inferencer/tf_inferencer.h>
#include <maps/wikimap/mapspro/services/autocart/libs/detection/include/detect_bld.h>
#include <maps/wikimap/mapspro/services/autocart/libs/post_processing/include/post_processing.h>
#include <maps/wikimap/mapspro/services/autocart/libs/post_processing/include/align_along_road.h>

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

#include <opencv2/opencv.hpp>

#include <tensorflow/core/public/session.h>
#include <tensorflow/core/framework/graph.pb.h>
#include <tensorflow/core/framework/tensor.h>
#include <tensorflow/core/common_runtime/dma_helper.h>

#include <vector>
#include <fstream>
#include <functional>
#include <unordered_map>

#include <util/generic/string.h>
#include <util/system/env.h>

namespace autocart = maps::wiki::autocart;
namespace tf_inferencer = maps::wiki::tf_inferencer;
namespace geolib3 = maps::geolib3;

namespace {

using BldRecognitionFunc =
    std::function<
        std::vector<autocart::Polygon>
            (const std::string&,
             const std::string&,
             const cv::Mat&)>;

std::vector<autocart::Polygon>
bldRecognitionByRectangle(const std::string& semSegmPath,
                          const std::string& edgeDetectPath,
                          const cv::Mat& image)
{
    REQUIRE(!semSegmPath.empty() && !edgeDetectPath.empty(),
            "sem_segm_path and edge_detector_path parameters "
            "must be specified for \"rectangle\" method");

    INFO() << "Loading semantic segmentation model: " << semSegmPath;
    tf_inferencer::TensorFlowInferencer segmentator(semSegmPath);

    INFO() << "Loading edge detection model: " << edgeDetectPath;
    tf_inferencer::TensorFlowInferencer edgeDetector(edgeDetectPath);

    INFO() << "Detect building";
    auto polygons = autocart::detectBldByMinAreaRect(segmentator, edgeDetector, image);
    return polygons;
}

std::vector<autocart::Polygon>
bldRecognitionByEdges(const std::string& edgeDetectPath, const cv::Mat& image)
{
    REQUIRE(!edgeDetectPath.empty(),
            "edge_detector_path parameter must be specified for \"edge\" method");

    INFO() << "Loading edge detection model: " << edgeDetectPath;
    tf_inferencer::TensorFlowInferencer edgeDetector(edgeDetectPath);

    INFO() << "Detect buildings";
    auto polygons = autocart::detectBldByEdges(edgeDetector, image);
    return polygons;
}

std::vector<autocart::Polygon>
bldRecognitionByVertsEdges(const std::string& semSegmPath,
                           const std::string& edgeVertDetectPath,
                           const cv::Mat& image)
{
    REQUIRE(!semSegmPath.empty() && !edgeVertDetectPath.empty(),
            "sem_segm_path and edge_vert_detector_path parameters "
            "must be specified for \"edge_vert\" method");

    INFO() << "Loading semantic segmentation model: " << semSegmPath;
    tf_inferencer::TensorFlowInferencer segmentator(semSegmPath);

    INFO() << "Loading edge and vert detection model: " << edgeVertDetectPath;
    tf_inferencer::TensorFlowInferencer edgeVertDetector(edgeVertDetectPath);

    INFO() << "Detect buildings";
    auto polygons = autocart::detectBldByVertsEdges(segmentator, edgeVertDetector, image);
    return polygons;
}

const std::unordered_map<std::string, BldRecognitionFunc>
    RECOGNITION_MODE_TO_FUNC  = {
        {std::string("rectangle"), bldRecognitionByRectangle},
        {std::string("edge_vert"), bldRecognitionByVertsEdges},
        {
            std::string("edge"),
            [](const std::string& /*semSegmPath*/,
               const std::string& edgeDetectPath,
               const cv::Mat& image)
            {
                return bldRecognitionByEdges(edgeDetectPath, image);
            }
        }
    };


std::vector<geolib3::Polygon2>
convertToGeolibPolygons(const std::vector<autocart::Polygon> &polygons)
{
    std::vector<geolib3::Polygon2> result;
    result.reserve(polygons.size());

    for (const auto &polygon : polygons) {
        std::vector<geolib3::Point2> ring;
        ring.reserve(polygon.size());
        for (const auto & point : polygon) {
            ring.emplace_back(point.x, point.y);
        }
        result.emplace_back(std::move(ring));
    }
    return result;
}

std::vector<geolib3::Polygon2>
rectifyAndAlignPolygons(const std::vector<geolib3::Polygon2>& polygons)
{
    std::vector<geolib3::Polygon2> boundingBoxes;
    boundingBoxes.reserve(polygons.size());
    std::transform(
        polygons.begin(), polygons.end(),
        std::back_inserter(boundingBoxes),
        [](const geolib3::Polygon2& polygon)
        {
            return autocart::getBoundingRectangle(polygon);
        }
    );

    std::vector<geolib3::Polygon2> result = autocart::alignRectangles(boundingBoxes, 3, 10, 200, 5);
    return result;
}

std::vector<geolib3::Polyline2> loadRoads(const std::string &roadsFilePath)
{
    std::ifstream ifs(roadsFilePath);
    if (!ifs.is_open()) {
        throw maps::RuntimeError() << "Failed to open file: " << roadsFilePath;
    }

    std::vector<geolib3::Polyline2> roads;
    while (!ifs.eof()) {
        std::string line;
        std::getline(ifs, line);
        if (line.empty()) {
            continue;
        }
        std::stringstream ss(line);
        std::string className;
        ss >> className;
        if ("road" != className) {
            continue;
        }
        size_t ptsCnt;
        ss >> ptsCnt;
        std::vector<geolib3::Point2> road;
        road.reserve(ptsCnt);
        for (size_t i = 0; i < ptsCnt; i++) {
            double x, y;
            ss >> x >> y;
            road.emplace_back(x, y);
        }
        roads.emplace_back(road);
    }
    return roads;
}

void saveBuildings(const std::string &outFilePath,
                   const std::vector<geolib3::Polygon2> &polygons,
                   const std::string &name = "bld")
{
    std::ofstream ofs(outFilePath);
    for (const auto &polygon : polygons) {
        size_t ptsCnt = polygon.pointsNumber();
        ofs << name << " " << ptsCnt << " ";
        for (size_t i = 0; i < ptsCnt; i++) {
            geolib3::Point2 pt = polygon.pointAt(i);
            ofs << pt.x() << " " << pt.y() << " ";
        }
        ofs << std::endl;
    }
}

void drawBuilding(cv::Mat &img, const geolib3::Polygon2 &polygon, const cv::Scalar &color)
{
    size_t ptsCnt = polygon.pointsNumber();
    for (size_t i = 0; i < ptsCnt; i++)
    {
        geolib3::Point2 pt0 = polygon.pointAt(i);
        geolib3::Point2 pt1 = polygon.pointAt((i + 1) % ptsCnt);

        cv::Point2f cvPt0(pt0.x(), pt0.y());
        cv::Point2f cvPt1(pt1.x(), pt1.y());

        cv::line(img, cvPt0, cvPt1, color, 3);
    }
}

void drawBuildings(const cv::Mat &img, const std::vector<geolib3::Polygon2> &polygons,
                   cv::Mat &markup, int colormap = cv::COLORMAP_JET)
{
    img.copyTo(markup);
    cv::Mat colors(256, 1, CV_8UC1);
    for (int i = 0; i < 256; i++)
        colors.at<uchar>(i, 0) = (uchar)i;
    cv::applyColorMap(colors, colors, colormap);

    for (size_t i = 0; i < polygons.size(); i++)
    {
        drawBuilding(markup, polygons[i], colors.at<cv::Vec3b>(i % colors.rows, 0));
    }
}

} //namespace


int main(int argc, char** argv)
try {
    maps::cmdline::Parser parser;
    auto imgPathParam = parser.string("img_path")\
            .required()\
            .help("path to the satellite image file");

    auto modeParam = parser.string("mode")\
            .defaultValue("rectangle")\
            .help("algorithm for building vectorization:\n"
                  "    rectangle - detect by minimum area rectangle\n"
                  "    edge - detect using only building edges\n"
                  "    edge_vert - detect using building edges and vertices");

    auto semSegmPathParam = parser.string("sem_segm_path")\
            .defaultValue("")\
            .help("path to the gdef file for semantic segmentation net");

    auto edgeDetectPathParam = parser.string("edge_detector_path")\
            .defaultValue("")\
            .help("path to the gdef file for edge detector net");

    auto edgeVertDetectPathParam = parser.string("edge_vert_detector_path")\
            .defaultValue("")\
            .help("path to the gdef file for edge and vertex detector net");

    auto roadsFilePathParam = parser.string("road_path")\
            .defaultValue("")\
            .help("path to txt file with roads polylines");

    auto outMarkupPathParam = parser.string("output_path")\
            .required()\
            .help("path to the output file with building rectangles");

    auto dumpMarkupImageParam = parser.string("dump_markup")\
            .defaultValue("")\
            .help("Dump image with rectangles markup on the source satellite image");

    parser.parse(argc, argv);

    INFO() << "Loading image: " << imgPathParam;
    cv::Mat image = cv::imread(imgPathParam);
    REQUIRE(!image.empty(), "Unable to load image from file: " << imgPathParam);

    if (!RECOGNITION_MODE_TO_FUNC.count(modeParam)) {
        REQUIRE(false, "Unknown mode: " << modeParam);
    }

    INFO() << "Using mode: " << modeParam;
    std::vector<autocart::Polygon> polygons =
            RECOGNITION_MODE_TO_FUNC.at(modeParam)(semSegmPathParam,
                                                   ("edge_vert" != modeParam) ? edgeDetectPathParam : edgeVertDetectPathParam,
                                                   image);

    INFO() << "Convert OpenCV polygons to geolib polygons";
    auto geolibPolygons = convertToGeolibPolygons(polygons);

    INFO() << "Rectifying and aligning polygons";
    auto alignedPolygons = rectifyAndAlignPolygons(geolibPolygons);

    INFO() << "Removing intersections between detected buildings";
    autocart::removeIntersections(alignedPolygons);

    std::vector<geolib3::Polyline2> roadExists;
    if (!roadsFilePathParam.empty()) {
        INFO() << "Loading existing roads";
        roadExists = loadRoads(roadsFilePathParam);
    }

    INFO() << "Aligning along existing roads";
    autocart::AlignAlongRoadsParams alignAlongRoadParams;
    auto resultPolygons = autocart::alignAlongRoads(alignedPolygons,
                                                    roadExists,
                                                    alignAlongRoadParams);
    INFO() << "Saving results";
    saveBuildings(outMarkupPathParam, resultPolygons);

    if (!dumpMarkupImageParam.empty()) {
        cv::Mat markup;
        drawBuildings(image, resultPolygons, markup);
        cv::imwrite(dumpMarkupImageParam, markup);
    }

    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;
}
