#include <iostream>
#include <math.h>
#include <cmath>
#include <algorithm>
#include <ros/ros.h>
#include <ros/console.h>
#include <sensor_msgs/Image.h>
#include <cv_bridge/cv_bridge.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>

#include <aruco/aruco.h>
#include <aruco/cvdrawingutils.h>

#include <opencv2/opencv.hpp>

#include <Eigen/Eigen>
#include <Eigen/SVD>
//Eigen SVD libnary, may help you solve SVD
//JacobiSVD<MatrixXd> svd(A, ComputeThinU | ComputeThinV);
#include <Eigen/Geometry>
#include <opencv2/core/eigen.hpp>


using namespace cv;
using namespace aruco;
using namespace Eigen;
using namespace std;
// global variables
double reproj_error;
MatrixXd cum_error = MatrixXd::Zero(6, 1);
int count_frame = 0;
int error_count = 0;

// Camera parameters
aruco::CameraParameters CamParam;
cv::Mat K, D;

// Aruco marker and detector
MarkerDetector MDetector;
vector<Marker> Markers;
float MarkerSize = 0.20 / 1.5 * 1.524;
float MarkerWithMargin = MarkerSize * 1.2;
BoardConfiguration TheBoardConfig;
BoardDetector TheBoardDetector;
Board TheBoardDetected;

// Topic subscribers and publishers
ros::Subscriber sub_img;
ros::Publisher pub_path_ref, pub_path;
nav_msgs::Path path_ref, path;

// This function publishes the pose results as a path.
void publishPath(
  const ros::Time& t,
  Vector3d& T,
  Quaterniond& Q,
  nav_msgs::Path& path,
  ros::Publisher& pub_path)
{
  geometry_msgs::PoseStampedPtr ps_ptr(new geometry_msgs::PoseStamped());
  ps_ptr->header.stamp = t;
  ps_ptr->header.frame_id = "world";
  ps_ptr->pose.position.x = T(0);
  ps_ptr->pose.position.y = -T(1);
  ps_ptr->pose.position.z = T(2);
  ps_ptr->pose.orientation.x = Q.x();
  ps_ptr->pose.orientation.y = Q.y();
  ps_ptr->pose.orientation.z = Q.z();
  ps_ptr->pose.orientation.w = Q.w();

  path.header.stamp = t;
  path.header.frame_id = ps_ptr->header.frame_id;
  path.poses.push_back(*ps_ptr);
  pub_path.publish(path);
}
// Obtain the
cv::Point3f getPositionFromIndex(int idx, int nth)
{
  int idx_x = idx % 6, idx_y = idx / 6;
  double p_x = idx_x * MarkerWithMargin - (3 + 2.5 * 0.2) * MarkerSize;
  double p_y = idx_y * MarkerWithMargin - (12 + 11.5 * 0.2) * MarkerSize;
  return cv::Point3f(p_x + (nth == 1 || nth == 2) * MarkerSize, p_y + (nth == 2 || nth == 3) * MarkerSize, 0.0);
}

/* THIS IS WHAT YOU NEED TO WORK WITH */
void calculateReprojectionError(
  const vector<cv::Point3f> &pts_3, // the input 3D points
  const vector<cv::Point2f> &pts_2, // the input 2D features that are corresponding to the 3D points
  const cv::Mat R, // the under-estimated rotation matrix
  const cv::Mat tt) // the under-estimated translation
{
    double error=0;
    for(int i=0;i<(int)pts_2.size();i++)
    {
        cv::Mat pts3=(cv::Mat_<double>(3,1)<<pts_3[i].x,pts_3[i].y,pts_3[i].z);
        cv::Mat pts2=(cv::Mat_<double>(2,1)<<pts_2[i].x,pts_2[i].y);
        cv::Mat pts1=R*pts3+tt;
        pts1=pts1/pts1.at<double>(2);
        cv::Mat pts=(cv::Mat_<double>(2,1)<<pts1.at<double>(0),pts1.at<double>(1));
        MatrixXd p;
        cv2eigen(pts2-pts,p);
        error=error+pow((pts_2[i].x-pts1.at<double>(0)),2)+pow((pts_2[i].y-pts1.at<double>(1)),2);
        reproj_error=sqrt(error);
    }
//    std::cout << "res error" << reproj_error << std::endl;
}

long int times1,times2 = 0;
double ref_error,res_error = 0;






Eigen::MatrixXd Jacobian(double x,double y,double z,double a,double b,double c ,double t1,double t2,double t3)
{
MatrixXd J(2, 6);
J(0,0)= ((x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a))*(t1 + z*sin(b) + x*cos(b)*cos(c) - y*cos(b)*sin(c)))/
        pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);

