#include "graph_pgn_extract.hpp"
#include "pp_pgns.hpp"
#include "utils.hpp"

#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"

#include <iostream>
#include <fstream>
#include <vector>
#include <list>
#include <set>
#include <algorithm>
#include <string>


using namespace maps::wiki::autocart;
static const cv::Scalar colors[18] =
{
    cv::Scalar(255, 0, 0),
    cv::Scalar(0, 255, 0),
    cv::Scalar(0, 0, 255),
    cv::Scalar(255, 255, 0),
    cv::Scalar(255, 0, 255),
    cv::Scalar(0, 255, 255),
    cv::Scalar(128, 0, 0),
    cv::Scalar(0, 128, 0),
    cv::Scalar(0, 0, 128),
    cv::Scalar(128, 128, 0),
    cv::Scalar(128, 0, 128),
    cv::Scalar(0, 128, 128),
    cv::Scalar(128, 255, 0),
    cv::Scalar(255, 128, 0),
    cv::Scalar(128, 0, 255),
    cv::Scalar(255, 0, 128),
    cv::Scalar(0, 128, 255),
    cv::Scalar(0, 255, 128),
};

static std::vector< std::vector<cv::Point> > testPolygons({
    std::vector<cv::Point>({
        cv::Point(19, 17),
        cv::Point(159, 29),
        cv::Point(162, 87),
        cv::Point(256, 87),
        cv::Point(249, 152),
        cv::Point(359, 153),
        cv::Point(347, 311),
        cv::Point(25, 305)}),
    std::vector<cv::Point>({
        cv::Point(52, 17),
        cv::Point(244, 13),
        cv::Point(241, 301),
        cv::Point(211, 298),
        cv::Point(207, 341),
        cv::Point(48, 341),
        cv::Point(50, 176),
        cv::Point(12, 176),
        cv::Point(12, 70),
        cv::Point(52, 65)}),
    std::vector<cv::Point>({
        cv::Point(316, 72),
        cv::Point(176, 28),
        cv::Point(84, 316),
        cv::Point(240, 372),
        cv::Point(284, 244),
        cv::Point(248, 216)}),
    std::vector<cv::Point>({
        cv::Point(117, 90),
        cv::Point(162, 155),
        cv::Point(92, 220),
        cv::Point(147, 370),
        cv::Point(302, 280),
        cv::Point(237, 205),
        cv::Point(307, 150),
        cv::Point(227, 30)}),
    std::vector<cv::Point>({
        cv::Point(10, 107),
        cv::Point(65, 212),
        cv::Point(105, 177),
        cv::Point(265, 362),
        cv::Point(390, 302),
        cv::Point(240, 102),
        cv::Point(200, 142),
        cv::Point(160, 37)}),
    std::vector<cv::Point>({
        cv::Point(56, 113),
        cv::Point(158, 275),
        cv::Point(248, 233),
        cv::Point(296, 359),
        cv::Point(344, 311),
        cv::Point(194, 41)}),
});

static const std::string filenames[] =
{
    "z18__29.010_41.070__29.018_41.075",
    "z18__29.010_41.080__29.018_41.085",
    "z18__29.018_41.075__29.026_41.080",
    "z18__29.026_41.070__29.034_41.075",
    "z18__29.026_41.080__29.034_41.085",
    "z18__29.010_41.075__29.018_41.080",
    "z18__29.018_41.070__29.026_41.075",
    "z18__29.018_41.080__29.026_41.085",
    "z18__29.026_41.075__29.034_41.080"
};

static bool loadPolygons(const std::string &filepath, std::vector< std::vector<cv::Point> > &polygons)
{
    std::ifstream ifs(filepath);
    if (!ifs.is_open())
        return false;

    for (; !ifs.eof();)
    {
        std::string line; std::getline(ifs, line);
        if (line.empty())
            continue;
        if ('#' == line[0])
            continue;
        std::stringstream ss(line);

        int pgn_cnt;
        ss >> pgn_cnt;
        if (0 == pgn_cnt)
            continue;

        std::vector<cv::Point> pgn(pgn_cnt);
        for (int i = 0; i < pgn_cnt; i++)
        {
            ss >> pgn[i].x >> pgn[i].y;
        }

        polygons.emplace_back(pgn);
    }
    return true;
}
static void savePolygons(const std::string &filepath, const std::vector< std::vector<cv::Point> > &polygons)
{
    std::ofstream ofs(filepath);
    for (size_t i = 0; i < polygons.size(); i++)
    {
        const std::vector<cv::Point> &pgn = polygons[i];
        if (0 == pgn.size())
            continue;
        ofs << pgn.size() << " ";
        for (size_t j = 0; j < pgn.size(); j++)
        {
            ofs << pgn[j].x << " " << pgn[j].y << " ";
        }
        ofs << std::endl;
    }
}

static cv::Scalar getColor(int idx, bool small)
{
    static cv::Mat colorsMap;
    if (colorsMap.empty())
    {
        colorsMap.create(256, 1, CV_8UC1);
        for (int i = 0; i < 256; i++)
            colorsMap.at<uchar>(i, 0) = (uchar)i;
        cv::applyColorMap(colorsMap, colorsMap, cv::COLORMAP_JET);
    }
    if (small)
        return colors[idx % 18];
    return colorsMap.at<cv::Vec3b>(idx % 256, 0);

}
static void drawPolygons(cv::Mat &img_show, const std::vector< std::vector<cv::Point> > &polygons, const std::set<size_t> &indices)
{
    bool smallColorMap = (polygons.size() <= 18);
    //bool smallColorMap = true;
    for (size_t i = 0; i < polygons.size(); i++)
    {
        if (!indices.empty() && (indices.find(i) == indices.end()))
            continue;
        const std::vector<cv::Point> &coords = polygons[i];
        //if (coords.size() <= 5)
        //    continue;
        cv::Scalar color = getColor((int)i, smallColorMap);
        for (size_t j = 0; j < coords.size() - 1; ++j)
        {
            cv::line(img_show, coords[j], coords[j + 1], color, 2);
        }
        for (size_t j = 0; j < coords.size() - 1; ++j)
        {
            cv::circle(img_show, coords[j], 1, color, 3);
        }
    }
}

