#include <Rcpp.h>

// OpenCV
#include <opencv2/opencv.hpp>
#include "opencv2/xfeatures2d.hpp"
#include "opencv2/features2d.hpp"
#include "opencv2/shape/shape_transformer.hpp"

// Internal functions
#include "auxiliary.h"

using namespace Rcpp;
using namespace std;
using namespace cv;
using namespace cv::xfeatures2d;

////
// processing
////

// process images before keypoint detection
cv::Mat preprocessImage(Mat &im, const bool invert, const char* flipflop, const char* rotate)
{
  // normalize
  Mat imNorm;
  cv::normalize(im, imNorm, 0, 255, cv::NORM_MINMAX, CV_8UC1);
  
  // rotate image
  Mat imRotate;
  if(atoi(rotate) > 0){
    cv::rotate(imNorm, imRotate, (atoi(rotate)/90)-1);
  } else {
    imRotate = imNorm;
  }
  
  // Flipflop image
  Mat imFlipFlop;
  if(strcmp(flipflop, "Flip") == 0){
    cv::flip(imRotate, imFlipFlop, 0);
  } else if(strcmp(flipflop, "Flop") == 0){
    cv::flip(imRotate, imFlipFlop, 1);
  } else if(strcmp(flipflop, "None") == 0){
    imFlipFlop = imRotate;
  }
  
  // invert/negate image and full processed image
  Mat imProcess;
  if(invert) {
    cv::bitwise_not(imFlipFlop, imProcess);
  } else {
    imProcess = imFlipFlop;
  }
  
  // return
  return imProcess;
}

// revert the processing on registrated images using the reference image
cv::Mat reversepreprocessImage(Mat &im, const char* flipflop, const char* rotate)
{
  
  // Flipflop image
  Mat imFlipFlop;
  if(strcmp(flipflop, "Flip") == 0){
    cv::flip(im, imFlipFlop, 0);
  } else if(strcmp(flipflop, "Flop") == 0){
    cv::flip(im, imFlipFlop, 1);
  } else if(strcmp(flipflop, "None") == 0){
    imFlipFlop = im;
  }
  
  // rotate image
  Mat imRotate;
  if(atoi(rotate) > 0){
    cv::rotate(imFlipFlop, imRotate, ((360-atoi(rotate))/90)-1);
  } else {
    imRotate = imFlipFlop;
  }
  
  // return
  return imRotate;
}

// [[Rcpp::export]]
Rcpp::RawVector warpImage(Rcpp::RawVector ref_image, Rcpp::RawVector query_image, 
                          Rcpp::List mapping,
                          const int width1, const int height1,
                          const int width2, const int height2)
{
  // Read reference image
  cv::Mat imReference = imageToMat(ref_image, width1, height1);
  
  // Read image to be aligned
  cv::Mat im = imageToMat(query_image, width2, height2);
  cv::Mat im_temp;
  
  // list 
  int n = mapping.size();
  
  // iterate over the list 
  for(int i=0; i<n; i++) {
    
    // get the current mapping
    Rcpp::List cur_mapping = mapping[i];
    
    // Get transformation matrix
    cv::Mat h = numericMatrixToMat(cur_mapping[0]);
    if(h.cols > 0){
      
      // transform coordinates
      cv::warpPerspective(im, im_temp, h, imReference.size()); 
      im = im_temp;
      
    } 
    
    // non-rigid warping
    if(cur_mapping[1] != R_NilValue){
      
      // get landmarks
      Rcpp::List keypoints = cur_mapping[1];
      std::vector<cv::Point2f> ref_mat = numericMatrixToPoint2f(keypoints[0]);
      std::vector<cv::Point2f> query_mat = numericMatrixToPoint2f(keypoints[1]);
      
      // get matches
      std::vector<cv::DMatch> matches;
      for (unsigned int i = 0; i < ref_mat.size(); i++)
        matches.push_back(cv::DMatch(i, i, 0));
      
      // calculate transformation
      Ptr<ThinPlateSplineShapeTransformer> tps = cv::createThinPlateSplineShapeTransformer(0);
      tps->estimateTransformation(ref_mat, query_mat, matches);
      
      // determine extension limits for both images
      int y_max = max(im.rows, imReference.rows);
      int x_max = max(im.cols, imReference.cols);
      
      // extend images
      cv::copyMakeBorder(im, im, 0.0, (int) (y_max - im.rows), 0.0, (x_max - im.cols), cv::BORDER_CONSTANT, Scalar(0, 0, 0));
      
      // transform image
      tps->warpImage(im, im_temp);
      
      // resize image
      cv::Mat im_temp_cropped  = im_temp(cv::Range(0,imReference.size().height), cv::Range(0,imReference.size().width));
      im_temp = im_temp_cropped.clone();
      
    } else {
      
      // pass registered object
      im_temp = im;
    }
    
    im = im_temp;
  }
  
  // return
  return matToImage(im);
} 


