//
// Created by fandikai on 17-10-5.
//

#include "OpticalFlow.h"
#include <Eigen/Dense>

float getGraysubPixel(float x, float y, cv::Mat img) {
    float xl = std::floor(x);
    float yl = std::floor(y);
    float xr = std::ceil(x);
    float yr = std::ceil(y);
    float f11 = img.at<uchar>(yl, xl);
    float f12 = img.at<uchar>(yr, xl);
    float f21 = img.at<uchar>(yl, xr);
    float f22 = img.at<uchar>(yr, xr);
    float a1 = x - xl, a2 = xr - x, b1 = y - yl, b2 = yr - y;
    return f11 * a2 * b2 + f22 * a1 * b1 + f12 * a2 * b1 + f21 * a1 * b2;
}

bool opticalPoint(cv::Mat img1, cv::Mat img2, cv::Point2f &px, int windowSize, int nIter, cv::Point2f& pxNext) {
    pxNext = px;
    std::vector<cv::Mat> imgTemplates(3, cv::Mat()), imgTargets(3, cv::Mat());
    imgTemplates[0] = img1;
    imgTargets[0] = img2;
    int winSize = windowSize / 2;
    Eigen::Vector2f g = Eigen::Vector2f::Zero();
    Eigen::Vector2f u(px.x, px.y);
    for(int i = 1; i < 3; ++i) {
        cv::pyrDown(imgTargets[i - 1], imgTargets[i]);
        cv::pyrDown(imgTemplates[i - 1], imgTemplates[i]);
    }

    for(int i = int(imgTargets.size()) - 1; i >= 0; --i) {
        Eigen::Vector2f ul = u / std::pow(2.0f, i);
        if(ul(0) - winSize -1 < 0 || ul(0) + winSize + 1 >= imgTargets[i].cols
           || ul(1) - winSize - 1< 0 || ul(1) + winSize + 1 >= imgTargets[i].rows )
            return false;
        Eigen::Matrix2f G = Eigen::Matrix2f::Zero();
        for(int p = -winSize; p <= winSize; ++p) {
            for(int q = -winSize; q <= winSize; ++q) {
                float Ix = 0.5f * (imgTemplates[i].at<uchar>(ul(1) + q, ul(0) + p + 1) -
                                   imgTemplates[i].at<uchar>(ul(1) + q, ul(0) + p - 1));
                float Iy = 0.5f * (imgTemplates[i].at<uchar>(ul(1) + q + 1, ul(0) + p) -
                                   imgTemplates[i].at<uchar>(ul(1) + q - 1, ul(0) + p));
                G(0, 0) += Ix * Ix;
                G(0, 1) += G(1, 0) = Ix * Iy;
                G(1, 1) += Iy * Iy;
            }
        }
        if(G.determinant() < 0.1)
            return false;

        Eigen::Vector2f v = Eigen::Vector2f::Zero();
        for(int it = 0; it < nIter; ++it) {
            Eigen::Vector2f b = Eigen::Vector2f::Zero();
            for(int p = -winSize; p <= winSize; ++p) {
                for(int q = -winSize; q <= winSize; ++q) {
                    if(ul(0) + p + g(0) + v(0) < 0 || ul(0) + p + g(0) + v(0) >= imgTargets[i].cols
                       || ul(1) + q + g(1) + v(1) < 0 || ul(1) + q + g(1) + v(1) >= imgTargets[i].rows)
                        return false;

                    float Ix = 0.5f * (imgTemplates[i].at<uchar>(ul(1) + q, ul(0) + p + 1) -
                                       imgTemplates[i].at<uchar>(ul(1) + q, ul(0) + p - 1));
                    float Iy = 0.5f * (imgTemplates[i].at<uchar>(ul(1) + q + 1, ul(0) + p) -
                                       imgTemplates[i].at<uchar>(ul(1) + q - 1, ul(0) + p));
                    float dI = imgTemplates[i].at<uchar>(ul(1) + q, ul(0) + p)
                               - getGraysubPixel(ul(0) + p + g(0) + v(0), ul(1) + q + g(1) + v(1), imgTargets[i]);
                    b(0) += Ix * dI;
                    b(1) += Iy * dI;
                }
            }
            Eigen::Vector2f eta = G.ldlt().solve(b);
            v += eta;
        }
        g += v;
    }
    pxNext.x += g(0);
    pxNext.y += g(1);

    std::cout << "px, pxNext = " << px.x << "," << px.y << ";" << pxNext.x << "," << pxNext.y << std::endl;
    return true;
}

