#include <iostream>
#include <opencv2/opencv_modules.hpp>
#ifdef HAVE_OPENCV_ARUCO
#include <opencv2/opencv.hpp>
#include <opencv2/aruco.hpp>

using namespace std;
using namespace cv;

namespace
{
enum Pattern { CHESSBOARD, CIRCLES_GRID, ASYMMETRIC_CIRCLES_GRID };

Scalar randomColor( RNG& rng )
{
  int icolor = (unsigned int) rng;
  return Scalar( icolor & 255, (icolor >> 8) & 255, (icolor >> 16) & 255 );
}

void calcChessboardCorners(Size boardSize, float squareSize, vector<Point3f>& corners, Pattern patternType = CHESSBOARD)
{
    corners.resize(0);

    switch (patternType)
    {
    case CHESSBOARD:
    case CIRCLES_GRID:
        for( int i = 0; i < boardSize.height; i++ )
            for( int j = 0; j < boardSize.width; j++ )
                corners.push_back(Point3f(float(j*squareSize),
                                          float(i*squareSize), 0));
        break;

    case ASYMMETRIC_CIRCLES_GRID:
        for( int i = 0; i < boardSize.height; i++ )
            for( int j = 0; j < boardSize.width; j++ )
                corners.push_back(Point3f(float((2*j + i % 2)*squareSize),
                                          float(i*squareSize), 0));
        break;

    default:
        CV_Error(Error::StsBadArg, "Unknown pattern type\n");
    }
}

void perspectiveCorrection(const string &img1Path, const string &img2Path, const Size &patternSize, RNG &rng)
{
    Mat img1 = imread(img1Path);
    Mat img2 = imread(img2Path);

    //! [find-corners]
    vector<Point2f> corners1, corners2;
    bool found1 = findChessboardCorners(img1, patternSize, corners1);
    bool found2 = findChessboardCorners(img2, patternSize, corners2);
    //! [find-corners]

    if (!found1 || !found2)
    {
        cout << "Error, cannot find the chessboard corners in both images." << endl;
        return;
    }

    //! [estimate-homography]
    Mat H = findHomography(corners1, corners2);
    cout << "H:\n" << H << endl;
    //! [estimate-homography]

    //! [warp-chessboard]
    Mat img1_warp;
    warpPerspective(img1, img1_warp, H, img1.size());
    //! [warp-chessboard]

    Mat img_draw_warp;
    hconcat(img2, img1_warp, img_draw_warp);
    imshow("Desired chessboard view / Warped source chessboard view", img_draw_warp);

    //! [compute-transformed-corners]
    Mat img_draw_matches;
    hconcat(img1, img2, img_draw_matches);
    for (size_t i = 0; i < corners1.size(); i++)
    {
        Mat pt1 = (Mat_<double>(3,1) << corners1[i].x, corners1[i].y, 1);
        Mat pt2 = H * pt1;
        pt2 /= pt2.at<double>(2);

        Point end( (int) (img1.cols + pt2.at<double>(0)), (int) pt2.at<double>(1) );
        line(img_draw_matches, corners1[i], end, randomColor(rng), 2);
    }

    imshow("Draw matches", img_draw_matches);
    waitKey();
    //! [compute-transformed-corners]
}

const char* params
    = "{ help h         |       | print usage }"
      "{ image1         | ../data/left02.jpg | path to the source chessboard image }"
      "{ image2         | ../data/left01.jpg | path to the desired chessboard image }"
      "{ width bw       | 9     | chessboard width }"
      "{ height bh      | 6     | chessboard height }";
}

int main(int argc, char *argv[])
{
    cv::RNG rng( 0xFFFFFFFF );
    CommandLineParser parser(argc, argv, params);

    if (parser.has("help"))
    {
        parser.about("Code for homography tutorial.\n"
            "Example 2: perspective correction.\n");
        parser.printMessage();
        return 0;
    }

    Size patternSize(parser.get<int>("width"), parser.get<int>("height"));
    perspectiveCorrection(parser.get<String>("image1"),
                          parser.get<String>("image2"),
                          patternSize, rng);

    return 0;
}
#else
int main()
{
    std::cerr << "FATAL ERROR: This sample requires opencv_aruco module (from opencv_contrib)" << std::endl;
    return 0;
}
#endif
