#include <Eigen/Core>
#include <iostream>

#include <opencv2/opencv.hpp>

#include "g2o/stuff/sampler.h"
#include "g2o/stuff/command_args.h"
#include "g2o/core/sparse_optimizer.h"
#include "g2o/core/block_solver.h"
#include "g2o/core/solver.h"
#include "g2o/core/optimization_algorithm_levenberg.h"
#include "g2o/core/base_vertex.h"
#include "g2o/core/base_unary_edge.h"
#include "g2o/solvers/dense/linear_solver_dense.h"

using namespace std;

class VertexHomographyParams : public g2o::BaseVertex<8, Eigen::VectorXd> 
{
  public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
    VertexHomographyParams()
    {
    }

    virtual bool read(std::istream& /*is*/)
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    virtual bool write(std::ostream& /*os*/) const
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    virtual void setToOriginImpl()
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
    }

    virtual void oplusImpl(const double* update)
    {
      Eigen::VectorXd::ConstMapType v(update, 8);
      // Eigen::VectorXd::ConstMapType v(update, VertexHomographyParams::Dimension);
      _estimate += v;
      // cout << "============================" << endl;
      // for(int i = 0; i < 8; i++) {
        // cout << update[i] << endl;
        // _estimate(i) += update[i];
      // }
    }
};

class EdgeSamePoint : public g2o::BaseUnaryEdge<2, Eigen::VectorXd, VertexHomographyParams>
{
  public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    EdgeSamePoint()
    {
    }
    virtual bool read(std::istream& /*is*/)
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }
    virtual bool write(std::ostream& /*os*/) const
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    void computeError()
    {
      const VertexHomographyParams* param = static_cast<const VertexHomographyParams*>(vertex(0));
      // const double h = param->estimate();
      const Eigen::VectorXd h = param->estimate();

      cv::Point2d src_point;
      src_point.x = measurement()(0);
      src_point.y = measurement()(1);
      cv::Point2d dst_point;
      dst_point.x = measurement()(2);
      dst_point.y = measurement()(3);

      double projected_point_on_im2[3];
      {
        projected_point_on_im2[0] = h(0)*src_point.x + 
                                    h(1)*src_point.y + 
                                    h(2)*(1.);
        projected_point_on_im2[1] = h(3)*src_point.x + 
                                    h(4)*src_point.y + 
                                    h(5)*(1.);
        projected_point_on_im2[2] = h(6)*src_point.x + 
                                    h(7)*src_point.y + 
                                    1.0*(1.);
        /*
         *  projected_point_on_im2 * S = H * point_on_im1(homogenious <==> z=1)
         */
      }

      _error[0] = projected_point_on_im2[0]/projected_point_on_im2[2] - dst_point.x;
      _error[1] = projected_point_on_im2[1]/projected_point_on_im2[2] - dst_point.y;

      /*
      const VertexHomographyParams* params = static_cast<const VertexHomographyParams*>(vertex(0));
      const double& a = params->estimate()(0);
      const double& b = params->estimate()(1);
      const double& lambda = params->estimate()(2);
      double fval = a * exp(-lambda * measurement()(0)) + b;
      _error(0) = fval - measurement()(1);
      */
    }
};

