#include <iostream>
#include <fstream>
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/stitching/stitcher.hpp"
#include "opencv/cv.h"
#include "opencv2/gpu/gpu.hpp"

#include <time.h>
#include <windows.h>                // for Windows APIs
using namespace std;
using namespace std;
using namespace cv;
using namespace cv::gpu;


bool try_use_gpu = false;
vector<Mat> imgs;
string result_name = "result.tif";

void printUsage();
int parseCmdArgs(int argc, char** argv);
void lensCorrectorBarrel( unsigned char*, int, int);

int main(int argc, char* argv[])
{

	/*DeviceInfo info = getDevice();
        cout << info.name() << endl;*/

LARGE_INTEGER frequency;        // ticks per second
    LARGE_INTEGER t1, t2;           // ticks
    double elapsedTime;

QueryPerformanceFrequency(&frequency);
    int retval = parseCmdArgs(argc, argv);
    if (retval) return -1;
	printf("Images loaded.\n");
    Mat pano;
	
	printf("Channels image1: %d\n", imgs[0].type());
	printf("Channels image2: %d\n", imgs[1].type());


	int width = 170;
	cv::Rect rect1 = cvRect( imgs[0].cols-width, 0, width, imgs[0].rows); //second half of the first image
	cv::vector<cv::Rect> roi1;
	roi1.push_back(rect1);


	cv::Rect rect12 =  cvRect(0, 0,width, imgs[1].rows);  //first half of the second image
	cv::Rect rect23 = cvRect(imgs[1].cols -width, 0, width, imgs[1].rows);  //second half of the first image
	cv::vector<cv::Rect> roi2;
	roi2.push_back(rect12);
	roi2.push_back(rect23);

	cv::Rect rect34 =  cvRect(0, 0,width, imgs[2].rows);  //first half of the second image
	cv::Rect rect45 = cvRect(imgs[2].cols -width, 0, width, imgs[2].rows);  //second half of the first image
	cv::vector<cv::Rect> roi3;
	roi3.push_back(rect34);
	roi3.push_back(rect45);

	cv::Rect rect56 =  cvRect(0, 0,width, imgs[3].rows);  //first half of the second image
	cv::Rect rect67 = cvRect(imgs[3].cols -width, 0, width, imgs[3].rows);  //second half of the first image
	cv::vector<cv::Rect> roi4;
	roi4.push_back(rect56);
	roi4.push_back(rect67);

	cv::Rect rect78 =  cvRect(0, 0,width, imgs[4].rows);  //first half of the second image
	cv::Rect rect89 = cvRect(imgs[4].cols -width, 0, width, imgs[4].rows);  //second half of the first image
	cv::vector<cv::Rect> roi5;
	roi5.push_back(rect78);
	roi5.push_back(rect89);

	cv::Rect rect910 = cvRect(0, 0, width, imgs[5].rows); //first half of the third image
	cv::vector<cv::Rect> roi6;
	roi6.push_back(rect910);



	cv::vector<cv::vector<cv::Rect>> rois;
	rois.resize(6);
	rois[0] = roi1;
	rois[1] = roi2;
	rois[2] = roi3;
	rois[3] = roi4;
	rois[4] = roi5;
	rois[5] = roi6;
	//cv::imshow("debug_img1", imgs[0]);
	Mat tmp1, tmp2, tmp3, tmp4, tmp5, tmp6;
	cvtColor(imgs[0], tmp1, CV_BGR2GRAY); 
	cvtColor(imgs[1], tmp2, CV_BGR2GRAY); 
	cvtColor(imgs[2], tmp3, CV_BGR2GRAY); 
	cvtColor(imgs[3], tmp4, CV_BGR2GRAY); 
	cvtColor(imgs[4], tmp5, CV_BGR2GRAY); 
	cvtColor(imgs[5], tmp6, CV_BGR2GRAY); 
printf("ROI ready\n");

	lensCorrectorBarrel(tmp1.data, tmp1.cols, tmp1.rows);
	lensCorrectorBarrel(tmp2.data, tmp2.cols, tmp2.rows);
	lensCorrectorBarrel(tmp3.data, tmp3.cols, tmp3.rows);
	lensCorrectorBarrel(tmp4.data, tmp4.cols, tmp4.rows);
	lensCorrectorBarrel(tmp5.data, tmp5.cols, tmp5.rows);
	lensCorrectorBarrel(tmp6.data, tmp6.cols, tmp6.rows);
printf("Lens correction completed\n");
	imwrite("corr1.tif", tmp1);
	imwrite("corr2.tif", tmp2);
	imwrite("corr3.tif", tmp3);
	imwrite("corr4.tif", tmp4);
	imwrite("corr5.tif", tmp5);
	imwrite("corr6.tif", tmp6);
	/*namedWindow( "debug_img_corr1", CV_WINDOW_AUTOSIZE );
	namedWindow( "debug_img_corr2", CV_WINDOW_AUTOSIZE );
	cv::imshow("debug_img_corr1", tmp1);
	cv::imshow("debug_img_corr2", tmp2);*/

	cvtColor(tmp1, imgs[0], CV_GRAY2BGR); 
	cvtColor(tmp2, imgs[1], CV_GRAY2BGR); 
	cvtColor(tmp3, imgs[2], CV_GRAY2BGR); 
	cvtColor(tmp4, imgs[3], CV_GRAY2BGR); 
	cvtColor(tmp5, imgs[4], CV_GRAY2BGR);
	cvtColor(tmp6, imgs[5], CV_GRAY2BGR);
	//char cc = cvWaitKey(0);

	/*cv::rectangle( imgs[0], rect1, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[1], rect12, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[1], rect23, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[2], rect34, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[2], rect45, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[3], rect56, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[3], rect67, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[4], rect78, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[4], rect89, CV_RGB(255,0,0), 1);
	cv::rectangle( imgs[5], rect910, CV_RGB(255,0,0), 1);
	cv::imshow("debug_img1", imgs[0]);
	cv::imshow("debug_img2", imgs[1]);
	cv::imshow("debug_img3", imgs[2]);
	cv::imshow("debug_img4", imgs[3]);
	cv::imshow("debug_img5", imgs[4]);
	cv::imshow("debug_img6", imgs[5]);*/
	char c = cvWaitKey(0);
printf("Stitching started...\n");
	QueryPerformanceCounter(&t1);
    Stitcher stitcher = Stitcher::createDefault(false);
	//stitcher.setRegistrationResol(0.5);
	printf("registrationResol=%f\n", stitcher.registrationResol());
	//stitcher.setPanoConfidenceThresh(0.03);
	//stitcher.setWaveCorrection(false);

	//stitcher.setFeaturesMatcher(new cv::detail::BestOf2NearestMatcher(true));
	//detail::SurfFeaturesFinder *featureFinder = new detail::SurfFeaturesFinder(100);
	//stitcher.setFeaturesMatcher(matcher);
	//stitcher.setFeaturesFinder(featureFinder);
    Stitcher::Status status = stitcher.stitch(imgs, rois, pano);
	//Stitcher::Status status = stitcher.composePanorama(imgs, pano);
	
    if (status != Stitcher::OK)
    {
        cout << "Can't stitch images, error code = " << status << endl;
        return -1;
    }


QueryPerformanceCounter(&t2);
elapsedTime = (t2.QuadPart - t1.QuadPart) * 1000.0 / frequency.QuadPart;
	printf("%.6f ms",elapsedTime);
    imwrite(result_name, pano);

	char ac = cvWaitKey(0);
	if(ac == 'ESC')
		return 0;
}


