#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);

}