static void runExtractPolygons(const std::string &filename, bool show = false)
{
    std::cout << filename << std::endl;

    static const maps::wiki::autocart::ExtractPolygonsParams ep_params;
    static const maps::wiki::autocart::PPPolygonsParams pp_params;

    const std::string folder_path = "C:/myfolder/work/TensorFlow/bld_edges_verts/temp/hed_512x512_edges_verts_sz5/ew1_vw10_lossvext_ep80/";
    const std::string img_path = folder_path + filename + ".jpg";
    const std::string verts_path = folder_path + filename + "_verts.png";
    const std::string edges_path = folder_path + filename + "_edges.png";
    const std::string pgn_path = folder_path + filename + "_pgn.txt";
    const std::string pgn_img_path = folder_path + filename + "_pgn.png";
    const std::string pgn_pp_img_path = folder_path + filename + "_pgn_pp.png";

    const std::string sem_seg_path = "C:/myfolder/work/TensorFlow/bld_edges_verts/temp/unet_finest_512x512_IoU_eq_hist_markup_refinement_150/" + filename + ".png";
    cv::Mat img = cv::imread(img_path, cv::IMREAD_COLOR);
    cv::Mat img_verts = cv::imread(verts_path, cv::IMREAD_GRAYSCALE);
    cv::Mat img_edges = cv::imread(edges_path, cv::IMREAD_GRAYSCALE);
    cv::Mat img_segm = cv::imread(sem_seg_path, cv::IMREAD_GRAYSCALE);

    cv::Rect roi(0, 0, img.cols, img.rows);
    //if (false)
    //    roi = cv::Rect(524, 372, 250, 250);
    //if (false)
    //    roi = cv::Rect(450, 0, 400, 450);
    //if (false)
    //    roi = cv::Rect(550, 100, 125, 175);
    //if (false)
    //    roi = cv::Rect(900, 380, 500, 400);
    //if (false)
    //    roi = cv::Rect(450, 600, 250, 250);
    cv::Mat img_roi; img(roi).copyTo(img_roi);
    cv::Mat img_verts_roi; img_verts(roi).copyTo(img_verts_roi);
    cv::Mat img_edges_roi; img_edges(roi).copyTo(img_edges_roi);
    cv::Mat img_segm_roi; img_segm(roi).copyTo(img_segm_roi);

    if (false)
    {
        cv::imshow("vert", img_verts_roi);
        cv::imshow("edges", img_edges_roi);
        cv::imshow("segm", img_segm_roi * 255);
        cv::Mat merged;
        cv::merge(std::vector<cv::Mat>({ img_verts_roi, img_edges_roi, cv::Mat::zeros(roi.size(), CV_8UC1) }), merged);
        cv::imshow("merged", merged);
    }

    std::vector< std::vector<cv::Point> > polygons;
    //if (true)
    if (false)
    {
        maps::wiki::autocart::extractPolygons(img_verts_roi, img_edges_roi, img_segm_roi, ep_params, polygons);
        if (false)
        {
            savePolygons(pgn_path, polygons);
            return;
        }
    }
    else
    {
        loadPolygons(pgn_path, polygons);
    }

    if (show)
    {
        cv::Mat img_show; img_roi.copyTo(img_show);
        drawPolygons(img_show, polygons, std::set<size_t>());
        cv::imshow("polygons", img_show);
    }


    maps::wiki::autocart::postprocessPolygons(pp_params, polygons);

    if (show)
    {
        cv::Mat img_show_pp; img_roi.copyTo(img_show_pp);
        drawPolygons(img_show_pp, polygons, std::set<size_t>());
        cv::imshow("polygons_pp", img_show_pp);
        cv::waitKey();
    }
    std::vector< std::vector<cv::Point> > bigpgn;
    for (size_t i = 0; i < polygons.size(); i++)
    {
        if (polygons[i].size() > 5)
            bigpgn.push_back(polygons[i]);
    }
    savePolygons("c:/temp/pgn1.txt", bigpgn);
}

int main(int argc, const char** argv)
{
#if 0
    for (const auto &pgn : testPolygons)
    {
        const size_t  pgn_cnt = pgn.size();
        cv::Mat img(400, 400, CV_8UC3);
        for (size_t i = 0; i < pgn_cnt; i++)
        {
            cv::line(img, pgn[i], pgn[(i + 1) % pgn_cnt], cv::Scalar(0, 0, 0));
        }
        cv::imshow("test", img);
        cv::waitKey(0);
    }
#endif
#if 1
    static const std::string filename = "z18__29.018_41.070__29.026_41.075";
    //static const std::string filename = "z18__29.018_41.080__29.026_41.085";
    //static const std::string file_name = "z18__29.026_41.070__29.034_41.075";
    runExtractPolygons(filename, true);
#else
    for (size_t i = 0; i < sizeof(filenames) / sizeof(std::string); i++)
    {
        runExtractPolygons(filenames[i], true);
    }
    return 0;
#endif
}
