/*
 * Copyright 2016 Jan Dorazil <deu439@gmail.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "register.hpp"

#include <opencv2/opencv.hpp>

using namespace cv;

Register::Register(bool dispFeatures)
  : mDispFeatures(dispFeatures)
{
}


void Register::registerFeature(Mat temp, Mat tempMask, Mat in,
                               OutputArray warp) const
{
    std::vector<KeyPoint> keypointsTemp, keypointsIn;
    Mat descriptorsTemp, descriptorsIn;
    std::vector<DMatch> matches;
    
    // Create ORB feature detector
    Ptr<ORB> detector = ORB::create(500, 1.2, 8, 31, 0, 2,
                                    ORB::HARRIS_SCORE, 31);

    // Detect features in both images
    detector->detectAndCompute(temp, tempMask, keypointsTemp,
                               descriptorsTemp, false);
    detector->detectAndCompute(in, noArray(), keypointsIn,
                               descriptorsIn, false);

    // Match the desctiptors
    BFMatcher bf = BFMatcher(NORM_HAMMING, true);
    bf.match(descriptorsTemp, descriptorsIn, matches);

    // Sort matches by distance
    std::map<double, DMatch> sorted;

    for( int i = 1; i < matches.size(); i++ )
    {
        double dist = matches[i].distance;
        sorted[dist] = matches[i];
    }

    // Select best matches
    int nMatches = 30;
    std::vector<DMatch> good_matches;
    typedef std::map<double, DMatch>::iterator mapIter;

    for(mapIter it = sorted.begin(); it != sorted.end() && nMatches > 0;
        ++it, --nMatches)
    {
        good_matches.push_back(it->second);
    }

    if( nMatches != 0 ) {
        std::cerr << "Not enough features matched\n";
        return;
    }

    // Display matches if needed
    if(mDispFeatures) {
        Mat matchesM;
        drawMatches(temp, keypointsTemp, in, keypointsIn, good_matches,
                    matchesM, Scalar::all(-1), Scalar::all(-1),
                    std::vector<char>(),
                    DrawMatchesFlags::NOT_DRAW_SINGLE_POINTS);

        namedWindow("Matches", WINDOW_NORMAL);
        imshow("Matches", matchesM);
        waitKey();
    }

    // Get points from matches
    std::vector<Point2d> tempPoints;
    std::vector<Point2d> inPoints;

    for( int i = 0; i < good_matches.size(); i++ ) {
        tempPoints.push_back( keypointsTemp[ good_matches[i].queryIdx ].pt );
        inPoints.push_back( keypointsIn[ good_matches[i].trainIdx ].pt );
    }

    // Calculate homography
    Mat warpMatrix = findHomography( tempPoints, inPoints, RANSAC);

    // Copy result
    warpMatrix.copyTo(warp);

    // Transform image
    //warpPerspective(in, out, H, temp.size());

}

void Register::registerECC(Mat temp, Mat in, InputOutputArray warp) const
{
    // Get input matrix
    Mat warpMatrix = warp.getMat();
    warpMatrix.convertTo(warpMatrix, CV_32F);
    std::cout << warpMatrix << std::endl;
    // Number of iterations
    int nIterations = 50;

    // Correlation error
    double eps = 1e-5;

    // Termination criteria
    TermCriteria criteria(TermCriteria::COUNT+TermCriteria::EPS,
                          nIterations, eps);

    // Run ECC algorithm
    findTransformECC( temp, in,
                      warpMatrix, MOTION_HOMOGRAPHY,
                      criteria );

    std::cout << warpMatrix << std::endl;
    warpMatrix.copyTo(warp);
    // Transform image
    //warpPerspective(in, out, warpMatrix, in.size(), INTER_LINEAR 
    //+ WARP_INVERSE_MAP);
}

void Register::process(DataUnit& data)
{
  Mat warp;
  
  assert(!data.mMat[TEMP].empty());
  assert(!data.mMat[TROI].empty());
  assert(!data.mMat[INPU].empty());
  
  // ORB need 8bit grayscale images
  Mat temp, in;
  data.mMat[TEMP].convertTo(temp, CV_8UC1);
  data.mMat[INPU].convertTo(in, CV_8UC1);
   
  // Initialize warp matrix through features
  registerFeature(temp, data.mMat[TROI], in, warp);
  
  // Optimise warp matrix with ECC
  registerECC(temp, in, warp);
  
  // Warp the mIn image and corp to mTemp size
  warpPerspective(data.mMat[INPU], data.mMat[INPU], warp, 
                  data.mMat[TEMP].size(),
                  INTER_LINEAR + WARP_INVERSE_MAP);
}