J(0,1)=- (z*cos(b) - x*cos(c)*sin(b) + y*sin(b)*sin(c))/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)) - ((z*cos(a)*sin(b) + x*cos(a)*cos(b)*cos(c) - y*cos(a)*cos(b)*sin(c))*(t1 + z*sin(b) + x*cos(b)*cos(c) - y*cos(b)*sin(c)))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);

J(0,2)=(y*cos(b)*cos(c) + x*cos(b)*sin(c))/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)) + ((x*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) - y*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)))*(t1 + z*sin(b) + x*cos(b)*cos(c) - y*cos(b)*sin(c)))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);


J(0,3)=-1/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b));
J(0,4)=0;
J(0,5)=(t1 + z*sin(b) + x*cos(b)*cos(c) - y*cos(b)*sin(c))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);

J(1,0)=(x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b))/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)) + ((x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a))*(t2 + x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a)))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);


J(1,1)=- (z*sin(a)*sin(b) + x*cos(b)*cos(c)*sin(a) - y*cos(b)*sin(a)*sin(c))/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)) - ((z*cos(a)*sin(b) + x*cos(a)*cos(b)*cos(c) - y*cos(a)*cos(b)*sin(c))*(t2 + x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a)))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);

J(1,2)=((x*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) - y*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)))*(t2 + x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a)))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2) - (x*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - y*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)))/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b));
J(1,3)=0;
J(1,4)=-1/(t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b));
J(1,5)=(t2 + x*(cos(a)*sin(c) + cos(c)*sin(a)*sin(b)) + y*(cos(a)*cos(c) - sin(a)*sin(b)*sin(c)) - z*cos(b)*sin(a))/pow((t3 + x*(sin(a)*sin(c) - cos(a)*cos(c)*sin(b)) + y*(cos(c)*sin(a) + cos(a)*sin(b)*sin(c)) + z*cos(a)*cos(b)),2);
return J;
}




