#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <libsgm.h>
#include "sample_common.h"

static const std::string keys =
"{ @left_img   | <none> | path to input left image                                                            }"
"{ @right_img  | <none> | path to input right image                                                           }"
"{ disp_size   |    256 | maximum possible disparity value                                                    }"
"{ P1          |     10 | penalty on the disparity change by plus or minus 1 between neighbor pixels          }"
"{ P2          |    120 | penalty on the disparity change by more than 1 between neighbor pixels              }"
"{ uniqueness  |   0.80 | margin in ratio by which the best cost function value should be at least second one }"
"{ num_paths   |      8 | number of scanlines used in cost aggregation                                        }"
"{ min_disp    |   -160 | minimum disparity value                                                             }"
"{ LR_max_diff |      1 | maximum allowed difference between left and right disparity                         }"
"{ census_type |      1 | type of census transform (0:CENSUS_9x7 1:SYMMETRIC_CENSUS_9x7)                      }"
"{ output_dir  |      . | directory to save disparity.xml and disparity_color.png                             }"
"{ no_display  |      0 | set to 1 to skip interactive display window (for pipeline/headless use)             }"
"{ help h      |        | display this help and exit                                                          }";

int main(int argc, char* argv[])
{
    double start_time = cv::getTickCount();  // Start total execution time

    cv::CommandLineParser parser(argc, argv, keys);
    if (parser.has("help")) {
        parser.printMessage();
        return 0;
    }

    double load_start = cv::getTickCount(); // Start loading time
    cv::Mat I1 = cv::imread(parser.get<cv::String>("@left_img"), cv::IMREAD_UNCHANGED);
    cv::Mat I2 = cv::imread(parser.get<cv::String>("@right_img"), cv::IMREAD_UNCHANGED);
    double load_end = cv::getTickCount();
    double load_time_s = (load_end - load_start) / cv::getTickFrequency();  // Seconds
    double load_time_ms = load_time_s * 1000.0;  // Milliseconds
    std::cout << "Image Loading Time: " << load_time_s << " s (" << load_time_ms << " ms)" << std::endl;

    if (I1.channels() > 1) cv::cvtColor(I1, I1, cv::COLOR_BGR2GRAY);
    if (I2.channels() > 1) cv::cvtColor(I2, I2, cv::COLOR_BGR2GRAY);

    const int disp_size = parser.get<int>("disp_size");
    const int P1 = parser.get<int>("P1");
    const int P2 = parser.get<int>("P2");
    const float uniqueness = parser.get<float>("uniqueness");
    const int num_paths = parser.get<int>("num_paths");
    const int min_disp = parser.get<int>("min_disp");
    const int LR_max_diff = parser.get<int>("LR_max_diff");
    const auto census_type = static_cast<sgm::CensusType>(parser.get<int>("census_type"));

    if (!parser.check()) {
        parser.printErrors();
        parser.printMessage();
        std::exit(EXIT_FAILURE);
    }

    ASSERT_MSG(!I1.empty() && !I2.empty(), "imread failed.");
    ASSERT_MSG(I1.size() == I2.size() && I1.type() == I2.type(), "input images must be same size and type.");
    ASSERT_MSG(I1.type() == CV_8U || I1.type() == CV_16U, "input image format must be CV_8U or CV_16U.");
    ASSERT_MSG(disp_size == 64 || disp_size == 128 || disp_size == 256, "disparity size must be 64, 128 or 256.");
    ASSERT_MSG(num_paths == 4 || num_paths == 8, "number of scanlines must be 4 or 8.");
    ASSERT_MSG(census_type == sgm::CensusType::CENSUS_9x7 || census_type == sgm::CensusType::SYMMETRIC_CENSUS_9x7, "census type must be 0 or 1.");

    const int src_depth = I1.type() == CV_8U ? 8 : 16;
    const int dst_depth = 16;
    const sgm::PathType path_type = num_paths == 8 ? sgm::PathType::SCAN_8PATH : sgm::PathType::SCAN_4PATH;

    const sgm::StereoSGM::Parameters param(P1, P2, uniqueness, false, path_type, min_disp, LR_max_diff, census_type);
    sgm::StereoSGM ssgm(I1.cols, I1.rows, disp_size, src_depth, dst_depth, sgm::EXECUTE_INOUT_HOST2HOST, param);

    cv::Mat disparity(I1.size(), CV_16S);

    double disparity_start = cv::getTickCount(); // Start disparity computation time
    ssgm.execute(I1.data, I2.data, disparity.data);
    double disparity_end = cv::getTickCount();
    double disparity_time_s = (disparity_end - disparity_start) / cv::getTickFrequency();  // Seconds
    double disparity_time_ms = disparity_time_s * 1000.0;  // Milliseconds
    std::cout << "Disparity Computation Time: " << disparity_time_s << " s (" << disparity_time_ms << " ms)" << std::endl;

    const std::string output_dir = parser.get<std::string>("output_dir");

    // Save disparity
    cv::FileStorage fs(output_dir + "/disparity.xml", cv::FileStorage::WRITE);
    fs << "disparity" << disparity;
    fs.release();

    // Convert disparity to 8-bit for visualization
    cv::Mat disparity_8u, disparity_color;
    disparity.convertTo(disparity_8u, CV_8U, 255.0 / disp_size);
    cv::applyColorMap(disparity_8u, disparity_color, cv::COLORMAP_TURBO);

    // Save colored disparity image
    cv::imwrite(output_dir + "/disparity_color.png", disparity_color);

    double total_end = cv::getTickCount();
    double total_time_s = (total_end - start_time) / cv::getTickFrequency();  // Seconds
    double total_time_ms = total_time_s * 1000.0;  // Milliseconds
    std::cout << "Total Execution Time: " << total_time_s << " s (" << total_time_ms << " ms)" << std::endl;

    // Display images
    const std::vector<cv::Mat> images = { disparity_8u, disparity_color, I1 };
    const std::vector<std::string> titles = { "Disparity", "Colored Disparity", "Input Image" };

    if (!parser.get<int>("no_display")) {
        std::cout << "Hot keys:\n";
        std::cout << "\tESC - Quit the program\n";
        std::cout << "\ts - Switch display (Disparity | Colored Disparity | Input Image)\n";

        int mode = 0;
        while (true) {
            cv::setWindowTitle("Image", titles[mode]);
            cv::imshow("Image", images[mode]);

            const char c = cv::waitKey(0);
            if (c == 's') mode = (mode < 2 ? mode + 1 : 0);
            if (c == 27) break;
        }
    }

    return 0;
}