void printUsage()
{
    cout <<
        "Rotation model images stitcher.\n\n"
        "stitching img1 img2 [...imgN]\n\n"
        "Flags:\n"
        "  --try_use_gpu (yes|no)\n"
        "      Try to use GPU. The default value is 'no'. All default values\n"
        "      are for CPU mode.\n"
        "  --output <result_img>\n"
        "      The default is 'result.jpg'.\n";
}


int parseCmdArgs(int argc, char** argv)
{
    if (argc == 1)
    {
        printUsage();
        return -1;
    }
    for (int i = 1; i < argc; ++i)
    {
        if (string(argv[i]) == "--help" || string(argv[i]) == "/?")
        {
            printUsage();
            return -1;
        }
        else if (string(argv[i]) == "--try_gpu")
        {
            if (string(argv[i + 1]) == "no")
                try_use_gpu = true;
            else if (string(argv[i + 1]) == "yes")
                try_use_gpu = true;
            else
            {
                cout << "Bad --try_use_gpu flag value\n";
                return -1;
            }
            i++;
        }
        else if (string(argv[i]) == "--output")
        {
            result_name = argv[i + 1];
            i++;
        }
        else
        {
            Mat img = imread(argv[i]);
            if (img.empty())
            {
                cout << "Can't read image '" << argv[i] << "'\n";
                return -1;
            }
            imgs.push_back(img);
        }
    }
    return 0;
}




void lensCorrectorBarrel( unsigned char *pixels, int cols, int rows)
{
	unsigned char *pixelsCopy;
	pixelsCopy = (unsigned char *)malloc( cols * rows );
	memcpy( (char*) pixelsCopy, (char*)pixels, cols*rows);

	
	// parameters for correction

	//for one good lense
   /* double paramA = 0.0; // affects only the outermost pixels of the image
    double paramB = -0.000000062; // most cases only require b optimization
    double paramC = -0.0000020; // most uniform correction
    double paramD = 1.0120;// - paramA - paramB - paramC; // describes the linear scaling of the image
 */
	//for Edmundoptics 58201
	double paramA = 0.0; // affects only the outermost pixels of the image
    double paramB = -0.000000150; // most cases only require b optimization
    double paramC = -0.0000085; // most uniform correction
    double paramD = 1.04400;// - paramA - paramB - paramC; // describes the linear scaling of the image


    // center of dst image
    double centerX = (cols - 1) / 2.0;
    double centerY = (rows - 1) / 2.0;   
    
	#pragma omp parallel for
    for(int x = 0; x < cols; x++) {
        for(int y = 0; y < rows; y++) {
        	int dstX = x;
            int dstY = y;
            // difference between center and point
            double diffX = centerX - dstX;
            double diffY = centerY - dstY;
//            printf("diffX=%f\n", diffX);
            // distance or radius of dst image
            double dstR = sqrt(diffX * diffX + diffY * diffY);

            // distance or radius of src image (with formula)
            double srcR = dstR*(paramA * dstR * dstR * dstR + paramB * dstR * dstR + paramC * dstR + paramD);

            // comparing old and new distance to get factor
            double factor =fabs( srcR / dstR);
            // coordinates in source image
            double srcXd = centerX + (diffX * factor);
            double srcYd = centerY + (diffY * factor);
			
            // no interpolation yet (just nearest point)
            int srcX = floor(srcXd+0.5);
            int srcY = floor(srcYd+0.5);     
            if(srcX >= 0 && srcY >= 0 && srcX < cols && srcY < rows) {
                int dstPos = dstY * cols + dstX;
                pixels[cols * rows -dstPos] = pixelsCopy[srcY * cols + srcX];
            }       
        }
    }
	free (pixelsCopy);
}