int main(int argc, char** argv)
{
  if(argc != 3) {
    cout << 
      "usage: this.out [/path/to/image1] [path/to/image2] "
      << endl;
    return -1;
  }

  cv::Mat image1 = cv::imread(argv[1]);
  cv::Mat image2 = cv::imread(argv[2]);

  vector<cv::KeyPoint> kpts_vec1, kpts_vec2;
  cv::Mat desc1, desc2;
  cv::Ptr<cv::AKAZE> akaze = cv::AKAZE::create();

  // extract feature points and calculate descriptors
  akaze -> detectAndCompute(image1, cv::noArray(), kpts_vec1, desc1);
  akaze -> detectAndCompute(image2, cv::noArray(), kpts_vec2, desc2);


  cv::BFMatcher* matcher = new cv::BFMatcher(cv::NORM_L2, false); 
  // cross check flag set to false
  // because i do cross-ratio-test match
  vector< vector<cv::DMatch> > matches_2nn_12, matches_2nn_21;
  matcher->knnMatch( desc1, desc2, matches_2nn_12, 2 );
  matcher->knnMatch( desc2, desc1, matches_2nn_21, 2 );
  const double ratio = 0.8;

  vector<cv::Point2f> selected_points1, selected_points2;

  for(int i = 0; i < matches_2nn_12.size(); i++) { // i is queryIdx
    if( matches_2nn_12[i][0].distance/matches_2nn_12[i][1].distance < ratio 
        and 
        matches_2nn_21[matches_2nn_12[i][0].trainIdx][0].distance 
          / matches_2nn_21[matches_2nn_12[i][0].trainIdx][1].distance < ratio ) 
    {
      if(matches_2nn_21[matches_2nn_12[i][0].trainIdx][0].trainIdx 
            == matches_2nn_12[i][0].queryIdx) 
      {
        selected_points1.push_back(kpts_vec1[matches_2nn_12[i][0].queryIdx].pt);
        selected_points2.push_back(
            kpts_vec2[matches_2nn_21[matches_2nn_12[i][0].trainIdx][0].queryIdx].pt
            );
      }
    }
  }

  Eigen::Vector2f* pf_src_points = new Eigen::Vector2f[selected_points1.size()];
  Eigen::Vector2f* pf_dst_points = new Eigen::Vector2f[selected_points2.size()];
  Eigen::Vector4d* points = new Eigen::Vector4d[selected_points1.size()];
  // selected_points1.size() == selected_points2.size()

  for(int i = 0; i < selected_points2.size(); i++) {
    pf_src_points[i].x() = selected_points1[i].x;
    pf_src_points[i].y() = selected_points1[i].y;
    pf_dst_points[i].x() = selected_points2[i].x;
    pf_dst_points[i].y() = selected_points2[i].y;
    points[i](0) = (double)selected_points1[i].x;
    points[i](1) = (double)selected_points1[i].y;
    points[i](2) = (double)selected_points2[i].x;
    points[i](3) = (double)selected_points2[i].y;
  }

  if(false) {
     cv::Mat src;
     cv::hconcat(image1, image2, src);
     for(int i = 0; i < selected_points1.size(); i++) {
       cv::line( src, selected_points1[i],
                 cv::Point2f(selected_points2[i].x + image1.cols, selected_points2[i].y),
                 1, 1, 0 );
     }
     cv::imshow("match-result.png", src);
     cv::waitKey(0);
  }

  // some handy typedefs
  typedef g2o::BlockSolver< g2o::BlockSolverTraits<Eigen::Dynamic, Eigen::Dynamic> >  MyBlockSolver;
  typedef g2o::LinearSolverDense<MyBlockSolver::PoseMatrixType> MyLinearSolver;

  // setup the solver
  g2o::SparseOptimizer optimizer;
  optimizer.setVerbose(false);

  g2o::OptimizationAlgorithmLevenberg* solver = new g2o::OptimizationAlgorithmLevenberg(
    g2o::make_unique<MyBlockSolver>(g2o::make_unique<MyLinearSolver>()));

  optimizer.setAlgorithm(solver);

  VertexHomographyParams* h = new VertexHomographyParams();
  h->setId(0);
  h->setEstimate(Eigen::VectorXd::Ones(8));
  optimizer.addVertex(h);

  Eigen::MatrixXd info = Eigen::Matrix<double,2,2>::Identity();

  cout << info << endl;

  for(int i = 0; i < selected_points1.size(); i++) {
    EdgeSamePoint* e = new EdgeSamePoint();
    e->setInformation(info);
    e->setVertex(0, h);
    e->setMeasurement(points[i]);
    optimizer.addEdge(e);
  }

  // perform the optimization
  optimizer.initializeOptimization();
  optimizer.setVerbose(true);
  optimizer.optimize(200);

  cout << "[g2o]H =\n"
       << h->estimate()(0) << " " << h->estimate()(1) << " " << h->estimate()(2) << endl
       << h->estimate()(3) << " " << h->estimate()(4) << " " << h->estimate()(5) << endl
       << h->estimate()(6) << " " << h->estimate()(7) << " " << 1.0 << endl;

  cv::Mat H = (cv::Mat_<double>(3,3) << h->estimate()(0),h->estimate()(1),h->estimate()(2),
                                        h->estimate()(3),h->estimate()(4),h->estimate()(5),
                                        h->estimate()(6),h->estimate()(7),1.0);

  cv::Mat m;
  cv::warpPerspective( image2, m, H, image2.size());
  cv::imshow("test", m);
  cv::waitKey(0);

  cv::Mat homography =  findHomography(selected_points1, selected_points2, 0, 3);
  cv::warpPerspective( image2, m, homography, image2.size());
  cv::imshow("test1", m);
  cv::waitKey(0);

  cout << "OpenCV H = \n" << homography << endl;

  cv::destroyAllWindows();

  delete[] points;

  return 0;
}