#include "detect_road.h"
#include <maps/wikimap/mapspro/libs/tf_inferencer/tf_inferencer.h>
#include <maps/wikimap/mapspro/services/autocart/libs/post_processing/include/post_processing.h>


#include <maps/libs/geolib/include/polygon.h>
#include <maps/libs/geolib/include/spatial_relation.h>
#include <maps/libs/geolib/include/intersection.h>
#include <maps/libs/geolib/include/vector.h>
#include <maps/libs/geolib/include/distance.h>
#include <maps/libs/geolib/include/convex_hull.h>
#include <maps/libs/geolib/include/segment.h>
#include <maps/libs/geolib/include/line.h>
#include <maps/libs/geolib/include/direction.h>
#include <maps/libs/geolib/include/bounding_box.h>
#include <maps/libs/geolib/include/static_geometry_searcher.h>

#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.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 <util/generic/string.h>
#include <util/system/env.h>
#include <stdlib.h>

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

void drawPolylines(cv::Mat* markup,
                   const std::vector<geolib3::Polyline2>& lines,
                   const cv::Mat& canvas) {
    const size_t THICKNESS = 5;
    const cv::Scalar COLOR = cv::Scalar(0, 0, 255);
    canvas.copyTo(*markup);
    for (const auto& line : lines) {
        cv::Scalar color(rand() % 256, rand() % 256, rand() % 256);
        for (size_t i = 0; i < line.pointsNumber() - 1; i++) {
            geolib3::Point2 pt0 = line.pointAt(i);
            geolib3::Point2 pt1 = line.pointAt(i + 1);
            cv::Point2d cvPt0(pt0.x(), pt0.y());
            cv::Point2d cvPt1(pt1.x(), pt1.y());
            cv::line(*markup, cvPt0, cvPt1, COLOR, THICKNESS);
        }
    }
}

void savePolylines(const std::string& path,
                   const std::vector<geolib3::Polyline2>& lines) {
    std::ofstream ofs(path);

    REQUIRE(ofs.is_open(), "Failed to open file: " << path);

    for (const auto& line : lines) {
        ofs << "road " << line.pointsNumber();
        for (size_t i = 0; i < line.pointsNumber(); ++i) {
            geolib3::Point2 pt = line.pointAt(i);
            ofs << " " << pt.x() << " " << pt.y();
        }
        ofs << "\n";
    }
    ofs.close();
}

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 modelPathParam = parser.string("model_path")\
        .required()\
        .help("path to the gdef file for road detection network");

    auto markupPathParam = parser.string("markup")\
        .required()\
        .help("path to txt file with detected roads");

    auto markupImgPathParam = parser.string("markup_image")\
        .required()\
        .help("path to image with result markup");

    parser.parse(argc, argv);

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

    INFO() << "loading road detection model " << modelPathParam;
    maps::wiki::tf_inferencer::TensorFlowInferencer road_detector(modelPathParam);

    auto roads = autocart::detectRoad(road_detector, image);
    INFO() << "Detect " << roads.size() << " roads";

    cv::Mat markup;
    drawPolylines(&markup, roads, image);
    cv::imwrite(markupImgPathParam, markup);
    savePolylines(markupPathParam, roads);
    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;
}