void process(
  const vector<int> &pts_id,
  const vector<cv::Point3f> &pts_3,
  const vector<cv::Point2f> &pts_2,
  const ros::Time& frame_time) {
    cv::Mat r, rvec, t;
    cv::solvePnPRansac(pts_3, pts_2, K, D, rvec, t);
    //cv::solvePnP(pts_3, pts_2, K, D, rvec, t);
    cv::Rodrigues(rvec, r);
    Eigen::Matrix3d R_ref;
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            R_ref(i, j) = r.at<double>(i, j);
        }
    }
   Eigen::Matrix3d R_ref_inv = R_ref.transpose();
    Quaterniond Q_ref;
    Q_ref = R_ref_inv;
   Eigen::Vector3d t_ref(0, 0, 0);
    for (size_t i = 0; i < 3; i++)
        t_ref(i) = t.at<double>(i, 0);
    Eigen::Vector3d t_ref_inv = R_ref_inv * t_ref;

    publishPath(frame_time, t_ref_inv, Q_ref, path_ref, pub_path_ref);
    vector<cv::Point2f> un_pts_2;

    cv::undistortPoints(pts_2, pts_2, K, D);



    
    const size_t size_pts = pts_2.size();
    MatrixXd A(2 * size_pts, 9);
    for (int i = 0; i < (int) size_pts; i++)
    {
        VectorXd u(9);
        VectorXd v(9);
        cv::Point3d world = pts_3[i];
        cv::Point2d image = pts_2[i];
        u << world.x, world.y, 1, 0, 0, 0, -world.x * image.x, -world.y * image.x, -image.x;
        v << 0, 0, 0, world.x, world.y, 1, -world.x * image.y, -world.y * image.y, -image.y;
        A.block<1, 9>(2 * i, 0) = u.transpose();
        A.block<1, 9>(2 * i + 1, 0) = v.transpose();
    }
    JacobiSVD<MatrixXd> svd(A, ComputeThinU | ComputeThinV);
    VectorXd h(9);
    h = svd.matrixV().col(8);
    Matrix3d H_hat;
    H_hat << h(0), h(1), h(2),
            h(3), h(4), h(5),
            h(6), h(7), h(8);
    Matrix3d k_eigen;
    cv2eigen(K, k_eigen);
    Vector3d h3 = H_hat.block(0, 2, 3, 1);
    if (h3(2) < 0)
    {
        H_hat = -H_hat;
    }
    Vector3d h1 = H_hat.block(0, 0, 3, 1);
    Vector3d h2 = H_hat.block(0, 1, 3, 1);
    h3 = H_hat.block(0, 2, 3, 1);
    Matrix3d H_solve;
    H_solve.block<3, 1>(0, 0) = h1;
    H_solve.block<3, 1>(0, 1) = h2;
    H_solve.block<3, 1>(0, 2) = h1.cross(h2);
    JacobiSVD<MatrixXd> svd_H(H_solve, ComputeThinU | ComputeThinV);
    Matrix3d R_H;
    cv::Mat R;
    R_H = svd_H.matrixU() * (svd_H.matrixV().transpose());
    eigen2cv(R_H, R);
    MatrixXd t_H(3, 1);
    cv::Mat tt;
    
    t_H = h3 / h1.norm();
    eigen2cv(t_H, tt);
    calculateReprojectionError(pts_3, pts_2, R, tt);
    std::cout << "error_1:  " << reproj_error << std::endl;

    
    MatrixXd gamma(2, size_pts);
    MatrixXd JacobianMatrix[size_pts];
    MatrixXd angle_trans(6, 1);
    int iter_times = 50;
    double a,b,c,t1,t2,t3,x,y,z;
    Vector3d euler_angles = R_H.eulerAngles(0, 1, 2);
    a = euler_angles(0);
    b = euler_angles(1);
    c = euler_angles(2);
    t1 = t_H(0, 0);
    t2 = t_H(1, 0);
    t3 = t_H(2, 0);
    angle_trans(0, 0) = a;
    angle_trans(1, 0) = b;
    angle_trans(2, 0) = c;
    angle_trans(3, 0) = t1;
    angle_trans(4, 0) = t2;
    angle_trans(5, 0) = t3;
    MatrixXd J(2, 6);
    Mat R_init,t_init;
    eigen2cv(R_H, R_init);
    eigen2cv(t_H, t_init);
    Matrix3d R_out;
    Vector3d t_out;
    for(int count=0;count<iter_times ;count++)
    {
        MatrixXd An = MatrixXd::Zero(6,6);
        MatrixXd bn = MatrixXd::Zero(6,1);
        for (int i=0; i<(int)size_pts;i++)
        {
            cv::Mat pts3 = (cv::Mat_<double>(3, 1) << pts_3[i].x, pts_3[i].y, pts_3[i].z);
            cv::Mat p = R_init * pts3 + t_init;
            p = p / p.at<double>(2);
            gamma(0,i) = pts_2[i].x - p.at<double>(0);
            gamma(1,i) = pts_2[i].y - p.at<double>(1);
            x = pts_3[i].x;
            y = pts_3[i].y;
            z = pts_3[i].z;
            J = Jacobian(x,y,z,a,b,c,t1,t2,t3);
            JacobianMatrix[i] = J;
            An += J.transpose()*J;
            bn -= J.transpose()*gamma.col(i);
        }
        MatrixXd delta = An.inverse()*bn*0.35;
        if (delta.cwiseAbs().sum() < 0.0001 )
        {
            break;
        }

        angle_trans = angle_trans+delta;

        R_out =  Eigen::AngleAxisd(angle_trans(0), Eigen::Vector3d::UnitX()) * Eigen::AngleAxisd(angle_trans(1), Eigen::Vector3d::UnitY()) * Eigen::AngleAxisd(angle_trans(2), Eigen::Vector3d::UnitZ());
        t_out(0)= angle_trans(3);
        t_out(1)= angle_trans(4);
        t_out(2)= angle_trans(5);
        a = angle_trans(0);
        b = angle_trans(1);
        c = angle_trans(2);
        t1=angle_trans(3);
        t2=angle_trans(4);
        t3=angle_trans(5);
        eigen2cv(R_out,R_init);
        eigen2cv(t_out,t_init);
    }
  Eigen::Matrix3d R_out_inv = R_out.transpose();


    Vector3d T=R_out_inv * t_out;
    Quaterniond Q_yourwork;
    Q_yourwork = R_out;
    publishPath(frame_time, T, Q_yourwork, path, pub_path);
    /* For quantitative evaluation */
    double diff_psi = R_ref.eulerAngles(2, 0, 1)(0) - R_out.eulerAngles(2, 0, 1)(0);
    if (diff_psi > M_PI)
    {
        diff_psi -= 2 * M_PI;
    }
    else if (diff_psi < - M_PI)
    {
        diff_psi += 2 * M_PI;
    }

    double diff_phi = R_ref.eulerAngles(2, 0, 1)(1) - R_out.eulerAngles(2, 0, 1)(1);
    if (diff_phi > M_PI)
    {
        diff_phi -= 2 * M_PI;
    }
    else if (diff_phi < -M_PI)
    {
        diff_phi += 2 * M_PI;
    }

    double diff_theta = R_ref.eulerAngles(2, 0, 1)(2) - R_out.eulerAngles(2, 0, 1)(2);
    if (diff_theta > M_PI)
    {
        diff_theta -= 2 * M_PI;
    }
    else if (diff_theta < -M_PI)
    {
        diff_theta += 2 * M_PI;
    }
    if(abs(T(0))<1e20&&abs(T(1))<1e20&&abs(T(2))<1e20)
    {
        count_frame = count_frame + 1;
        cum_error(0, 0) = cum_error(0, 0) + pow(diff_psi, 2);
        cum_error(1, 0) = cum_error(1, 0) + pow(diff_phi, 2);
        cum_error(2, 0) = cum_error(2, 0) + pow(diff_theta, 2);
        cum_error(3, 0) = cum_error(3, 0) + pow(t.at<double>(0, 0) - T(0), 2);
        cum_error(4, 0) = cum_error(4, 0) + pow(t.at<double>(1, 0) - T(1), 2);
        cum_error(5, 0) = cum_error(5, 0) + pow(t.at<double>(2, 0) - T(2), 2);
    }
    ROS_INFO("RMSE X, Y, Z, roll, pitch, yaw: \n %f, %f, %f, %f, %f, %f",
             sqrt(cum_error(3, 0) / count_frame), sqrt(cum_error(4, 0) / count_frame), sqrt(cum_error(5, 0) / count_frame),
             sqrt(cum_error(1, 0) / count_frame), sqrt(cum_error(2, 0) / count_frame), sqrt(cum_error(0, 0) / count_frame));

    cv::Mat R_res_cv(3, 3, CV_64FC1);
    cv::eigen2cv(R_out, R_res_cv);

    cv::Mat t_res_cv(3, 1, CV_64FC1);
    cv::eigen2cv(T, t_res_cv);
    
    calculateReprojectionError(pts_3, pts_2, r, t);
    times1 ++;
    std::cout << "ref error:  " << reproj_error << std::endl;
    ref_error +=reproj_error;
    calculateReprojectionError(pts_3,pts_2,R_res_cv,t_res_cv);
    times2 ++;
    std::cout << "res error:  " << reproj_error << std::endl;
    cout<<"================================================================"<<endl;
    res_error += reproj_error;  
    cout<<"average ref_error: "<<ref_error/times1<<endl;
    cout<<"average res_error: "<<res_error/times2<<endl;
}