void OpticalFlow::handleFrame(cv::Mat img1, cv::Mat img2,
                              std::vector<cv::Point2f> fastPoints, int windowSize,
                              int nIter) {
    std::vector<bool> fastPointsValid(fastPoints.size(), false);
    std::vector<cv::Point2f> fastPointsNext(fastPoints.size(), cv::Point2f(0, 0));
    size_t cnt = 0;
    for(size_t i = 0; i < fastPoints.size(); ++i) {
        if(opticalPoint(img1, img2, fastPoints[i], windowSize, nIter, fastPointsNext[i])) {
            fastPointsValid[i] = true;
            cnt++;
        }
    }
    std::cout << "get points " << cnt << std::endl;

    static cv::Scalar colorLRMatcher[3] = {
            cv::Scalar(0,255,0),
            cv::Scalar(0,0,255),
            cv::Scalar(255,0,0)
    };

    std::vector<cv::Point2f> curCorners = fastPoints;
    std::vector <uchar> status;
    std::vector <float> err;

    cv::Size winSize = cv::Size(7, 7);
    int maxLevel = 3;
    cv::TermCriteria criteria = cv::TermCriteria(
            (cv::TermCriteria::COUNT + cv::TermCriteria::EPS), 20, 0.05);

    int flag = 0;

    double minEigThreshold = 0.002;

    cv::calcOpticalFlowPyrLK(img1, img2,
                             fastPoints, curCorners,
                             status, err, winSize, maxLevel,
                             criteria, flag, minEigThreshold);


    cv::Mat matchimg(img1.rows * 2, img1.cols, CV_8UC3), match(img1.rows * 2, img1.cols, CV_8UC3);
    cv::Mat img1C3, img2C3;
    cv::cvtColor(img1, img1C3, CV_GRAY2RGB);
    cv::cvtColor(img2, img2C3, CV_GRAY2RGB);
    img1C3.copyTo(matchimg(cv::Rect(0, 0, img1.cols, img1.rows)));
    img2C3.copyTo(matchimg(cv::Rect(0, img2.rows, img2.cols, img2.rows)));
    img1C3.copyTo(match(cv::Rect(0, 0, img1.cols, img1.rows)));
    img2C3.copyTo(match(cv::Rect(0, img2.rows, img2.cols, img2.rows)));
    int var = 0;
    for(size_t itmp = 0; itmp < fastPointsValid.size(); itmp++){
        if(fastPointsValid[itmp]) {
            cv::circle(matchimg, fastPoints[itmp], 3, cv::Scalar(255,0,0), -1);
            cv::circle(matchimg, cv::Point(fastPointsNext[itmp].x, fastPointsNext[itmp].y + img2.rows), 3, cv::Scalar(0,255,0), -1);
            cv::line(matchimg, fastPoints[itmp],
                     cv::Point(fastPointsNext[itmp].x, fastPointsNext[itmp].y + img2.rows),
                     colorLRMatcher[(var++) % 3]);
        }
    }

    var = 0;
    for(size_t itmp = 0; itmp < fastPoints.size(); ++itmp) {
        if(status[itmp]) {
            cv::circle(match, fastPoints[itmp], 3, cv::Scalar(255,0,0), -1);
            cv::circle(match, cv::Point(curCorners[itmp].x, curCorners[itmp].y + img2.rows), 3, cv::Scalar(0,255,0), -1);
            cv::line(match, fastPoints[itmp],
                     cv::Point(curCorners[itmp].x, curCorners[itmp].y + img2.rows),
                     colorLRMatcher[(var++) % 3]);
        }
    }

    cv::imshow("opencv matching", match);
    cv::imshow("my matchimg", matchimg);
    cv::waitKey(0);
}