//Jaiyam Sharma, Yanmar

#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/highgui.hpp>

#include <vpi/Image.h>
#include <vpi/Context.h>
#include <vpi/Status.h>
#include <vpi/Stream.h>
#include <vpi/algo/StereoDisparity.h>

#include <cstring> // for memset
#include <iostream>
#include <sstream>
#include <chrono>


#define CHECK_STATUS(STMT)                                    \
    do                                                        \
    {                                                         \
        VPIStatus status = (STMT);                            \
        if (status != VPI_SUCCESS)                            \
        {                                                     \
            char buffer[VPI_MAX_STATUS_MESSAGE_LENGTH];       \
            vpiGetLastStatusMessage(buffer, sizeof(buffer));  \
            std::ostringstream ss;                            \
            ss << vpiStatusGetName(status) << ": " << buffer; \
            throw std::runtime_error(ss.str());               \
        }                                                     \
    } while (0);


static std::string format_string(const char* fmt, int args)
{
    const int BUF_SIZE = 1024;
    char buf[BUF_SIZE];
    std::snprintf(buf, BUF_SIZE, fmt, args);
    return std::string(buf);
}

int main(int argc, char* argv[])
{
    VPIContext ctx = nullptr;
    VPIContext cctx = nullptr;

    CHECK_STATUS(vpiContextCreate(0, &ctx));
    std::cout << "Created context at: " << ctx << std::endl;

    CHECK_STATUS(vpiContextSetCurrent(ctx));
    vpiContextPush(ctx);

    VPIImage left      = NULL;
    VPIImage right     = NULL;
    VPIImage disparity = NULL;
    VPIStream stream   = NULL;
    VPIPayload stereo  = NULL;

    std::string strBackend = argv[1];

    // Now parse the backend
    VPIBackend backendType;

    if (strBackend == "cpu")
    {
        backendType = VPI_BACKEND_CPU;
    }
    else if (strBackend == "cuda")
    {
        backendType = VPI_BACKEND_CUDA;
    }
    else if (strBackend == "pva")
    {
        backendType = VPI_BACKEND_PVA;
    }
    else
    {
        throw std::runtime_error("Backend '" + strBackend +
            "' not recognized, it must be either cpu, cuda or pva.");
    }

    // Create the stream for the given backend.
    CHECK_STATUS(vpiStreamCreate(backendType, &stream));

    // Create the payload for Harris Corners Detector algorithm
    VPIStereoDisparityEstimatorParams params;
    params.windowSize = 5;
    params.maxDisparity = 64;

    const int initframe = 0;
    int retval = 0;
    bool first = true;

    VPIImageData data;
    VPIImageData imgData;
    memset(&imgData, 0, sizeof(imgData));
    //vpiContextPop(&ctx);

    cv::Mat cvImageLeft, cvImageRight, cvImageLeft_8U, cvImageRight_8U, cvOut;

    try
    {

    while(true)
    {
        cvImageLeft_8U = cv::imread("/opt/nvidia/vpi/samples/assets/chair_stereo_left.png");
        cvImageRight_8U = cv::imread("/opt/nvidia/vpi/samples/assets/chair_stereo_right.png");
        //cv::imwrite("leftimg.jpg", cvImageLeft_8U);
        //cv::imwrite("rightimg.jpg", cvImageRight_8U);

        // Currently we only accept unsigned 16bpp inputs.
        cvImageLeft_8U.convertTo(cvImageLeft, CV_16UC1);
        cvImageRight_8U.convertTo(cvImageRight, CV_16UC1);
        //CHECK_STATUS(vpiContextSetCurrent(ctx));
        
        // We now wrap the loaded images into a VPIImage object to be used by VPI.
        {
            // First fill VPIImageData with the, well, image data...
            vpiContextPush(ctx);
            CHECK_STATUS(vpiContextSetCurrent(ctx));
            imgData.type                 = VPI_IMAGE_FORMAT_U16;
            imgData.numPlanes            = 1;
            imgData.planes[0].width      = cvImageLeft.cols;
            imgData.planes[0].height     = cvImageLeft.rows;
            imgData.planes[0].pitchBytes = cvImageLeft.step[0];
            imgData.planes[0].data       = cvImageLeft.data;

            // Wrap it into a VPIImage. VPI won't make a copy of it, so the original
            // image must be in scope at all times.
            CHECK_STATUS(vpiImageCreateHostMemWrapper(&imgData, 0, &left));

            imgData.planes[0].width      = cvImageRight.cols;
            imgData.planes[0].height     = cvImageRight.rows;
            imgData.planes[0].pitchBytes = cvImageRight.step[0];
            imgData.planes[0].data       = cvImageRight.data;

            CHECK_STATUS(vpiImageCreateHostMemWrapper(&imgData, 0, &right));
        }

        if(first)
        {
            // Create the image where the disparity map will be stored.
            CHECK_STATUS(vpiContextSetCurrent(ctx));
	    CHECK_STATUS(vpiImageCreate(cvImageLeft.cols, cvImageLeft.rows, VPI_IMAGE_FORMAT_U16, 0, &disparity));

            CHECK_STATUS(vpiCreateStereoDisparityEstimator(backendType, cvImageLeft.cols, cvImageLeft.rows,
                                                           VPI_IMAGE_FORMAT_U16, params.maxDisparity, &stereo));

            first=false;
        }


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

        //CHECK_STATUS(vpiContextSetCurrent(ctx));

        vpiContextPush(ctx);
	// Submit it with the input and output images

        CHECK_STATUS(vpiContextGetCurrent(&cctx));

        std::cout << "Current context while submitting to disparity estimator:" << cctx << std::endl;

        CHECK_STATUS(vpiSubmitStereoDisparityEstimator(stream, stereo, left, right, disparity, &params));

        //CHECK_STATUS(vpiContextSetCurrent(ctx));
        CHECK_STATUS(vpiContextGetCurrent(&cctx));
        std::cout << "Current context while syncing stream:" << cctx << std::endl;
        // Wait until the algorithm finishes processing
        CHECK_STATUS(vpiStreamSync(stream));
        //vpiContextPop(&ctx);
        const auto t2 = std::chrono::system_clock::now();

        const auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();

            // Now let's retrieve the output
        {
            // Lock output to retrieve its data on cpu memory
            vpiContextPush(ctx);
            CHECK_STATUS(vpiImageLock(disparity, VPI_LOCK_READ, &data));

            // Make an OpenCV matrix out of this image
            cv::Mat cvOut_8U(data.planes[0].height, data.planes[0].width, CV_16UC1, data.planes[0].data,
                          data.planes[0].pitchBytes);

            // Scale result and write it to disk
            double min, max;
            minMaxLoc(cvOut_8U, &min, &max);
            cvOut_8U.convertTo(cvOut_8U, CV_8UC1, 255.0 / (max - min), -min);
            cv::applyColorMap(cvOut_8U, cvOut, cv::COLORMAP_JET);

            cv::putText(cvOut, format_string("sgm execution time: %4.1f[msec]", 1e-3 * duration),
            cv::Point(50, 50), 2, 0.75, cv::Scalar(255, 255, 255));

            //cv::imshow("disparity", cvOut);
            //char c=cv::waitKey(1);
            //if (c==27 /*ESC*/){break;}
            cv::imwrite("output.png", cvOut);
            // Done handling output, don't forget to unlock it.
            CHECK_STATUS(vpiImageUnlock(disparity));
            //vpiContextPop(&ctx);
        }

	} //endwhile
}

catch (std::exception &e)
    {
        std::cerr << e.what() << std::endl;
        retval = 1;
    }

    if (stream != NULL)
    {
        vpiStreamSync(stream);
    }

    //vpiImageDestroy(left);
    //vpiImageDestroy(right);
    //vpiImageDestroy(disparity);
    //vpiPayloadDestroy(stereo);
    //vpiStreamDestroy(stream);
    vpiContextDestroy(ctx);
    return retval;

}