// Callback function for processing the incoming images.
//
void img_callback(const sensor_msgs::ImageConstPtr &img_msg)
{
  /* Detect markers and obtain 3D-2D data association */
  double t = clock();
  cv_bridge::CvImagePtr bridge_ptr = cv_bridge::toCvCopy(img_msg, sensor_msgs::image_encodings::MONO8);
  MDetector.detect(bridge_ptr->image, Markers);
  float probDetect = TheBoardDetector.detect(Markers, TheBoardConfig, TheBoardDetected, CamParam, MarkerSize);
  ROS_DEBUG("p: %f, time cost: %f\n", probDetect, (clock() - t) / CLOCKS_PER_SEC);

  vector<int> pts_id;
  vector<cv::Point3f> pts_3;
  vector<cv::Point2f> pts_2;
  for (unsigned int i = 0; i < Markers.size(); i++)
  {
    // Obtain the ID of the marker
    int idx = TheBoardConfig.getIndexOfMarkerId(Markers[i].id);

    char str[100];
    sprintf(str, "%d", idx);
    cv::putText(bridge_ptr->image, str, Markers[i].getCenter(), CV_FONT_HERSHEY_COMPLEX, 0.5, cv::Scalar(-1));
    for (unsigned int j = 0; j < 4; j++)
    {
      sprintf(str, "%d", j);
      cv::putText(bridge_ptr->image, str, Markers[i][j], CV_FONT_HERSHEY_COMPLEX, 0.5, cv::Scalar(0,0,1));
    }
    // j here is index of four corners
    for (unsigned int j = 0; j < 4; j++)
    {
      pts_id.push_back(Markers[i].id * 4 + j);
      pts_3.push_back(getPositionFromIndex(idx, j));
      pts_2.push_back(Markers[i][j]);
    }
  }

  /* Call your pose estimation function */
  if (pts_id.size() > 5)
    process(pts_id, pts_3, pts_2, img_msg->header.stamp);

  /* Render the detection result on the raw image */
  cv::imshow("in", bridge_ptr->image);
  cv::waitKey(10);
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "tag_detector");
    ros::NodeHandle n("~");

    sub_img = n.subscribe("image_raw", 100, img_callback);
    pub_path_ref = n.advertise<nav_msgs::Path>("path_ref", 10);
    pub_path = n.advertise<nav_msgs::Path>("path_yourwork", 10);
    //init aruco detector
    string cam_cal, board_config;
    n.getParam("cam_cal_file", cam_cal);
    n.getParam("board_config_file", board_config);
    CamParam.readFromXMLFile(cam_cal);
    TheBoardConfig.readFromFile(board_config);

    //init intrinsic parameters
    cv::FileStorage param_reader(cam_cal, cv::FileStorage::READ);
    param_reader["camera_matrix"] >> K;
    param_reader["distortion_coefficients"] >> D;

    //init window for visualization
    cv::namedWindow("in", 0);
    ros::spin();
}