// g++ test_v2.cpp -o test_v2 `pkg-config --cflags --libs opencv4 gstreamer-1.0 gstreamer-rtsp-server-1.0`


#include <iostream>
//#include "opencv2/highgui.hpp"
#include <opencv4/opencv2/highgui.hpp>
//#include "opencv2/imgproc.hpp"
#include <opencv4/opencv2/imgproc.hpp>

#include <opencv4/opencv2/core/core.hpp>    
#include <opencv4/opencv2/imgproc/imgproc.hpp>    
#include <opencv4/opencv2/calib3d/calib3d.hpp>    
#include <opencv4/opencv2/highgui/highgui.hpp>    
#include <iostream>    
#include <fstream>    

using namespace cv;
using namespace std;


#define SRC_WIDTH  1280
#define SRC_HEIGHT 720

void InitMat(float* num , Mat& m)
{
	for (int i = 0; i<m.rows; i++)
	for (int j = 0; j<m.cols; j++)
		m.at<float>(i, j) = *(num + i * m.rows + j);
}

Mat mapx = Mat(Size(SRC_WIDTH, SRC_HEIGHT), CV_32FC1);
Mat mapy = Mat(Size(SRC_WIDTH, SRC_HEIGHT), CV_32FC1);
Mat newimage ; //Output picture after correction

void jiaozheng_init(void)
{
	int OK = 0;


	Mat R = Mat::eye(3, 3, CV_32F);
	//Parameter matrix
	
	float neican_data[] = { 668.611073, 0.000000, 601.394929, 0.000000, 666.724859, 340.687923, 0.000000, 0.000000, 1.000000 }; //AR0231
	Mat cameraMatrix(3, 3, CV_32FC1);
	InitMat(neican_data,cameraMatrix);

	cout << "cameraMatrix= " << endl << " " << cameraMatrix << endl << endl;
	//Measured distortion coefficient	
	float jibian_data[] = { -0.299505, 0.066764, -0.001612, 0.000428, 0.000000};//AR0231
	Mat distCoeffs(1, 5, CV_32FC1); /* 5 distortion coefficients of the camera：k1,k2,p1,p2,k3 */
	InitMat( jibian_data, distCoeffs);
	cout << "distCoeffs= " << endl << " " << distCoeffs << endl << endl;

	/********Camera correction*******************************************************************************/
	//cout << "mapx= " << endl << " " << mapx << endl << endl;
	//cout << "mapy= " << endl << " " << mapy << endl << endl;
	initUndistortRectifyMap(cameraMatrix, distCoeffs, R, cameraMatrix, Size(SRC_WIDTH, SRC_HEIGHT), CV_32FC1, mapx, mapy);
	//cout << "mapx= " << endl << " " << mapx << endl << endl;
	//cout << "mapy= " << endl << " " << mapy << endl << endl;

}

int jibianjiaozheng(Mat src_picture, Mat & dst_picture)
{
	int OK = 0;
	static unsigned int INIT_MAT_OK = 1;

	if (INIT_MAT_OK)
	{
		newimage = src_picture.clone(); //Output picture after correction
		INIT_MAT_OK = 0;
	}

	remap(src_picture, newimage, mapx, mapy, INTER_LINEAR);
	dst_picture= newimage+1;
	return OK;
}


int main()
{

    cv::VideoCapture cap("/home/nvidia/deepstream-distortion/fov120_distortion_1min.mp4");
/*
    const char* gst = 	"v4l2src device=/dev/video2 ! "                                
                                " videoconvert ! "
                                "appsink";

    
    VideoCapture cap(gst);
*/
    int fps = cap.get(cv::CAP_PROP_FPS);
    int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
    int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
    Size size(w,h);

        
    VideoWriter out("appsrc ! videoconvert ! x264enc tune=zerolatency speed-preset=superfast ! h264parse ! rtph264pay ! udpsink host=127.0.0.1 port=5000",CAP_GSTREAMER,0,fps,size,true);

    if(!cap.isOpened() || !out.isOpened())
    {
        std::cout<<"VideoCapture or VideoWriter not opened"<<std::endl;
        exit(-1);
    }

    //cv::Mat frame;
    Mat frame;
    Mat after;
    int hrst = 0;
    unsigned int jiaozheng_init_flag = 0;
    while(true) {
        cap.read(frame);
        
        //imshow("MyCameraPreview",frame);
        if(frame.empty())
            break;
 
        /* remove distortion */
        if (!jiaozheng_init_flag)
        {
           printf("initialization\n");
           jiaozheng_init();
           jiaozheng_init_flag = 1;
        }

        hrst = jibianjiaozheng(frame, after);
        if (hrst != 0)
        {
           printf("Distortion correction failed\n");
        }
        else
        {
           //printf("Distortion correction successful\n");
        }

        
        //out.write(frame);
        out.write(after);
        //imshow("MyCameraPreview",frame);
        waitKey(50);
    }
}