/*
Copyright 2016 Fixstars Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http ://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include <iostream>
#include <chrono>
#include <stdexcept>

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <fstream> // Add this line to use std::ofstream for file output

#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            }"
"{ start_number        |      0 | index to start reading                      }"
"{ help h              |        | display this help and exit                  }";

class ImagePreprocessor {
public:
    void preprocess_image_pair(cv::Mat& img_left, cv::Mat& img_right) {
        // Get the shape of both images
        int h1 = img_left.rows, w1 = img_left.cols;
        int h2 = img_right.rows, w2 = img_right.cols;

        // Find the minimum height and width between the two images
        int min_height = std::min(h1, h2);
        int min_width = std::min(w1, w2);

        // Crop both images to match the minimum height and width
        img_left = img_left(cv::Rect(0, 0, min_width, min_height));
        img_right = img_right(cv::Rect(0, 0, min_width, min_height));

        // Convert to CV_8U grayscale
        //cv::cvtColor(img_left, img_left, cv::COLOR_BGR2GRAY);
        img_left.convertTo(img_left, CV_8U); // Ensure it's in CV_8U format
        //cv::cvtColor(img_right, img_right, cv::COLOR_BGR2GRAY);
        img_right.convertTo(img_right, CV_8U); // Ensure it's in CV_8U format
    }
};

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

    const std::string image_format_L = parser.get<cv::String>("@left-image-format");
    const std::string image_format_R = parser.get<cv::String>("@right-image-format");
    const int disp_size = parser.get<int>("disp_size");
    const int start_number = parser.get<int>("start_number");

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

    cv::Mat I1, I2;

    ImagePreprocessor preprocessor; // Create an instance of the ImagePreprocessor class

    for (int frame_no = start_number;; frame_no++) {
        I1 = cv::imread(cv::format(image_format_L.c_str(), frame_no), cv::IMREAD_GRAYSCALE);
        I2 = cv::imread(cv::format(image_format_R.c_str(), frame_no), cv::IMREAD_GRAYSCALE);
        
        // Check if images are empty, if so break the loop
        if (I1.empty() || I2.empty()) {
            std::cout << "No more images to process or image pair not found." << std::endl;
            break;
        }

        // Preprocess the images
        preprocessor.preprocess_image_pair(I1, I2);

        const int width = I1.cols;
        const int height = I1.rows;

        const int src_depth = I1.type() == CV_8U ? 8 : 16;
        const int dst_depth = disp_size < 256 ? 8 : 16;
        const int src_bytes = src_depth * width * height / 8;
        const int dst_bytes = dst_depth * width * height / 8;

        sgm::StereoSGM sgm(width, height, disp_size, src_depth, dst_depth, sgm::EXECUTE_INOUT_CUDA2CUDA);

        device_buffer d_I1(src_bytes), d_I2(src_bytes), d_disparity(dst_bytes);
        cv::Mat disparity(height, width, dst_depth == 8 ? CV_8S : CV_16S), disparity_color;

        const int invalid_disp = sgm.get_invalid_disparity();

        d_I1.upload(I1.data);
        d_I2.upload(I2.data);

        const auto t1 = std::chrono::system_clock::now();

        sgm.execute(d_I1.data, d_I2.data, d_disparity.data);
        cudaDeviceSynchronize();

        const auto t2 = std::chrono::system_clock::now();
        const auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();
        const double fps = 1e6 / duration;

        d_disparity.download(disparity.data);
        cv::imwrite(cv::format("disparity_output_%04d.png", frame_no), disparity);

        // Save disparity map as text file with pixel values
        //std::ofstream disparity_file(cv::format("disparity_output_%04d.txt", frame_no));
        //if (disparity_file.is_open()) {
         //   for (int y = 0; y < disparity.rows; ++y) {
          //      for (int x = 0; x < disparity.cols; ++x) {
           //         disparity_file << disparity.at<short>(y, x) << " "; // Assuming disparity is CV_16S
           //     }
            //    disparity_file << std::endl;
          //  }
           // disparity_file.close();
        //} else {
          //  std::cerr << "Error: Could not open text file for disparity output." << std::endl;
        //}
        

        // Print the size of the disparity map in MB
        double disparity_size_mb = static_cast<double>(dst_bytes) / (1024 * 1024);
        std::cout << "Size of disparity map: " << disparity_size_mb << " MB" << std::endl;

        // Draw results
        if (I1.type() != CV_8U)
            cv::normalize(I1, I1, 0, 255, cv::NORM_MINMAX, CV_8U);

        colorize_disparity(disparity, disparity_color, disp_size, disparity == invalid_disp);
        cv::putText(disparity_color, cv::format("sgm execution time: %4.1f[msec] %4.1f[FPS]",
            1e-3 * duration, fps), cv::Point(50, 50), 2, 0.75, cv::Scalar(255, 255, 255));

        cv::imshow("left image", I1);
        cv::imshow("disparity", disparity_color);

        cv::waitKey(0); // Hold the window open for inspection; press any key to continue
    }

    return 0;
}