// [[Rcpp::export]]
Rcpp::RawVector warpImageAuto(Rcpp::RawVector ref_image, Rcpp::RawVector query_image, 
                              Rcpp::List mapping,
                              const int width1, const int height1,
                              const int width2, const int height2)
{
  // Get transformation matrix
  cv::Mat h = numericMatrixToMat(mapping[0]);
  
  // Read reference image
  cv::Mat imReference = imageToMat(ref_image, width1, height1);
  
  // Read image to be aligned
  cv::Mat im = imageToMat(query_image, width2, height2);
  
  // transform coordinates
  Mat imWarp;
  cv::warpPerspective(im, imWarp, h, imReference.size()); 
  
  // non-rigid warping
  cv::Mat imReg;
  if(mapping[1] != R_NilValue){
    
    // get landmarks
    Rcpp::List keypoints = mapping[1];
    std::vector<cv::Point2f> ref_mat = numericMatrixToPoint2f(keypoints[0]);
    std::vector<cv::Point2f> query_mat = numericMatrixToPoint2f(keypoints[1]);
    
    // get matches
    std::vector<cv::DMatch> matches;
    for (unsigned int i = 0; i < ref_mat.size(); i++)
      matches.push_back(cv::DMatch(i, i, 0));
    
    // calculate transformation
    Ptr<ThinPlateSplineShapeTransformer> tps = cv::createThinPlateSplineShapeTransformer(0);
    tps->estimateTransformation(ref_mat, query_mat, matches);
    
    // determine extension limits for both images
    int y_max = max(imWarp.rows, imReference.rows);
    int x_max = max(imWarp.cols, imReference.cols);
    
    // extend images
    cv::copyMakeBorder(imWarp, imWarp, 0.0, (int) (y_max - imWarp.rows), 0.0, (x_max - imWarp.cols), cv::BORDER_CONSTANT, Scalar(0, 0, 0));
    
    // transform image
    tps->warpImage(imWarp, imReg);
    
    // resize image
    // cv::resize(im1Reg, im1Reg, im2.size());
    cv::Mat imReg_cropped  = imReg(cv::Range(0,imReference.size().height), cv::Range(0,imReference.size().width));
    imReg = imReg_cropped.clone();
    
  } else {
    
    // pass registered object
    imReg = imWarp;
  }
  
  // return
  return matToImage(imReg);
}


// [[Rcpp::export]]
Rcpp::RawVector warpImageManual(Rcpp::RawVector ref_image, Rcpp::RawVector query_image, 
                                Rcpp::List mapping,
                                const int width1, const int height1,
                                const int width2, const int height2)
{
  // Get landmarks as Point2f and transformation matrix
  cv::Mat h = numericMatrixToMat(mapping[0]);
  Rcpp::List keypoints = mapping[1];
  std::vector<cv::Point2f> ref_mat = numericMatrixToPoint2f(keypoints[0]);
  std::vector<cv::Point2f> query_mat = numericMatrixToPoint2f(keypoints[1]);
  
  // Read reference image
  cv::Mat imReference = imageToMat(ref_image, width1, height1);
  
  // Read image to be aligned
  cv::Mat im = imageToMat(query_image, width2, height2);
  
  // transform coordinates
  Mat imWarp;
  if(h.cols > 0){
    cv::warpPerspective(im, imWarp, h, imReference.size()); 
  } else {
    imWarp = im;
  }
  
  // get matches
  std::vector<cv::DMatch> matches;
  for (unsigned int i = 0; i < ref_mat.size(); i++)
    matches.push_back(cv::DMatch(i, i, 0));
  
  // calculate transformation
  Ptr<ThinPlateSplineShapeTransformer> tps = cv::createThinPlateSplineShapeTransformer(0);
  tps->estimateTransformation(ref_mat, query_mat, matches);
  
  // determine extension limits for both images
  int y_max = max(imWarp.rows, imReference.rows);
  int x_max = max(imWarp.cols, imReference.cols);
  
  // extend images
  cv::copyMakeBorder(imWarp, imWarp, 0.0, (int) (y_max - imWarp.rows), 0.0, (x_max - imWarp.cols), cv::BORDER_CONSTANT, Scalar(0, 0, 0));
  
  // transform image
  cv::Mat imReg;
  tps->warpImage(imWarp, imReg);
  
  // resize image
  // cv::resize(im1Reg, im1Reg, im2.size());
  cv::Mat imReg_cropped  = imReg(cv::Range(0,imReference.size().height), cv::Range(0,imReference.size().width));
  imReg = imReg_cropped.clone();
  
  // return
  return matToImage(imReg);;
}