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

static const std::string keys =
"{ @left-image-format  | <none> | format string for path to input left image  }"
"{ @right-image-format | <none> | format string for path to input right image }"
"{ disp_size           |    256 | maximum possible disparity value            }"
"{ P1                  |     10 | penalty on the disparity change by plus or minus 1 }"
"{ P2                  |    120 | penalty on the disparity change by more than 1 }"
"{ uniqueness          |   0.80 | margin in ratio for best cost function value }"
"{ num_paths           |      8 | number of scanlines used in cost aggregation }"
"{ min_disp            |   -160 | minimum disparity value                     }"
"{ LR_max_diff         |      1 | max allowed difference between L/R disparity }"
"{ census_type         |      1 | type of census transform                    }"
"{ interval            |      1 | polling interval in seconds                 }"
"{ help h              |        | display this help and exit                  }";

class ImagePreprocessor {
public:
    void preprocess_image_pair(cv::Mat& img_left, cv::Mat& img_right) {
        if (img_left.channels() > 1) cv::cvtColor(img_left, img_left, cv::COLOR_BGR2GRAY);
        if (img_right.channels() > 1) cv::cvtColor(img_right, img_right, cv::COLOR_BGR2GRAY);

        int min_height = std::min(img_left.rows, img_right.rows);
        int min_width = std::min(img_left.cols, img_right.cols);
        img_left = img_left(cv::Rect(0, 0, min_width, min_height));
        img_right = img_right(cv::Rect(0, 0, min_width, min_height));
    }
};

bool disparityAlreadyProcessed(int frame_no) {
    std::string xml_path = cv::format("output/disparity_%04d.xml", frame_no);
    return cv::utils::fs::exists(xml_path);
}

int main(int argc, char* argv[]) {
    cv::CommandLineParser parser(argc, argv, keys);
    if (parser.has("help")) {
        parser.printMessage();
        return 0;
    }

    const std::string format_L = parser.get<cv::String>("@left-image-format");
    const std::string format_R = parser.get<cv::String>("@right-image-format");
    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 int interval = parser.get<int>("interval");
    const auto census_type = static_cast<sgm::CensusType>(parser.get<int>("census_type"));

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

    if (!cv::utils::fs::exists("output")) {
        cv::utils::fs::createDirectory("output");
    }

    ImagePreprocessor preprocessor;
    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);

    int last_checked = 0;
    while (true) {
        const std::string left_path = cv::format(format_L.c_str(), last_checked);
        const std::string right_path = cv::format(format_R.c_str(), last_checked);

        if (cv::utils::fs::exists(left_path) && cv::utils::fs::exists(right_path) && !disparityAlreadyProcessed(last_checked)) {
            cv::TickMeter timer;
            timer.start();
            std::cout << "Processing frame " << last_checked;

            cv::Mat I1 = cv::imread(left_path, cv::IMREAD_UNCHANGED);
            cv::Mat I2 = cv::imread(right_path, cv::IMREAD_UNCHANGED);

            if (I1.empty() || I2.empty()) {
                std::cerr << "Error reading images." << std::endl;
                break;
            }

            preprocessor.preprocess_image_pair(I1, I2);

            ASSERT_MSG(I1.size() == I2.size() && I1.type() == I2.type(), "Mismatched image size/type.");
            ASSERT_MSG(I1.type() == CV_8U || I1.type() == CV_16U, "Images must be CV_8U or CV_16U.");

            const int src_depth = I1.type() == CV_8U ? 8 : 16;
            const int dst_depth = 16;

            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);
            ssgm.execute(I1.data, I2.data, disparity.data);

            cv::FileStorage fs(cv::format("output/disparity_%04d.xml", last_checked), cv::FileStorage::WRITE);
            fs << "disparity" << disparity;
            fs.release();

            timer.stop();
            std::cout << " - " << timer.getTimeSec() << " seconds" << std::endl;
        }

        last_checked++;
        cv::waitKey(interval * 1000);  // Sleep for polling interval
    }

    return 0;
}