//Eigen
#include <Eigen/Core>

//OpenCV
#include <opencv2/core/core.hpp>
#include <opencv2/core/eigen.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/features2d/features2d.hpp>
#include <opencv2/calib3d.hpp>
#include <opencv2/imgproc.hpp>

//PCL
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/visualization/common/common.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/filters/statistical_outlier_removal.h>

#include <chrono>
#include <iostream>

#include "estimate_motion.h"

using namespace p3dv;

bool MotionEstimator::estimate2D2D_E5P_RANSAC(frame_t &cur_frame_1, frame_t &cur_frame_2,
                                              std::vector<cv::DMatch> &matches, std::vector<cv::DMatch> &inlier_matches,
                                              Eigen::Matrix4f &T, double ransac_thre, double ransac_prob, bool show)
{
    std::chrono::steady_clock::time_point tic = std::chrono::steady_clock::now();

    std::vector<cv::Point2f> pointset1;
    std::vector<cv::Point2f> pointset2;

    for (int i = 0; i < (int)matches.size(); i++)
    {
        pointset1.push_back(cur_frame_1.keypoints[matches[i].queryIdx].pt);
        pointset2.push_back(cur_frame_2.keypoints[matches[i].trainIdx].pt);
    }

    cv::Mat camera_mat;
    cv::eigen2cv(cur_frame_1.K_cam, camera_mat);

    cv::Mat essential_matrix;

    cv::Mat inlier_matches_indicator;

    essential_matrix = cv::findEssentialMat(pointset1, pointset2, camera_mat,
                                            CV_RANSAC, ransac_prob, ransac_thre, inlier_matches_indicator);

    std::cout << "essential_matrix is " << std::endl
              << essential_matrix << std::endl;

    for (int i = 0; i < (int)matches.size(); i++)
    {
        if (inlier_matches_indicator.at<bool>(0, i) == 1)
        {
            inlier_matches.push_back(matches[i]);
        }
    }

    cv::Mat R;
    cv::Mat t;

    cv::recoverPose(essential_matrix, pointset1, pointset2, camera_mat, R, t, inlier_matches_indicator);

    std::chrono::steady_clock::time_point toc = std::chrono::steady_clock::now();
    std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(toc - tic);
    std::cout << "Estimate Motion [2D-2D] cost = " << time_used.count() << " seconds. " << std::endl;
    std::cout << "Find [" << inlier_matches.size() << "] inlier matches from [" << matches.size() << "] total matches." << std::endl;

    Eigen::Matrix3f R_eigen;
    Eigen::Vector3f t_eigen;
    cv::cv2eigen(R, R_eigen);
    cv::cv2eigen(t, t_eigen);
    T.block(0, 0, 3, 3) = R_eigen;
    T.block(0, 3, 3, 1) = t_eigen;
    T(3, 0) = 0;
    T(3, 1) = 0;
    T(3, 2) = 0;
    T(3, 3) = 1;

    std::cout << "Transform is " << std::endl
              << T << std::endl;

    if (show)
    {
        cv::Mat ransac_match_image;
        cv::namedWindow("RANSAC inlier matches", 0);
        cv::drawMatches(cur_frame_1.rgb_image, cur_frame_1.keypoints, cur_frame_2.rgb_image, cur_frame_2.keypoints, inlier_matches, ransac_match_image);
        cv::imshow("RANSAC inlier matches", ransac_match_image);
        cv::waitKey(0);
    }

    return 1;
}

bool MotionEstimator::estimate2D3D_P3P_RANSAC(frame_t &cur_frame, pointcloud_sparse_t &cur_map_3d, double ransac_thre,
                                              int iterationsCount, double ransac_prob, bool show)
{
    std::chrono::steady_clock::time_point tic = std::chrono::steady_clock::now();

    cv::Mat camera_mat;
    cv::eigen2cv(cur_frame.K_cam, camera_mat);

    //std::cout<< "K Mat:" <<std::endl<< camera_mat<<std::endl;

    std::vector<cv::Point2f> pointset2d;
    std::vector<cv::Point3f> pointset3d;
    std::vector<int> pointset3d_index;

    int count = 0;
    float dist_thre = 300;

    std::cout << "2D points: " << cur_frame.unique_pixel_ids.size() << std::endl
              << "3D points: " << cur_map_3d.unique_point_ids.size() << std::endl;

    // Construct the 2d-3d initial matchings
    for (int i = 0; i < cur_frame.unique_pixel_ids.size(); i++)
    {
        for (int j = 0; j < cur_map_3d.unique_point_ids.size(); j++)
        {
            if (cur_frame.unique_pixel_ids[i] == cur_map_3d.unique_point_ids[j])
            {

                float x_3d = cur_map_3d.rgb_pointcloud->points[j].x;
                float y_3d = cur_map_3d.rgb_pointcloud->points[j].y;
                float z_3d = cur_map_3d.rgb_pointcloud->points[j].z;
                float x_2d = cur_frame.keypoints[i].pt.x;
                float y_2d = cur_frame.keypoints[i].pt.y;

                if (std::abs(x_3d) < dist_thre && std::abs(y_3d) < dist_thre && std::abs(z_3d) < dist_thre)
                {
                    // Assign value for pointset2d and pointset3d
                    pointset2d.push_back(cv::Point2f(x_2d, y_2d));
                    //pointset2d.push_back(pixel2cam(cur_frame.keypoints[i].pt, camera_mat));

                    pointset3d.push_back(cv::Point3f(x_3d, y_3d, z_3d));
                    pointset3d_index.push_back(j);

                    count++;
                }
                //std::cout << "2D: " << cur_frame.unique_pixel_ids[i] << " " << cur_frame.keypoints[i].pt << std::endl;
                //std::cout << "3D: " << cur_map_3d.unique_point_ids[j] << " " << cv::Point3f(x_3d, y_3d, z_3d) << std::endl;
            }
        }
    }

    std::cout << count << " initial correspondences are used." << std::endl;

    // Use RANSAC P3P to estiamte the optimal transformation

    cv::Mat distort_para = cv::Mat::zeros(1, 4, CV_64FC1); // Assuming no lens distortion
    cv::Mat r_vec;
    cv::Mat t_vec;
    cv::Mat inliers;

    cv::solvePnPRansac(pointset3d, pointset2d, camera_mat, distort_para, r_vec, t_vec,
                       false, iterationsCount, ransac_thre, ransac_prob, inliers, cv::SOLVEPNP_EPNP);

    //cv::solvePnP(pointset3d, pointset2d, camera_mat, distort_para, r_vec, t_vec, false, cv::SOLVEPNP_EPNP);

    std::cout << "Inlier count: " << inliers.rows << std::endl;

    std::vector<int> outlier;
    for (int i = 0; i < inliers.rows; i++)
    {
        pointset3d_index[inliers.at<float>(i, 0)] = -1; // inlier 's index
    }

    for (int i = 0; i < pointset3d_index.size(); i++)
    {
        if (pointset3d_index[i] >= 0)
            outlier.push_back(pointset3d_index[i]);
    }
    for (int i = 0; i < outlier.size(); i++)
    {
        cur_map_3d.is_inlier[outlier[i]] = 0;
    }

    cv::Mat R_mat;
    cv::Rodrigues(r_vec, R_mat);

    Eigen::Matrix3f R_eigen;
    Eigen::Vector3f t_eigen;
    Eigen::Matrix4f T_mat;

    cv::cv2eigen(R_mat, R_eigen);
    cv::cv2eigen(t_vec, t_eigen);

    T_mat.block(0, 0, 3, 3) = R_eigen;
    T_mat.block(0, 3, 3, 1) = t_eigen;
    T_mat(3, 0) = 0;
    T_mat(3, 1) = 0;
    T_mat(3, 2) = 0;
    T_mat(3, 3) = 1;

    // std::cout << "Transform is: " << std::endl
    //           << T_mat << std::endl;

    cur_frame.pose_cam = T_mat;

    // Calculate the reprojection error
    std::vector<cv::Point2f> proj_points;
    cv::projectPoints(pointset3d, R_mat, t_vec, camera_mat, distort_para, proj_points);

    float reproj_err = 0.0;
    for (int i = 0; i < proj_points.size(); i++)
    {
        float cur_repro_error = norm(proj_points[i] - pointset2d[i]);
        reproj_err += cur_repro_error;

        //std::cout << cur_repro_error << std::endl;
    }

    reproj_err /= proj_points.size();
    double inlier_ratio = 1.0 * inliers.rows / count;
    std::cout << "Mean reprojection error: " << reproj_err << std::endl;

    std::chrono::steady_clock::time_point toc = std::chrono::steady_clock::now();
    std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(toc - tic);
    std::cout << "Estimate Motion [3D-2D] cost = " << time_used.count() << " seconds. " << std::endl;

    if (reproj_err > 10 && inlier_ratio < 0.5) // mean reprojection error is too big (may be some problem)
    {
        std::cout << "[Warning] pnp may encounter some problem, the inlier ratio is [ " << inlier_ratio * 100 << " % ], the mean reprojection error is [ " << reproj_err << " ]." << std::endl;
        return 0;
    }
    else
        return 1;
}

bool MotionEstimator::getDepthFast(frame_t &cur_frame_1, frame_t &cur_frame_2, Eigen::Matrix4f &T_21,
                                   const std::vector<cv::DMatch> &matches, double &appro_depth, int random_rate)
{
    cv::Mat T1_mat;
    cv::Mat T2_mat;
    cv::Mat camera_mat;

    Eigen::Matrix4f Teye = Eigen::Matrix4f::Identity();
    Eigen::Matrix<float, 3, 4> T1 = Teye.block<3, 4>(0, 0);
    Eigen::Matrix<float, 3, 4> T2 = (T_21 * Teye).block<3, 4>(0, 0);

    cv::eigen2cv(cur_frame_1.K_cam, camera_mat);
    cv::eigen2cv(T1, T1_mat);
    cv::eigen2cv(T2, T2_mat);

    std::vector<cv::Point2f> pointset1;
    std::vector<cv::Point2f> pointset2;

    for (int i = 0; i < matches.size(); i++)
    {
        if (i % random_rate == 0)
        {
            pointset1.push_back(pixel2cam(cur_frame_1.keypoints[matches[i].queryIdx].pt, camera_mat));
            pointset2.push_back(pixel2cam(cur_frame_2.keypoints[matches[i].trainIdx].pt, camera_mat));
        }
    }

    cv::Mat pts_3d_homo;
    if (pointset1.size() > 0)
        cv::triangulatePoints(T1_mat, T2_mat, pointset1, pointset2, pts_3d_homo);

    // De-homo and calculate mean depth
    double depth_sum = 0;
    for (int i = 0; i < pts_3d_homo.cols; i++)
    {
        cv::Mat pts_3d = pts_3d_homo.col(i);

        pts_3d /= pts_3d.at<float>(3, 0);

        Eigen::Vector3f pt_temp;
        pt_temp(0) = pts_3d.at<float>(0, 0);
        pt_temp(1) = pts_3d.at<float>(1, 0);
        pt_temp(2) = pts_3d.at<float>(2, 0);

        depth_sum = depth_sum + pt_temp.norm();
    }
    appro_depth = depth_sum / pts_3d_homo.cols;

    std::cout << "Mean relative depth is about " << appro_depth << " * baseline length. " << std::endl;
}

bool MotionEstimator::doTriangulation(frame_t &cur_frame_1, frame_t &cur_frame_2,
                                      const std::vector<cv::DMatch> &matches,
                                      pointcloud_sparse_t &sparse_pointcloud, bool show)
{
    std::chrono::steady_clock::time_point tic = std::chrono::steady_clock::now();

    cv::Mat T1_mat;
    cv::Mat T2_mat;
    cv::Mat camera_mat;

    Eigen::Matrix<float, 3, 4> T1 = cur_frame_1.pose_cam.block(0, 0, 3, 4);
    Eigen::Matrix<float, 3, 4> T2 = cur_frame_2.pose_cam.block(0, 0, 3, 4);

    cv::eigen2cv(cur_frame_1.K_cam, camera_mat);
    cv::eigen2cv(T1, T1_mat);
    cv::eigen2cv(T2, T2_mat);

    // std::cout<<camera_mat<<std::endl;
    // std::cout<<T1_mat<<std::endl;
    // std::cout<<T2_mat<<std::endl;

    std::vector<cv::Point2f> pointset1;
    std::vector<cv::Point2f> pointset2;

    int count_newly_triangu = 0;
    for (int i = 0; i < matches.size(); i++)
    {
        bool already_in_world = 0;
        for (int k = 0; k < sparse_pointcloud.unique_point_ids.size(); k++)
        {
            if (sparse_pointcloud.unique_point_ids[k] == cur_frame_1.unique_pixel_ids[matches[i].queryIdx])
            {
                already_in_world = 1;
                break;
            }
        }
        if (!already_in_world)
        {
            sparse_pointcloud.unique_point_ids.push_back(cur_frame_1.unique_pixel_ids[matches[i].queryIdx]);
            sparse_pointcloud.is_inlier.push_back(1);

            pointset1.push_back(pixel2cam(cur_frame_1.keypoints[matches[i].queryIdx].pt, camera_mat));
            pointset2.push_back(pixel2cam(cur_frame_2.keypoints[matches[i].trainIdx].pt, camera_mat));

            count_newly_triangu++;
        }
    }

    cv::Mat pts_3d_homo;
    if (pointset1.size() > 0)
        cv::triangulatePoints(T1_mat, T2_mat, pointset1, pointset2, pts_3d_homo);

    // De-homo and assign color
    for (int i = 0; i < pts_3d_homo.cols; i++)
    {
        cv::Mat pts_3d = pts_3d_homo.col(i);

        pts_3d /= pts_3d.at<float>(3, 0);

        pcl::PointXYZRGB pt_temp;
        pt_temp.x = pts_3d.at<float>(0, 0);
        pt_temp.y = pts_3d.at<float>(1, 0);
        pt_temp.z = pts_3d.at<float>(2, 0);

        // check here if(pt_temp.x> )

        cv::Point2f cur_key_pixel = cur_frame_1.keypoints[matches[i].queryIdx].pt;

        uchar blue = cur_frame_1.rgb_image.at<cv::Vec3b>(cur_key_pixel.y, cur_key_pixel.x)[0];
        uchar green = cur_frame_1.rgb_image.at<cv::Vec3b>(cur_key_pixel.y, cur_key_pixel.x)[1];
        uchar red = cur_frame_1.rgb_image.at<cv::Vec3b>(cur_key_pixel.y, cur_key_pixel.x)[2];

        pt_temp.r = 1.0 * red;
        pt_temp.g = 1.0 * green;
        pt_temp.b = 1.0 * blue;
        sparse_pointcloud.rgb_pointcloud->points.push_back(pt_temp);
    }

    std::cout << "Triangulate [ " << count_newly_triangu << " ] new points, [ " << sparse_pointcloud.rgb_pointcloud->points.size() << " ] points in total." << std::endl;

    std::chrono::steady_clock::time_point toc = std::chrono::steady_clock::now();
    std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(toc - tic);
    std::cout << "Triangularization done in " << time_used.count() << " seconds. " << std::endl;

    if (show)
    {
        // Show 2D image pair and correspondences
        cv::Mat match_image_pair;
        cv::drawMatches(cur_frame_1.rgb_image, cur_frame_1.keypoints, cur_frame_2.rgb_image, cur_frame_2.keypoints, matches, match_image_pair);
        cv::imshow("Triangularization matches", match_image_pair);
        cv::waitKey(0);

        boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer(new pcl::visualization::PCLVisualizer("Sfm Viewer"));
        viewer->setBackgroundColor(0, 0, 0);

        // Draw camera
        char t[256];
        std::string s;
        int n = 0;
        float frame_color_r, frame_color_g, frame_color_b;
        float sphere_size = 0.2;
        float line_size_cam = 0.4;

        pcl::PointXYZ pt_cam1(cur_frame_1.pose_cam(0, 3), cur_frame_1.pose_cam(1, 3), cur_frame_1.pose_cam(2, 3));
        pcl::PointXYZ pt_cam2(cur_frame_2.pose_cam(0, 3), cur_frame_2.pose_cam(1, 3), cur_frame_2.pose_cam(2, 3));

        sprintf(t, "%d", n);
        s = t;
        viewer->addSphere(pt_cam1, sphere_size, 1.0, 0.0, 0.0, s);
        n++;

        sprintf(t, "%d", n);
        s = t;
        viewer->addSphere(pt_cam2, sphere_size, 0.0, 0.0, 1.0, s);
        n++;

        sprintf(t, "%d", n);
        s = t;
        viewer->addLine(pt_cam1, pt_cam2, 0.0, 1.0, 0.0, s);
        n++;

        // for (int i = 0; i < sparse_pointcloud->points.size(); i++)
        // {
        //     char sparse_point[256];
        //     pcl::PointXYZ ptc_temp;
        //     ptc_temp.x = sparse_pointcloud->points[i].x;
        //     ptc_temp.y = sparse_pointcloud->points[i].y;
        //     ptc_temp.z = sparse_pointcloud->points[i].z;
        //     sprintf(sparse_point, "SP_%03u", i);
        //     viewer->addSphere(ptc_temp, 0.2, 1.0, 0.0, 0.0, sparse_point);
        // }

        viewer->addPointCloud(sparse_pointcloud.rgb_pointcloud, "sparsepointcloud");

        std::cout << "Click X(close) to continue..." << std::endl;
        while (!viewer->wasStopped())
        {
            viewer->spinOnce(100);
            boost::this_thread::sleep(boost::posix_time::microseconds(100000));
        }
    }

    std::cout << "Generate new sparse point cloud done." << std::endl;
    return true;
}

bool MotionEstimator::doUnDistort(frame_t &cur_frame, cv::Mat distort_coeff)
{
    cv::Mat camera_mat;
    cv::eigen2cv(cur_frame.K_cam, camera_mat);
    cv::Mat undistorted_img;
    cv::undistort(cur_frame.rgb_image, undistorted_img, camera_mat, distort_coeff);
    cur_frame.rgb_image = undistorted_img;

    std::cout << "Undistort the image done." << std::endl;
    return true;
}

/**
* \brief Transform a Point Cloud using a given transformation matrix
* \param[in]  cloud_in : A pointer of the Point Cloud before transformation
* \param[out] cloud_out : A pointer of the Point Cloud after transformation
* \param[in]  trans : A 4*4 transformation matrix
*/
bool MotionEstimator::transformCloud(pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud_in,
                                     pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud_out,
                                     Eigen::Matrix4f &trans)
{
    Eigen::Matrix4Xf PC;
    Eigen::Matrix4Xf TPC;
    PC.resize(4, cloud_in->size());
    TPC.resize(4, cloud_in->size());
    for (int i = 0; i < cloud_in->size(); i++)
    {
        PC(0, i) = cloud_in->points[i].x;
        PC(1, i) = cloud_in->points[i].y;
        PC(2, i) = cloud_in->points[i].z;
        PC(3, i) = 1;
    }
    TPC = trans * PC;
    for (int i = 0; i < cloud_in->size(); i++)
    {
        pcl::PointXYZ pt;
        pt.x = TPC(0, i);
        pt.y = TPC(1, i);
        pt.z = TPC(2, i);
        cloud_out->points.push_back(pt);
    }
    //cout << "Transform done ..." << endl;
}

bool MotionEstimator::outlierFilter(pointcloud_sparse_t &sparse_pointcloud, int MeanK, double std)
{
    // Create the filtering object
    pcl::StatisticalOutlierRemoval<pcl::PointXYZRGB> sor;

    std::vector<int> filtered_indice;

    sor.setInputCloud(sparse_pointcloud.rgb_pointcloud);
    sor.setMeanK(MeanK);         //50
    sor.setStddevMulThresh(std); //1.0
    sor.filter(filtered_indice);

    pcl::PointCloud<pcl::PointXYZRGB>::Ptr output_pointcloud(new pcl::PointCloud<pcl::PointXYZRGB>);
    std::vector<int> output_unique_points_id;
    std::vector<int> output_is_inlier;
    for (int i = 0; i < filtered_indice.size(); i++)
    {
        output_pointcloud->points.push_back(sparse_pointcloud.rgb_pointcloud->points[filtered_indice[i]]);
        output_unique_points_id.push_back(sparse_pointcloud.unique_point_ids[filtered_indice[i]]);
        output_is_inlier.push_back(sparse_pointcloud.is_inlier[filtered_indice[i]]);
    }

    std::cout << "apply outlier filter: [ " << sparse_pointcloud.unique_point_ids.size() << " ] points before filtering, [ " << filtered_indice.size() << " ] points after filtering." << std::endl;

    sparse_pointcloud.rgb_pointcloud->points.swap(output_pointcloud->points);
    sparse_pointcloud.unique_point_ids.swap(output_unique_points_id);
    sparse_pointcloud.is_inlier.swap(output_is_inlier);

    return 1;
}

bool MotionEstimator::estimateE8Points(std::vector<cv::KeyPoint> &keypoints1,
                                       std::vector<cv::KeyPoint> &keypoints2,
                                       std::vector<cv::DMatch> &matches,
                                       Eigen::Matrix3f &K,
                                       Eigen::Matrix4f &T)
{
#if 0
    std::vector<cv::Point2f> pointset1;
    std::vector<cv::Point2f> pointset2;

    for (int i = 0; i < (int)matches.size(); i++)
    {
        pointset1.push_back(keypoints1[matches[i].queryIdx].pt);
        pointset2.push_back(keypoints2[matches[i].trainIdx].pt);
    }

    cv::Point2d principal_point(K(0, 2), K(1, 2));
    double focal_length = (K(0, 0) + K(1, 1)) * 0.5;
    cv::Mat essential_matrix;
    essential_matrix = cv::findEssentialMat(pointset1, pointset2, focal_length, principal_point);
    std::cout << "essential_matrix is " << std::endl
              << essential_matrix << std::endl;

    cv::Mat R;
    cv::Mat t;

    cv::recoverPose(essential_matrix, pointset1, pointset2, R, t, focal_length, principal_point);

    Eigen::Matrix3f R_eigen;
    Eigen::Vector3f t_eigen;
    cv::cv2eigen(R, R_eigen);
    cv::cv2eigen(t, t_eigen);
    T.block(0, 0, 3, 3) = R_eigen;
    T.block(0, 3, 3, 1) = t_eigen;
    T(3, 0) = 0;
    T(3, 1) = 0;
    T(3, 2) = 0;
    T(3, 3) = 1;

    std::cout << "Transform is " << std::endl
              << T << std::endl;
#endif
}


// #pragma GCC optimize ("O0")

#include <boost/numeric/odeint.hpp>       // odeint function definitions

using namespace boost::numeric::odeint;

#include "eigenIncluder.hpp"
#include "forceModels.hpp"
#include "streamTrace.hpp"
#include "instrument.hpp"
#include "navigation.hpp"
#include "satRefSys.hpp"
#include "constants.hpp"
#include "acsConfig.hpp"
#include "gravity.hpp"
#include "common.hpp"
#include "jplEph.hpp"
#include "tides.hpp"
#include "enums.h"

typedef Eigen::Vector<double, 48> VectorXSTMS;			//todo aaron, delete this

map<int, OrbitPropagator>	orbitPropagatorMap;

GravityModel OrbitPropagator::gravityModel;

/* Frac: Gives the fractional part of a number
*/
double Frac (double x)
{
	return x - floor(x);
}

/* Computes the Sun's geocentric position using a low precision analytical series
*/
Vector3d SunPos(
	double mjdTT)		///< Terrestrial time: modified Julian date
{
	/* Constants
	*/
	const double eps	= 23.43929111 * D2R;				// Obliquity of J2000 ecliptic
	const double T		= (mjdTT - mjdJ2000) / 36525.0;		// Julian cent. since J2000

	/* Mean anomaly, ecliptic longitude and radius
	*/
	double M = PI2 * Frac ( 0.9931267 + 99.9973583 * T);										// [rad]
	double L = PI2 * Frac ( 0.7859444 + M/PI2 + ( 6892 * sin(M) + 72 * sin(2 * M) ) / 1296e3 );	// [rad]
	double R = 149.619e9 - 2.499e9 * cos(M) - 0.021e9 * cos(2 * M);								// [m]

	/* Solar position vector [m] with respect to the mean equator and equinox of J2000 (EME2000, ICRF)
	*/
	return R_x(-eps) * R * Vector3d(cos(L), sin(L), 0);
}

/* Computes the fractional illumination of a spacecraft in the vicinity of the Earth assuming a cylindrical shadow model
*
*/
double Illumination(
	const Vector3d& 		rSat,			///< Spacecraft position vector [m]
	const Vector3d& 		rSun )			///< Sun position vector [m]
{
	Vector3d eSun = rSun / rSun.norm();    // Sun direction unit vector
	double s     	= rSat.dot(eSun);      // Projection of s/c position

	/* Illumination factor:
	*  1: Spacecraft fully illuminated by the Sun
	*  0: Spacecraft in Earth shadow
	*/
	return ( ( s > 0 || (rSat - s * eSun).norm() > RE_WGS84 ) ?  1 : 0 );
}

/* Class type: The class of solar radiation pressure
*
*/
SolarRadPressure::SolarRadPressure(
	SRPPara 				paraSRP)
{
	mSRPPara		= paraSRP;
}

Vector3d SolarRadPressure::directSolarRadiationAcc(
	Trace&				trace,		 				///< Trace to output to
	double 				mjdTT,		 				///< Terrestrial time (modified Julian date)
	const Vector3d& 	rSat)        				///< Satellite position vector, unit: m, m/s
{

	/* Relative position vector of spacecraft w.r.t. Sun
	*/
	Vector3d rSun;
	jplEphPos(nav.jplEph_ptr, mjdTT + JD2MJD, E_ThirdBody::SUN, rSun);

	// trace << "Calculated sun position: " << std::setw(14) << mjdTT << std::setw(14) << rSun.transpose() << std::endl;

	Vector3d rDis = rSat - rSun;
	double illumination = Illumination(rSat, rSun);

	switch (mSRPPara.srpMdlName)
	{
		case E_SRPModels::CANNONBALL:		return illumination * mSRPPara.srpCoef * (mSRPPara.srpArea / mSRPPara.satMass) * PSOL * (AU * AU) * rDis / pow(rDis.norm(), 3);
		case E_SRPModels::BOXWING:			return Vector3d::Zero();
		case E_SRPModels::ECOM:				return Vector3d::Zero();
		case E_SRPModels::ECOM2:			return Vector3d::Zero();
		default:							return Vector3d::Zero();
	}
}

Vector3d SolarRadPressure::indirectSolarRadiationAcc(
	Trace&				trace,		 					///< Trace to output to (similar to cout)
	double 				mjdTT,		 					///< Terrestrial time (modified Julian date)
	const Vector3d& 	rSat)        					///< Satellite position vector, unit: m, m/s
{
	return Vector3d::Zero();
}

Vector3d antennarThrustAcc()
{
	return Vector3d::Zero();
}

Vector3d empiricalAcc()
{
	return Vector3d::Zero();
}

Vector3d manoeuvreAcc()
{
	return Vector3d::Zero();
}

/* Set options for force models in the propagator
*
*/
void OrbitPropagator::setPropOption(
	ForceModels				forceMdl)
{
	/* Options from yaml file */
	propOpt.optEarthGravMdl.earthGravMdl		= E_GravMdl::GGM03S;

	/* Parameters from yaml file */
	propOpt.optEarthGravMdl.earthGravAccDeg.mMax = forceMdl.egmAccDeg;
	propOpt.optEarthGravMdl.earthGravAccDeg.nMax = forceMdl.egmAccOrd;
	propOpt.optEarthGravMdl.earthGravSTMDeg.mMax = forceMdl.egmSTMDeg;
	propOpt.optEarthGravMdl.earthGravSTMDeg.nMax = forceMdl.egmSTMOrd;

	propOpt.paraSRP.srpMdlName	= forceMdl.srp_model;
	propOpt.paraSRP.satMass		= forceMdl.sat_mass;
	propOpt.paraSRP.srpArea		= forceMdl.srp_area;
	propOpt.paraSRP.srpCoef		= forceMdl.srp_coef;
}

/* initialise the propagator with time, state and necessary parameters
*
*/
void OrbitPropagator::init(
	Vector6d				rvSatECI,
	double 					mjdUTC,
	double					leapSec,
	ERP*					erpSrc,
	EGMCoef					egmCoef)
{
	inertialState 		= rvSatECI;
	mMJDUTC 			= mjdUTC;
	erp 				= erpSrc;

	geterp(*erp, mMJDUTC, erpv);

	gravityModel = GravityModel(propOpt.optEarthGravMdl, egmCoef);

	if	(acsConfig.forceModels.solar_radiation_pressure)
	{
		solarRadPressure = SolarRadPressure(propOpt.paraSRP);
	}
}

void OrbitPropagator::update(
	double				mjdUTC)
{
	// update time
	mMJDUTC = mjdUTC;

	// update erp
	geterp(*erp, mMJDUTC, erpv);
	double dUTC_TAI	= -(19 + erpv.leaps);
	double xp		= erpv.xp;
	double yp		= erpv.yp;
	double dUT1_UTC	= erpv.ut1_utc;
	double lod		= erpv.lod;

	// update iers
	iers = IERS(dUT1_UTC, dUTC_TAI, xp, yp, lod);

	//update srp parameters
	solarRadPressure.mSRPPara = propOpt.paraSRP;

	//update third body positions
	double mjdTT = mMJDUTC + iers.TT_UTC() / 86400.0;
	for (int i = 0; i < E_ThirdBody::_size(); i++)
	{
		E_ThirdBody	body 		= E_ThirdBody::_values()[i];
		string		bodyName	= E_ThirdBody::_names()[i];

		if (acsConfig.forceModels.process_third_body[body] == false)
		{
			continue;
		}

		// Relative position vector of satellite w.r.t. point mass
		Vector3d thirdBodyPos;
		jplEphPos(nav.jplEph_ptr, mjdTT + JD2MJD, body, thirdBodyPos);

		thirdBodyPositionMap[body] = thirdBodyPos;
	}

	//update reference frame matrices
	eci2ecef_sofa(mjdUTC, iers, mECI2ECEF, mdECI2ECEF);

	//update egm coefficients for tides
	gravityModel.correctEgmCoefficients(std::cout, mjdUTC, erpv, mECI2ECEF);
}

Vector3d  OrbitVEQPropagator::calculateAccelNGradient(
	Trace&    				trace,	 			///< Trace to output to (similar to cout)
	const Vector3d& 		rSat,		 		///< Inertial position of satellite (m)
	const Vector3d& 		vSat,		 		///< Inertial velocity of satellite (m/s)
	Matrix3d&				mGradient,			///< Gradient (G=da/dr) in the ICRF/J2000 system
	Vector3d&				dadCr,       		///< Partials of acceleration w.r.t. to the solar radiation coefficient
	const Matrix3d& 		mECI2ECEF)	 		///< Transformation matrix from ECI coordinate to ECEF
{
// 	bool bVarEq = true;
// 	Vector3d earthGravityAcc = gravityModel.centralBodyGravityAcc(trace, mMJDUTC, mERPv, rSat, mECI2ECEF, bVarEq);
// 	Vector3d drSat;
// 	const double dinc = 1;   // Position increment [m]
// 	Vector3d daSat;
//
// 	/* Gradient
// 	*/
//   	for (int i = 0; i < 3; i++)
// 	{
//     	drSat = Vector3d::Zero();
// 		drSat(i) = dinc;	// Set offset in i-th component of the position vector
//     	daSat = gravityModel.centralBodyGravityAcc (trace, mMJDUTC, mERPv, rSat + drSat, mECI2ECEF, bVarEq) -  earthGravityAcc;	// Acceleration difference
//     	mGradient.col(i) = daSat / dinc;	// Derivative with respect to i-th component
//   	}
//
//   	/* Radiation pressure coefficient partials
// 	*/
// 	double mjdTT = mIERS.TT_UTC(mMJDUTC) / 86400.0 + mMJDUTC;
// 	solarRadPressure.mSRPPara.srpCoef = 1;
// 	dadCr = solarRadPressure.directSolarRadiationAcc(trace, mjdTT, rSat);
	Vector3d acc;
	return acc;
}
/** Propagate orbits by ODE functor.
 * Used by RK - input is a 6 element inertial state and time, and output is the derivative of the state at the provided time.
 * This function requires the update function to be called first to set planetary ephemerides, erp values, etc.
 */
void OrbitPropagator::operator ()(
	Vector6d&		 inertialState,				 ///< Inertial position and velocity of satellite (m, m/s)
	Vector6d&		dInertialState,
	const double	mjdUTCinSec)
{
	Instrument instrument(__FUNCTION__);

	//Get sub vectors from the state
	Vector3d rSat = inertialState.head(3);
	Vector3d vSat = inertialState.tail(3);

	auto& trace = std::cout;

	double mjdUTC = mjdUTCinSec / 86400;

	// double erpv[4] = {};
	// geterp_from_utc(&nav.erp, mjdUTC, erpv);
	// double leapSec	= nav.leaps;
	// double dUT1_UTC = erpv[2];
	// double dUTC_TAI	= -(19 + leapSec);
	// double xp = erpv[0];
	// double yp = erpv[1];
	// double lod = erpv[3];
	// clIERS iersInstance;
	// iersInstance.Set(dUT1_UTC, dUTC_TAI, xp, yp, lod);
	// double mjdTT = mjdUTC + iersInstance.TT_UTC(mjdUTC) / 86400;

	// Matrix3d mECI2ECEF = Matrix3d::Identity();
	// Matrix3d mdECI2ECEF = Matrix3d::Identity();
	// eci2ecef_sofa(mjdUTC, iersInstance, mECI2ECEF, mdECI2ECEF);
	// std::cout << "mECI2ECEF: " << mjdUTC << "  " << std::setw(14) << mECI2ECEF << std::endl;
	// Vector3d aSat = calculateAcceleration(std::cout, mjdTT, rSat, vSat, mECI2ECEF, egm.cmn, egm.smn, 12, 12);
	// std::cout << "Accelerations: " << mjdUTC << "  " << std::setw(14) << aSat.transpose() << std::endl;

	// std::cout << "mECI2ECEF: " << orbitProp.mMJDUTC << "  " << std::setw(14) << mECI2ECEF << std::endl;

	//calculate acceleration components
	Vector3d aSat = Vector3d::Zero();

	bool bVarEq = false;
	if (acsConfig.forceModels.earth_gravity)
	{
		Instrument instrument("Grav");

		Vector3d earthgravityAcc = -GM_Earth * rSat.normalized() / rSat.squaredNorm();// = gravityModel.centralBodyGravityAcc(trace, mMJDUTC, erpv, rSat, mECI2ECEF, bVarEq);

// 		trace << "Calculated accleration due to the Earth's central body gravity: " << std::setw(14) << mMJDUTC << std::setw(14) << earthgravityAcc.transpose() << std::endl;

		aSat += earthgravityAcc;
	}

	double mjdTT = iers.TT_UTC() / 86400.0 + mMJDUTC;
	for (int i = 0; i < E_ThirdBody::_size(); i++)
	{
		E_ThirdBody	body 		= E_ThirdBody::_values()[i];
		string		bodyName	= E_ThirdBody::_names()[i];

		if (acsConfig.forceModels.process_third_body[body] == false)
		{
			continue;
		}

		Vector3d& bodyPos = thirdBodyPositionMap[body];

		Vector3d thirdbodyAcc = accelPointMassGravity(trace, mjdTT, rSat, body, bodyPos);

// 		trace << "Calculated acceleration due to " << bodyName << "'s attraction: " << std::setw(14) << mMJDUTC << std::setw(14) << thirdbodyAcc.transpose() << std::endl;

		aSat += thirdbodyAcc;
	}

	if (acsConfig.forceModels.relativity_effect)
	{
		Instrument instrument("Rel");

		Vector3d relativityAcc = gravityModel.relativityEffectsAcc(trace, rSat, vSat);

// 		trace << "Calculated acceleration due to the relativity effect: " << std::setw(14) << mMJDUTC << std::setw(14) << relativityAcc.transpose() << std::endl;

		aSat += relativityAcc;
	}

	if (acsConfig.forceModels.solar_radiation_pressure)
	{
		Instrument instrument("SRP");

		Vector3d directSRPAcc = solarRadPressure.directSolarRadiationAcc(trace, mMJDUTC, rSat);

// 		trace << "Calculated acceleration due to the direct solar radiation: " << std::setw(14) << mMJDUTC << std::setw(14) << directSRPAcc.transpose() << std::endl;

		aSat += directSRPAcc;
	}

	// std::cout << "Total accelerations: " << mjdUTC << "  " << std::setw(14) << aSat.transpose() << std::endl;

	//set ODE output from sub vectors
	dInertialState.head(3) = vSat;
	dInertialState.tail(3) = aSat;

	std::cout << "\n" << (aSat.dot(rSat.normalized()));
}

/** Observer, prints time and state when called (during integration)
*
*/
void obsvrSatMotion(
	const Vector6d& x,
	const double mjdUTC)
{
	// std::cout << std::endl;
	// std::cout << std::setw(8) << mjdUTC << std::setw(14) << x[0] << std::setw(14) << x[1] << std::setw(14) << x[2]
	// 		  << std::setw(14) << x[3] << std::setw(14) << x[4] << std::setw(14) << x[5] << std::setw(14) <<  std::endl;
}

/** ODE of variational equation to be solved, i.e. the derivative of the state vector and the state transition matrix
*
*/
// struct VarEqPropagator
// {
// 	void operator()(
// 		const	VectorXSTMS&	rvPhiS,
// 				VectorXSTMS&	drvPhiSdt,
// 		const	double			mjdUTCinSec)
// 	{
// 		double mjdUTC = mjdUTCinSec / 86400;	// Time
//
// 		Vector3d rSat = rvPhiS.segment(0, 3);	// Position components
// 		Vector3d vSat = rvPhiS.segment(3, 3);	// Velocity components
//
// 		MatrixXd Phi = MatrixXd::Identity(6, 6);	// State transition matrix
// 		// for (int j = 0; j < 6; j++)
// 		// {
// 		// 	Phi.col(j) = rvPhiS.segment(6 * (j + 1), 6);
// 		// }
//
// 		MatrixXd S = MatrixXd::Zero(6, 1);	// Sensitivity matrix
// 		// for (int j = 0; j < 1; j++)
// 		// {
// 		// 	S.col(j) = rvPhiS.segment(6 * (j + 7), 6);
// 		// }
//
// 		Matrix3d mECI2ECEF	= Matrix3d::Identity();
// 		Matrix3d mdECI2ECEF	= Matrix3d::Identity();
// 		updPropagator(mjdUTC); //time epoch inside the integrator
// 		eci2ecef_sofa(mjdUTC, mIERS, mECI2ECEF, mdECI2ECEF);
//
// 		Matrix3d mGradient;
// 		Vector3d dadCr;
// 		Vector3d aSat = calculateAccelNGradient(std::cout, rSat, vSat, mGradient, dadCr, mECI2ECEF);
//
// 		/* Time derivative of state transition matrix
// 		*
// 		*/
// 		MatrixXd dfdy = MatrixXd::Zero(6, 6);
// 		for (int i = 0; i < 3; i++)
// 		for (int j = 0; j < 3; j++)
// 		{
// 			dfdy(i, 	j	 ) = 0;						// dv/dr(i, j)
// 			dfdy(i + 3, j	 ) = mGradient(i, j);		// da/dr(i, j)
// 			dfdy(i,     j + 3) = ( i == j ? 1 : 0 );	// dv/dv(i, j)
// 			dfdy(i + 3, j + 3) = 0;						// da/dv(i, j)
// 		}
// 		MatrixXd dPhi = dfdy * Phi;	// Time derivative of state transition matrix
//
//
// 		/* Time derivative of sensitivity matrix
// 		*
// 		*/
// 		MatrixXd dfdp(6, 1);
// 		for (int i = 0; i < 3; i++)
// 		{
// 			dfdp(i    ) = 0;							// dv/dCr(i)
// 			dfdp(i + 3) = dadCr(i);						// da/dCr(i)
// 		}
//
// 		MatrixXd dS = MatrixXd::Zero(6, 1);
// 		dS = dfdy * S + dfdp;
//
// 		/* Derivative of combined state vector and state transition matrix
// 		*
// 		*/
// 		for (int i = 0; i < 3; i++)
// 		{
// 			drvPhiSdt(i    ) = vSat(i);                    // dr/dt(i)
// 			drvPhiSdt(i + 3) = aSat(i);                    // dv/dt(i)
// 		}
//
// 		for (int i = 0; i < 6; i++)
// 		for (int j = 0; j < 6; j++)
// 		{
// 			drvPhiSdt(6 * (j + 1) + i  ) = dPhi(i, j);     // dPhi/dt(i,j)
// 		}
//
// 		for (int i = 0; i < 6; i++)
// 		for (int j = 6; j < 7; j++)
// 		{
// 			drvPhiSdt(6 * (j + 1) + i) = dS(i, j - 6);     // dS/dt(i,j)
// 		}
// 	}
// };

void updateOrbits(
	Trace&			trace,
	KFState&		kfState,
	GTime           time)
{
	{
		//get current inertial states from the kfState
		for (auto& [kfKey, index] : kfState.kfIndexMap)
		{
			if	( kfKey.type	!= KF::SAT_POS
				||kfKey.num		!= 0)
			{
				continue;
			}

			auto& orbitPropagator = orbitPropagatorMap[kfKey.Sat];

			Vector6d rvECI;
			for (int i = 0; i < 3; i++)
			{
				KFKey posKey = kfKey;

				posKey.num = i;
				kfState.getKFValue(posKey, rvECI(0 + i));

				posKey.type = KF::SAT_POS_RATE;
				kfState.getKFValue(posKey, rvECI(3 + i));
			}

			double mjdUTC	= gpst2mjd(tsync);
			double leapSec	= nav.leaps;
			orbitPropagator.setPropOption(acsConfig.forceModels);
			orbitPropagator.init(rvECI, mjdUTC, leapSec, &nav.erp, nav.egm);
		}
	}

	for (auto& [satId, orbitPropagator] : orbitPropagatorMap)
	{
		SatSys Sat;
		Sat.fromHash(satId);

		double dt = time - kfState.time;			//time indicates the current epoch, kfState.time indicates the last epoch

		double t0 = gpst2mjd(kfState.time)	* 86400;
		double t1 = gpst2mjd(time)			* 86400;

		if (dt == 0)
		{
			return;
		}

// 		orbitPropagator.stateTime = time;

		VectorXd oldState = orbitPropagator.inertialState;

		double t_mid = (t0 + t1) / 2 / 86400;

	// 	orbitPropagator.updPropagator(t0);//time epoch before the orbital propagation

		std::cout << std::setprecision(2) << std::fixed;
		std::cout
		<< "ICRF coordinates before the orbital propagation step: "
		<< std::setprecision(14) << orbitPropagator.inertialState.transpose() <<  std::endl;

		orbitPropagator.update(t_mid); //time epoch inside the integrator

		//Run the propagator using the functor

		if (acsConfig.forceModels.ode_integrator == +E_Integrator::RKF78)
		{
			typedef runge_kutta_fehlberg78<Vector6d>	rkf78;			// Error stepper, used to create the controlled stepper
// 			typedef controlled_runge_kutta<rkf78>		ctrl_rkf78;		// Controlled stepper: it's built on an error stepper and allows us to have the output at each internally defined (refined) timestep, via integrate_adaptive call

			double errAbs = 1.0e-16; // Error bounds
			double errRel = 1.0e-13;

			//integrate_adaptive(ctrl_rkf78(), OrbitPropagator(), rvECI, t0, t1, dt);
			auto controller = make_controlled(errAbs, errRel, rkf78());
			integrate_adaptive(controller, orbitPropagator, orbitPropagator.inertialState, t0, t1, dt);
		}

		std::cout
		<< "ICRF coordinates after the orbital propagation step:  "
		<< std::setprecision(14) << orbitPropagator.inertialState.transpose() <<  std::endl;
	// 	VectorXSTMS	rvPhiS	= kfState.xSTMSM;

	// 	if (0)
	// 	{
	// 		// typedef runge_kutta4<VectorXSTMS> rk4;
	// 		// integrate_const(rk4(), odeVarEquation, rvPhiS, t0, t1, dt);
	// 		integrate(VarEqPropagator(), rvPhiS, t0, t1, dt);
	// 		// typedef runge_kutta_dopri5<Vector6d> rk5;
	// 		// typedef controlled_runge_kutta<rk5> ctrl_rk5;
	// 		// double errAbs = 1.0e-10; // Error bounds
	// 		// double errRel = 1.0e-8;
	// 		// integrate_adaptive(ctrl_rk5(), odeVarEquation, rvPhiS, t0, t1, dt);
	//
	// 		std::cout << "ICRF coordinates before the variatioanal equation propagation: " << std::setprecision(14) << rvPhiS.head(3).transpose() <<  std::endl;
	// 		kfState.xSTMSM = rvPhiS;
	// 	}

	// 	kfState.x		= rvECI;
	// 	kfState.time	= time;

		// ECI to ECEF transformation

		// double erpv[4] = {};
		// geterp(&nav.erp, time + dt, erpv);
		// double mjdUTC = gpst2mjd(time + dt);
		// double leapSec	= nav.leaps;
		// double dUT1_UTC = erpv[2];
		// double dUTC_TAI	= -(19 + orbitProp.mLeapSec);
		// double xp = erpv[0];
		// double yp = erpv[1];
		// double lod = erpv[3];
		// clIERS iersInstance;
		// iersInstance.Set(dUT1_UTC, dUTC_TAI, xp, yp, lod);
		// double mjdTT = mjdUTC + iersInstance.TT_UTC(mjdUTC) / 86400;
		// Matrix3d mECI2ECEF = Matrix3d::Identity();
		// Matrix3d mdECI2ECEF = Matrix3d::Identity();
		// Vector6d rvECEF;
		// eci2ecefVec_sofa(mjdUTC, iersInstance, rvECI, rvECEF);
		// std::cout << "ITRF coordinates: " << std::endl << std::setprecision(14) << rvECEF.transpose() <<  std::endl;

	// 	Matrix3d mECI2ECEF	= Matrix3d::Identity();
	// 	Matrix3d mdECI2ECEF	= Matrix3d::Identity();
		orbitPropagator.update(gpst2mjd(time));//time epoch after the orbital propagation;

		Vector3d rECI = orbitPropagator.inertialState.head(3);
		Vector3d vECI = orbitPropagator.inertialState.tail(3);
// 		std::cout
// 		<< "ICRF coordinates after the orbital propagation step:  "
// 		<< std::setprecision(14) << rECI.transpose() << " " << vECI.transpose() <<  std::endl;

		Vector3d rECEF2;
		Vector3d vECEF2;
		eci2ecef_sofa(orbitPropagator.mMJDUTC, orbitPropagator.iers, rECI, vECI, rECEF2, vECEF2);
// 		std::cout
// 		<< "ITRF coordinates after the orbital propagation step:  "
// 		<< std::setprecision(14) << rECEF2.transpose() << " " << vECEF2.transpose() <<  std::endl;

// 		std::ofstream fileOPResults("./ex01/OPResults.txt", std::ios_base::app | std::ios_base::in);
// 		if (!fileOPResults)
// 		{
// 			std::cout << "Error openinng results file!\n";
// 			return;
// 		}

// 		fileOPResults << std::setprecision(6) << std::fixed;

		// fileOPResults << setw(16) << orbitProp.mMJDUTC << setw(16) << rvECI.transpose() << setw(16) << rvECEF2.transpose() << endl;
// 		fileOPResults << std::setw(16) << orbitPropagator.mMJDUTC << " " << rECEF2.transpose() << vECEF2.transpose() << std::endl;

		//get the change in state ready for use in the kalman filter's state transition
// 		orbitPropagator.stateDelta		= orbitPropagator.inertialState - oldState;

		//get the noise in the state delta according to uncertainty in the inputs, and propagation uncertainty for use in the filters' covariance transition
// 		orbitPropagator.stateDeltaNoise	= Matrix6d::Zero();

		Matrix6d stateTransition = Matrix6d::Identity();
		stateTransition.topRightCorner(3,3) = Matrix3d::Identity() * dt;

		Vector6d transitionedState = stateTransition * oldState;

		Vector6d missingDynamics = orbitPropagator.inertialState - transitionedState;

		for (int i = 0; i < 3; i++)
		{
			KFKey oneKey = {.type = KF::ONE};

			KFKey kfKey;
			kfKey.Sat = Sat;
			kfKey.num = i;

			kfKey.type = KF::SAT_POS;

			kfState.setKFTrans(kfKey, oneKey, missingDynamics(0 + i));

			kfKey.type = KF::SAT_POS_RATE;

			kfState.setKFTrans(kfKey, oneKey, missingDynamics(3 + i));
		}
	}
}


#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <iterator>
#include <map>
#include <set>
#include <sstream>
#include <string>
#include <vector>

#include <Eigen/Core>

#include "igl/AABB.h"
#include "igl/adjacency_list.h"
#include "igl/adjacency_matrix.h"
#include "igl/ambient_occlusion.h"
#include "igl/components.h"
#include "igl/cotmatrix.h"
#include "igl/decimate.h"
#include "igl/dijkstra.h"
#include "igl/embree/ambient_occlusion.h"
#include "igl/embree/EmbreeIntersector.h"
#include "igl/embree/unproject_onto_mesh.h"
#include "igl/gaussian_curvature.h"
#include "igl/invert_diag.h"
#include "igl/jet.h"
#include "igl/per_vertex_attribute_smoothing.h"
#include "igl/per_vertex_normals.h"
#include "igl/point_mesh_squared_distance.h"
#include "igl/polygon_mesh_to_triangle_mesh.h"
#include "igl/principal_curvature.h"
#include "igl/ray_mesh_intersect.h"
#include "igl/readOFF.h"
#include "igl/unproject_onto_mesh.h"
#include "igl/viewer/Viewer.h"
#include "igl/writeOFF.h"

#include "GLFW/glfw3.h"

#include "dijkstra.hxx"

const double mOCC_MAX_RAY_DIST = 10.0;
const int mOCC_NUM_RAYS = 20;
const double mDECIMATE_PERC = 0.6;

void getMeanCurvature(const Eigen::MatrixXd &V, const Eigen::MatrixXi &F,
                      Eigen::VectorXd &VCs) {
  Eigen::MatrixXd PD1,PD2;
  Eigen::VectorXd PV1,PV2;
  //igl::principal_curvature(V,F,PD1,PD2,PV1,PV2, 5);
  igl::principal_curvature(V,F,PD1,PD2,PV1,PV2, 8);
  //mean_curve = 0.5 * (PV1 + PV2);
  VCs = PV1.cwiseProduct(PV1) + PV2.cwiseProduct(PV2);
  VCs = VCs.cwiseMin(0.1);
  VCs /= VCs.maxCoeff();

  Eigen::VectorXd ones = Eigen::VectorXd::Constant(VCs.rows(), 1, 1);
  VCs = ones - VCs;

  /*
  // Get the laplacian
  Eigen::SparseMatrix<double> L,M,Minv;
  //   from the cotangent matrix
  igl::cotmatrix(V,F, L);
  //   and the mass matrix
  igl::massmatrix(V,F, igl::MASSMATRIX_TYPE_VORONOI, M);
  igl::invert_diag(M, Minv);

  Eigen::MatrixXd HN = -Minv*(L * V);
  // Mean curviture is this value.
  Eigen::VectorXd VC = HN.rowwise().norm();

  // Smooth this out a bit.
  igl::per_vertex_attribute_smoothing(VC,F, VCs);
  // Cutoff at 1 (to avoid overpowering by high curvature)
  VCs = VCs.cwiseMin(1.0);

  // Normalize
  VCs /= VCs.maxCoeff();
  // Invert so high curvature has low weight.
  Eigen::VectorXd ones = Eigen::VectorXd::Constant(VCs.rows(), 1, 1);
  // Sqrt before switching
  VCs = VCs.cwiseSqrt();
  VCs = ones - VCs;
  //std::cout << VCs;
  */

}

void printHelp() {
  printf("Press 'U' to undo last point\n");
  printf("Press 'C' to show colors from curvature\n");
  printf("Press 'D' to show colors from ambient occlusion\n");
}

bool file_exists(const char* fn) {
  FILE* ifs = fopen(fn, "r");
  if (ifs) {
    fclose(ifs);
    return true;
  }
  return false;
}

int main(int argc, char* argv[]) {
  if (argc < 3) {
    fprintf(stderr, "Used for selecting points on the surface of a mesh. Will write the\n"
                    "corresponding vertices to output.off\n");
    fprintf(stderr, "usage: %s <input.off> <output.off> [nowrite]\n",
            argv[0]);
    return -1;
  }

  bool writeable = true;
  if (argc == 4) {
    writeable = false;
  }

  Eigen::MatrixXd V,V2;
  Eigen::MatrixXi F,F2;
  igl::embree::EmbreeIntersector ei;

  printf("Reading in mesh...\n");
  // Read in the input file.
  igl::readOFF(argv[1], V, F);

  // Colors are all white to start.
  Eigen::MatrixXd VC = Eigen::MatrixXd::Constant(V.rows(),3,1);
  // Initialize this thing so we can get the vertices.
  ei.init(V.cast<float>(), F);

  // Ambient occlusion
  Eigen::VectorXd AO;
  Eigen::MatrixXd N;
  igl::per_vertex_normals(V,F,N);
  // Compute ambient occlusion factor using embree
  printf("Building AABB tree...\n");
  igl::AABB<Eigen::MatrixXd, 3> aabb;
  aabb.init(V, F);
  const auto & shoot_ray = [&aabb,&V,&F](
      const Eigen::Vector3f& _s,
      const Eigen::Vector3f& dir)->bool
  {
    Eigen::Vector3f s = _s+1e-4*dir;
    igl::Hit hit;
    if (aabb.intersect_ray(V,F, s.cast<double>().eval(),dir.cast<double>().eval(), hit)) {
      return hit.t < mOCC_MAX_RAY_DIST;
    }
    return false;
  };

  printf("Now setting ambient occlusion vertices to improve contrast (patience)...\n");
  //igl::embree::ambient_occlusion(V,F,V,N,300,AO);
  igl::ambient_occlusion(shoot_ray, V,N, mOCC_NUM_RAYS,AO);
  AO = 1.0 - AO.array();
  for (int i = 0; i < V.rows(); ++i) {
    VC.row(i) *= AO(i);
  }

  printf("Constructing adjacency list...\n");
  // Need the adjacency list
  std::vector<std::vector<int> > VV;
  igl::adjacency_list(F, VV);

  std::vector<std::vector<double> > darkWeight;
  darkWeight.resize(VV.size());
  //std::vector<std::vector<double> > curveWeight;
  //curveWeight.resize(VV.size());
  for (int i = 0; i < VV.size(); ++i) {
    for (int j = 0; j < VV[i].size(); ++j) {
      // Just use the value of the ambient occlusion.
      darkWeight[i].push_back(AO(VV[i][j]));
      //curveWeight[i].push_back(1.0 - AO(VV[i][j]));
    }
  }

  // Also caluculate the curvature values.
  printf("Calculating mean curvature...\n");
  Eigen::VectorXd mean_curve;
  getMeanCurvature(V,F, mean_curve);
  std::vector<std::vector<double> > curveWeight;
  curveWeight.resize(VV.size());
  double maxCurve = 0;
  double minCurve = 0;
  for (int i = 0; i < VV.size(); ++i) {
    for (int j = 0; j < VV[i].size(); ++j) {
      // Difference in gausssian curvature
      //double diff = std::abs(VCs(i) - VCs(j));
      double diff = mean_curve(VV[i][j]);
      curveWeight[i].push_back(diff);
      maxCurve = std::max(diff, maxCurve);
      minCurve = std::min(diff, minCurve);
    }
  }

  // Create the colors for mean curviture to use
  Eigen::MatrixXd curve_colors;
  igl::jet(mean_curve, true, curve_colors);

  std::vector<int> all_points;  // will contain all selected points.

  igl::viewer::Viewer viewer;
  viewer.data.set_mesh(V, F);
  viewer.data.set_colors(curve_colors);
  viewer.core.show_lines = false;

  if (file_exists(argv[2])) {
    printf("File exists! Adding points from previous run!\n");
    Eigen::MatrixXd sp;
    Eigen::MatrixXd spf;
    igl::readOFF(argv[2], sp, spf);

    // Create indices.
    Eigen::VectorXi ids(V.rows());
    for (int i = 0; i < V.rows(); ++i) {
      ids(i) = i;
    }

    Eigen::VectorXd D;
    Eigen::VectorXi I;
    Eigen::MatrixXd CP;
    igl::point_mesh_squared_distance(sp, V,ids, D,I,CP);

    // Add them to the viewer.
    for (int i = 0; i < sp.rows(); ++i) {
      all_points.push_back(I(i));
      VC.row(I(i)) << 1,0,0;

      std::stringstream label;
      label << all_points.size();
      viewer.data.add_label(sp.row(i), label.str());
    }
  }

  viewer.callback_key_down = [&](igl::viewer::Viewer& viewer, unsigned char key, int modifier)->bool {
    switch(key) {
      case 'U':
      {
        printf("Undoing last point...\n");
        // Undo the last choice
        int last = all_points.back();
        printf("Undoing last choice... %d\n", last);
        all_points.pop_back();
        VC.row(last) << 1,1,1;
        VC.row(last) *= AO(last);

        // Set the colors back to what they were.
        viewer.data.set_colors(VC);
        viewer.data.labels_positions.conservativeResize(all_points.size(), 3);
        viewer.data.labels_strings.pop_back();
        break;
      }
      case 'C':
        printf("Setting curvature colors...\n");
        viewer.data.set_colors(curve_colors);
        break;
      case 'D':
        printf("Setting colors to highlight darkness...\n");
        viewer.data.set_colors(VC);
        break;
      case 'W':
      {
        if (!writeable) {
          return false;
        }
        printf("Writing output to %s\n", argv[2]);
        // Write to output file
        FILE* ofs = fopen(argv[2], "w");
        fprintf(ofs, "OFF\n");
        fprintf(ofs, "%d 0 0\n", all_points.size());
        for (int idx : all_points) {
          fprintf(ofs, "%f %f %f\n", V(idx, 0), V(idx, 1), V(idx, 2));
        }
        fclose(ofs);
        // You can quit now.
        printf("Finished! Please close the program.\n");
        break;
      }
      default:
        // Question mark ('?')
        if (key == '/' && modifier == GLFW_MOD_SHIFT) {
          printHelp();
        }
    }

    // Make sure this stays the same.
    viewer.core.lighting_factor =
        std::min(std::max(viewer.core.lighting_factor,0.f),1.f);
    return false;
  };
  viewer.callback_mouse_down =
      [&](igl::viewer::Viewer &viewer, int, int)->bool
      {
        if (!writeable) {
          return false;
        }
        // fid will be the face ID. Need to get the vertex ID
        int fid, vid;
        // Cast a ray in the view direction starting from the mouse position
        double x = viewer.current_mouse_x;
        double y = viewer.core.viewport(3) - viewer.current_mouse_y;
        if(igl::embree::unproject_onto_mesh(Eigen::Vector2f(x,y), F,
                                            viewer.core.view * viewer.core.model,
                                            viewer.core.proj,
                                            viewer.core.viewport,
                                            ei,
                                            fid, vid)) {
          // If it already exists, ignore it.
          for (int i : all_points) {
            if (vid == i) return false;
          }

          all_points.push_back(vid);
          VC.row(vid) << 1,0,0;
          viewer.data.set_colors(VC);
          std::stringstream label;
          label << all_points.size();
          viewer.data.add_label(V.row(vid), label.str());

          return true;
        }
        return false;
      };

  printHelp();
  viewer.launch();
}


/******************************
      Author: Joel Veness
        Date: 2013
******************************/

#include "skipcts.hpp"
#include "fastmath.hpp"

#include <boost/random.hpp>

// adaptive KT estimator parameters
static const bool UseDiscounting = true;
static const count_t Gamma       = 0.02f;
static const count_t Discount    = 1.0f - Gamma;
static const count_t KT_Alpha = 0.0625f;
static const count_t KT_Alpha2 = KT_Alpha + KT_Alpha;

// initialise weights for switching model
static const double SplitPrior        = 0.925;
static const double LogSplitPrior   = std::log(SplitPrior);
static const double LogStopPrior    = std::log(1.0 - SplitPrior);

static const double LogOneHalf = std::log(0.5);

struct prior_t {
    double stop;
    double split;
    double skip;
};

// prior weights for either splitting, stopping or skipping
static const prior_t SplitSkipPrior[] = {
    { LogStopPrior,    LogSplitPrior,   std::log(0.0)   },
    { std::log(0.075), std::log(0.85),  std::log(0.075) },
    { std::log(0.075), std::log(0.85),  std::log(0.075) },
    { std::log(0.075), std::log(0.85),  std::log(0.075) },
    { std::log(0.075), std::log(0.85),  std::log(0.075) },
    { std::log(0.075), std::log(0.85),  std::log(0.075) },
};

// zobrist key generation for context hashing
typedef boost::mt19937 randsrc_t;
typedef boost::variate_generator<randsrc_t &, boost::uniform_int<uint64_t> > zobrng_t;
static randsrc_t randsrc(666);
static boost::uniform_int<uint64_t> uint64_random_range(0, std::numeric_limits<uint64_t>::max());
static zobrng_t zobrng(randsrc, uint64_random_range);

// precomputed static tables
zobhash_t SkipCTS::s_zobtbl[MaxDepth][2];
weight_t SkipCTS::s_log_tbl[LogTblSize];

/* adds two numbers represented in the logarithmic domain */
inline static double fast_logadd(double log_x, double log_y) {

    if (log_x < log_y) {
        return fast_jacoblog(log_y - log_x) + log_x;
    } else {
        return fast_jacoblog(log_x - log_y) + log_y;
    }
}

/* skip nodes */
SkipNode::SkipNode() :
    m_log_prob_est(LogStopPrior),
    m_log_prob_split(LogSplitPrior),
    m_log_prob_weighted(0.0),
    m_log_skip_lik(NULL),
    m_depth(-1),
    m_buf(0.0)
{
    m_count[0] = 0.0f;
    m_count[1] = 0.0f;
}

SkipNode::~SkipNode() {
    delete [] m_log_skip_lik;
}

/* compute the logarithm of the KT-estimator update multiplier */
inline double SkipNode::logKTMul(bit_t b) const {

    return fast_log(ktMul(b));
}

/* compute the logarithm of the KT-estimator update multiplier */
inline double SkipNode::ktMul(bit_t b) const {

    count_t kt_mul_numer = m_count[b] + KT_Alpha;
    count_t kt_mul_denom = m_count[0] + m_count[1] + KT_Alpha2;

    return kt_mul_numer / kt_mul_denom;
}

/* update the KT estimates */
inline void SkipNode::updateKT(bit_t b, double log_est_mul) {

    if (UseDiscounting) {
        m_count[0] *= Discount;
        m_count[1] *= Discount;
    }
    m_count[b]++;
}

/* precompute the zobrist hash keys used for context hashing. */
void SkipCTS::initZobrist() {

    static bool init = false;

    if (!init) {
        for (int i=0; i < SkipCTS::MaxDepth; i++) {
            s_zobtbl[i][0] = zobrng();
            s_zobtbl[i][1] = zobrng();
        }
        init = true;
    }
}

/* initialise a table of precomputed logarithms */
void SkipCTS::initLogTbl() {

    static bool init = false;

    if (!init) {
        for (int i=0; i < SkipCTS::LogTblSize; i++) {
            s_log_tbl[i] = std::log(static_cast<double>(i));
        }
        init = true;
    }
}

/* initialise the hash deltas to save having to recompompute each
   zobrist hash key from scratch. */
void SkipCTS::initHashDeltas() {

    // compute the hash deltas
    indices_t delta;

    for (int i=m_depth; i >= 0; i--) {

        indices_list_t &il = m_indices[i];

        for (int j=0; j < il.size(); j++) {

            indices_t &idxs = il[j];
            indices_t diff(m_depth);

            indices_t::iterator it = std::set_symmetric_difference(
                delta.begin(), delta.end(), idxs.begin(),
                idxs.end(), diff.begin()
            );
            diff.resize(it - diff.begin());
            delta = idxs;
            idxs = diff;
        }
    }
}

/* initialise auxilary information such as number of remaining skips,
   and the depth of the context. */
void SkipCTS::initAuxInfo() {

    // initialise the auxilary information
    for (int i=0; i <= m_depth; i++) {

        m_auxinfo.push_back(aux_info_list_t());
        indices_list_t &il = m_indices[i];

        for (int j=0; j < il.size(); j++) {

            m_auxinfo.back().push_back(aux_info_t());

            // precompute last index
            m_auxinfo.back().back().last_idx   = -1;
            if (!il[j].empty())
                m_auxinfo.back().back().last_idx = static_cast<int>(il[j].back());

            // precompute number of skips remaining
            m_auxinfo.back().back().skips_left = m_skips;
            for (size_t k=0; k < il[j].size(); k++) {
                if ((k == 0 &&  il[j][k] > 0) || (k > 0 &&  il[j][k] >  il[j][k-1]+1))
                    m_auxinfo.back().back().skips_left--;
            }
        }
    }
}

/* compute the indices for the zobrist hashes, on a per depth basis */
void SkipCTS::initIndices() {

    for (size_t i=0; i <= m_depth; i++)
        m_indices.push_back(indices_list_t());

    skip_contexts_t::const_iterator it = m_skip_contexts.begin();
    for ( ; it != m_skip_contexts.end(); ++it) {

        const std::string &s = *it;
        int depth = static_cast<int>(std::count(s.begin(), s.end(), 'b'));
        m_indices[depth].push_back(indices_t());

        for (size_t j=0; j < s.length(); j++) {
            if (s[j] == 'b') m_indices[depth].back().push_back(j);
        }
    }
}

/* precomputes the context indices */
void SkipCTS::initContexts() {

    // generate contexts
    m_skip_contexts.clear();
    genBoundedSkipContexts("", m_depth, m_skips);

    initIndices();

    m_skip_contexts.clear();
}

/* create the skip context tree */
SkipCTS::SkipCTS(history_t &history, int depth, int max_skips, size_t log2_slots) :
    m_history(history),
    m_depth(depth),
    m_skips(max_skips),
    m_log_skip_preds(MaxDepth),
    m_nodes(new SkipNode[size_t(1) << log2_slots]),
    m_mask((size_t(1) << log2_slots)-1)
{
    initZobrist();
    initContexts();
    initAuxInfo();
    initHashDeltas();
    initLogTbl();
}

SkipCTS::~SkipCTS() {
    delete [] m_nodes;
}

/* the logarithm of the probability of all processed experience */
double SkipCTS::logBlockProbability() const {

    return getNode(0, 0, numSubmodels(-1, m_skips)).logProbWeighted();
}

/* compute the switching rate for a given time t */
double SkipCTS::switchRate(size_t t) const {

    return 1.0 / double(t - m_depth + 3);
}

int SkipCTS::numSubmodels(int position, int skips) const {

    int n = m_depth - position;
    return (skips == 0) ? std::min(n, 2) : n;
}

/* the probability of seeing a particular symbol next */
double SkipCTS::prob(bit_t b) {

    // We proceed as with update(), except that we keep track of the symbol probability at
    // each node instead of actually updating parameters
    getContext();

    zobhash_t hash = 0;
    int skips_left, last_idx;
    double symbolLogProb = LogOneHalf;

    // propagate the symbol probability from leaves to root
    for (int i=m_depth; i >= 0; i--) {

        // update the KT statistics, then the weighted
        // probability for every node on this level
        const indices_list_t &il = m_indices[i];

        for (int j=0; j < il.size(); j++) {

            getContextInfo(hash, il[j]);
            skips_left = m_auxinfo[i][j].skips_left;
            last_idx   = m_auxinfo[i][j].last_idx;

            // update the node
            int n_submodels = numSubmodels(last_idx, skips_left);

            SkipNode &n = getNode(hash, i, n_submodels);

            // handle the stop case
            double log_est_mul = n.logKTMul(b);
            if (n_submodels == 1) {
                n.m_buf = symbolLogProb = log_est_mul;
                continue;
            }

            // Here we rely on the property that log_prob_est and the like are unnormalized
            // posteriors. we accumulate the symbol log probability under 'symbolLogProb'
            symbolLogProb = n.m_log_prob_est + log_est_mul;

            // handle the split case
            zobhash_t delta = s_zobtbl[last_idx+1][m_context[last_idx+1]];
            const SkipNode &nn = getNode(hash ^ delta, i+1, numSubmodels(last_idx+1, skips_left));
            // recall that m_buf contains the symbol probability at the child node
            double log_split_pred = nn.m_buf;
            symbolLogProb = fast_logadd(symbolLogProb, n.m_log_prob_split + log_split_pred);

            // handle the skipping case
            if (n_submodels > 2) {

                // if we did not yet allocate these models, this node must never have been
                // updated. We assume (perhaps incorrectly) that none of this node's children
                // exist and pretend they return a symbol probability of 0.5
                if (n.m_log_skip_lik == NULL) {

                    const prior_t &p = SplitSkipPrior[skips_left];
                    symbolLogProb = fast_logadd(symbolLogProb, p.skip + LogOneHalf);
                }

                // mix in the symbol probability from the skipping models
                else for (int k=last_idx+2; k < m_depth; k++) {

                    zobhash_t h = hash ^ s_zobtbl[k][m_context[k]];
                    SkipNode &sn = getNode(h, i+1, numSubmodels(k, skips_left - 1));

                    double log_skip_pred = sn.m_buf;
                    int z = k - last_idx - 2;
                    symbolLogProb = fast_logadd(symbolLogProb, n.m_log_skip_lik[z] + log_skip_pred);
                }
            }

            // Finally we normalize by the mixture probability at this node
            symbolLogProb -= n.m_log_prob_weighted;
            n.m_buf = symbolLogProb;

            assert(n.m_buf < 0.0);
        }
    }

    // our scheme assumes that the last node processed is the root; the variable 'symbolLogProb'
    // contains its symbol probability
    return fast_exp(symbolLogProb);
}

/* gets the node's index into the hash table */
inline SkipNode &SkipCTS::getNode(zobhash_t hash, int depth, int submodels) const {

    size_t key = static_cast<size_t>(hash & m_mask);

    // do a linear scan till we find either an empty slot, or
    // a populated slot with matching depth
    do {
        if (m_nodes[key].m_depth == -1) break;
        if (m_nodes[key].m_depth == depth && m_nodes[key].m_submodels == submodels) break;
        key = (key + 1) & m_mask;
    } while (true);

    // mark the slot as used
    m_nodes[key].m_depth = depth;
    m_nodes[key].m_submodels = submodels;

    return m_nodes[key];
}

/* performs the update operation to maintain a switching posterior. */
void SkipCTS::posteriorUpdate(
    SkipNode &n, double log_scale, double log_alpha,
    double log_K, double log_mul, double &log_post
) const {

    log_post = log_scale +
        fast_logadd(
            log_alpha + n.m_log_prob_weighted,
            log_K + log_post + log_mul
        );
}

/* lazy allocation of skipping posterior weights. */
void SkipCTS::lazyAllocate(SkipNode &n, int n_submodels, int skips_left) const {

    n.m_log_skip_lik = new weight_t[n_submodels-2];

    const prior_t &p = SplitSkipPrior[skips_left];
    n.m_log_prob_est    = p.stop;
    n.m_log_prob_split  = p.split;
    for (int k=0; k < n_submodels-2; k++) {
        n.m_log_skip_lik[k] = p.skip - s_log_tbl[n_submodels-2];
    }
}

/* compute the information needed to update the current context stats. */
void SkipCTS::getContextInfo(zobhash_t &hash, const indices_t &idxs) const {

    // compute the hash
    for (int k=0; k < idxs.size(); k++) {
        size_t x = idxs[k];
        hash ^= s_zobtbl[x][m_context[x]];
    }
}

/* process a new piece of sensory experience */
void SkipCTS::update(bit_t b) {

    getContext();

    double alpha = switchRate(m_history.size());
    double log_alpha = fast_log(alpha);

    zobhash_t hash = 0;
    int skips_left, last_idx;

    // update nodes from deepest to shallowest
    for (int i=m_depth; i >= 0; i--) {

        // update the KT statistics, then the weighted
        // probability for every node on this level
        const indices_list_t &il = m_indices[i];

        for (int j=0; j < il.size(); j++) {

            m_log_skip_preds.clear();

            getContextInfo(hash, il[j]);
            skips_left = m_auxinfo[i][j].skips_left;
            last_idx   = m_auxinfo[i][j].last_idx;

            // update the node
            int n_submodels = numSubmodels(last_idx, skips_left);

            SkipNode &n = getNode(hash, i, n_submodels);
            n.m_buf = n.m_log_prob_weighted;

            // lazy allocation of skipping prior weights
            if (n_submodels > 2 && n.m_log_skip_lik == NULL)
                lazyAllocate(n, n_submodels, skips_left);

            // handle the stop case
            double log_est_mul = n.logKTMul(b);
            if (n_submodels == 1) {
                n.updateKT(b, log_est_mul);
                n.m_log_prob_weighted += log_est_mul;
                n.m_buf = log_est_mul;
                continue;
            }

            double log_acc = n.m_log_prob_est + log_est_mul;
            n.updateKT(b, log_est_mul);

            // handle the split case
            zobhash_t delta = s_zobtbl[last_idx+1][m_context[last_idx+1]];
            const SkipNode &nn = getNode(hash ^ delta, i+1, numSubmodels(last_idx+1, skips_left));
            double log_split_pred = nn.m_buf;
            log_acc = fast_logadd(log_acc, n.m_log_prob_split + log_split_pred);

            // handle the skipping case
            if (n_submodels > 2) {

                // update the skipping models
                for (int k=last_idx+2; k < m_depth; k++) {

                    zobhash_t h = hash ^ s_zobtbl[k][m_context[k]];
                    SkipNode &sn = getNode(h, i+1, numSubmodels(k, skips_left - 1));

                    double log_skip_pred = sn.m_buf;
                    m_log_skip_preds.push_back(log_skip_pred);
                    int z = k - last_idx - 2;
                    log_acc = fast_logadd(log_acc, n.m_log_skip_lik[z] + log_skip_pred);
                }
            }

            // store the weighted probability
            n.m_log_prob_weighted = log_acc;

            assert(n.m_log_prob_weighted < n.m_buf);
            // Store the *difference* in log probability in m_buf
            n.m_buf = n.m_log_prob_weighted - n.m_buf;

            updatePosteriors(n, n_submodels, alpha, log_alpha, log_est_mul, log_split_pred);
        }
    }

    m_history.push_back(b != 0);
}

/* update the switching posterior weights */
void SkipCTS::updatePosteriors(SkipNode &n, int n_submodels, double alpha,
    double log_alpha, double log_stop_mul, double log_split_mul) {

    // update switching log-posteriors
    double dn        =  static_cast<double>(n_submodels);
    double K         =  (1.0 - alpha) * dn - 1.0;
    double log_K     =  fast_log(K);
    double log_scale = -s_log_tbl[n_submodels-1];

    posteriorUpdate(n, log_scale, log_alpha, log_K, log_stop_mul, n.m_log_prob_est);
    posteriorUpdate(n, log_scale, log_alpha, log_K, log_split_mul, n.m_log_prob_split);
    for (int k=0; k < m_log_skip_preds.size(); k++) {
        posteriorUpdate(n, log_scale, log_alpha, log_K, m_log_skip_preds[k], n.m_log_skip_lik[k]);
    }
}

/* generate possible context strings with bounded # of skips */
void SkipCTS::genBoundedSkipContexts(std::string buf, int depth, int skips) {

    m_skip_contexts.push_back(buf);

    if (depth == 0) return;

    if (skips > 0) {
        std::string skipstr = "*";
        for (int i=0; i < depth-1; i++) {
            genBoundedSkipContexts(buf + skipstr + "b", depth - (i+1) - 1, skips-1);
            skipstr.append("*");
        }
    }

    genBoundedSkipContexts(buf + "b", depth - 1, skips);
}

/* prints a set of skips contexts */
void SkipCTS::printContexts(skip_contexts_t &l) const {

    skip_contexts_t::const_iterator it = l.begin();
    for (; it != l.end(); ++it) {
        std::cout << *it << std::endl;
    }
}

/* compute the current binary context */
inline void SkipCTS::getContext() {

    size_t offset = m_history.size();
    m_context.clear();

    for (size_t i=0; i < m_depth; ++i) {
        m_context.push_back(m_history[offset-i-1]);
    }
}


/*
 * EigenProxy.hpp
 *
 *  Created on: 12 Apr 2018
 *      Author: Viktor Csomor
 */

#ifndef C_ATTL3_CORE_EIGENPROXY_H_
#define C_ATTL3_CORE_EIGENPROXY_H_

#define EIGEN_USE_THREADS

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <Eigen/Dense>
#include <iostream>
#include <string>
#include <thread>
#include <unsupported/Eigen/CXX11/Tensor>

/**
 * The namespace containing all classes and typedefs of the C-ATTL3 library.
 */
namespace cattle {

/**
 * An alias for a single row matrix of an arbitrary scalar type.
 */
template<typename Scalar>
using RowVector = Eigen::Matrix<Scalar,1,Eigen::Dynamic,Eigen::RowMajor, 1,Eigen::Dynamic>;

/**
 * An alias for a single column matrix of an arbitrary scalar type.
 */
template <typename Scalar>
using ColVector = Eigen::Matrix<Scalar,Eigen::Dynamic,1,Eigen::ColMajor,Eigen::Dynamic,1>;

/**
 * An alias for a dynamically sized matrix of an arbitrary scalar type.
 */
template<typename Scalar>
using Matrix = Eigen::Matrix<Scalar,Eigen::Dynamic,Eigen::Dynamic,Eigen::ColMajor,
		Eigen::Dynamic,Eigen::Dynamic>;

/**
 * An alias for a class that can be used to map raw pointer data to a dynamically
 * sized Matrix of an arbitrary scalar type.
 */
template<typename Scalar>
using MatrixMap = Eigen::Map<Matrix<Scalar>>;

/**
 * An alias for a tensor of arbitrary rank and scalar type with dynamic dimensionality.
 */
template<typename Scalar, std::size_t Rank>
using Tensor = Eigen::Tensor<Scalar,Rank,Eigen::ColMajor,std::size_t>;

/**
 * An for a class that can be used to map raw pointer data to a tensor of arbitrary
 * rank and scalar type with dynamic dimensionality.
 */
template<typename Scalar, std::size_t Rank>
using TensorMap = Eigen::TensorMap<Tensor<Scalar,Rank>>;

/**
 * An alias for permutation matrices.
 */
using PermMatrix = Eigen::PermutationMatrix<Eigen::Dynamic,Eigen::Dynamic>;

/**
 * An alias for self-adjoint eigen solvers.
 */
template<typename Scalar>
using EigenSolver = Eigen::SelfAdjointEigenSolver<Matrix<Scalar>>;

/**
 * An alias for Eigen's bi-diagonal divide and conquer singular-value decomposition.
 */
template<typename Scalar>
using SVD = Eigen::BDCSVD<Matrix<Scalar>>;

/**
 * An alias for Eigen's singular-value decomposition options.
 */
using SVDOptions = Eigen::DecompositionOptions;

/**
 * @return The number of threads used by Eigen to accelerate operations
 * supporting multithreading.
 */
inline int num_of_eval_threads() {
	return Eigen::nbThreads();
}

/**
 * @param num_of_threads The number of threads Eigen should use to accelerate
 * operations supporting multithreading. The lower bound of the actual value
 * applied is 1 while the upper bound is the maximum of 1 and the level of
 * hardware concurrency detected.
 */
inline void set_num_of_eval_threads(int num_of_threads) {
	int max = std::max(1, (int) std::thread::hardware_concurrency());
	Eigen::setNbThreads(std::max(1, std::min(num_of_threads, max)));
}

/**
 * It serializes the matrix in a format such that the first two numbers denote the
 * matrix's number of rows and columns respectively and the remaining numbers represent
 * the coefficients of the matrix in column-major order.
 *
 * @param matrix The matrix to serialize.
 * @param out_stream The non-binary stream to serialize the matrix to.
 */
template<typename Scalar>
inline void serialize(const Matrix<Scalar>& matrix, std::ostream& out_stream) {
	out_stream << sizeof(Scalar);
	out_stream << " " << matrix.rows();
	out_stream << " " << matrix.cols();
	for (std::size_t i = 0; i < matrix.size(); ++i)
		out_stream << " " << *(matrix.data() + i);
	out_stream << std::flush;
}

/**
 * It serializes the matrix into a file at the specified file path.
 *
 * @param matrix The matrix to serialize.
 * @param file_path The path to the file to which the matrix is to be serialized.
 */
template<typename Scalar>
inline void serialize(const Matrix<Scalar>& matrix, const std::string& file_path) {
	std::ofstream out_stream(file_path);
	assert(out_stream.is_open());
	serialize<Scalar>(matrix, out_stream);
}

/**
 * It serializes the matrix in a format such that the first 2 bytes denote the size of
 * a single coefficient of the matrix in bytes, the second and third 4 bytes denote the
 * matrix's number of rows and columns respectively, and the remaining bytes contain
 * the coefficients of the matrix in column-major order.
 *
 * @param matrix The matrix to serialize.
 * @param out_stream The binary stream to serialize the matrix to.
 */
template<typename Scalar>
inline void serialize_binary(const Matrix<Scalar>& matrix, std::ostream& out_stream) {
	unsigned short scalar_size = static_cast<unsigned short>(sizeof(Scalar));
	out_stream.write(reinterpret_cast<const char*>(&scalar_size),
			std::streamsize(sizeof(unsigned short)));
	unsigned rows = static_cast<unsigned>(matrix.rows());
	unsigned cols = static_cast<unsigned>(matrix.cols());
	out_stream.write(reinterpret_cast<const char*>(&rows), std::streamsize(sizeof(unsigned)));
	out_stream.write(reinterpret_cast<const char*>(&cols), std::streamsize(sizeof(unsigned)));
	out_stream.write(reinterpret_cast<const char*>(matrix.data()),
			std::streamsize(matrix.size() * sizeof(Scalar)));
	out_stream << std::flush;
}

/**
 * It serializes the matrix into a binary file at the specified file path.
 *
 * @param matrix The matrix to serialize.
 * @param file_path The path to the binary file to which the matrix is to be serialized.
 */
template<typename Scalar>
inline void serialize_binary(const Matrix<Scalar>& matrix, const std::string& file_path) {
	std::ofstream out_stream(file_path, std::ios::binary);
	assert(out_stream.is_open());
	serialize_binary<Scalar>(matrix, out_stream);
}

/**
 * It deserializes a matrix assuming the serialized format matches that used by the
 * serialize() method.
 *
 * @param in_stream The stream to the serialized matrix.
 * @return The unserialized matrix.
 */
template<typename Scalar>
inline Matrix<Scalar> deserialize(std::istream& in_stream) {
	unsigned rows, cols;
	in_stream >> rows;
	in_stream >> cols;
	Matrix<Scalar> matrix(rows, cols);
	for (std::size_t i = 0; i < matrix.size(); ++i)
		in_stream >> *(matrix.data() + i);
	return matrix;
}

/**
 * It deserializes a matrix from the file at the provided file path.
 *
 * @param file_path The path to the file containing the serialized matrix.
 * @return The deserialized matrix.
 */
template<typename Scalar>
inline Matrix<Scalar> deserialize(const std::string& file_path) {
	std::ifstream in_stream(file_path);
	assert(in_stream.is_open());
	return deserialize<Scalar>(in_stream);
}

/**
 * It deserializes a matrix assuming the serialized format matches that used by the
 * serialize_binary() method.
 *
 * @param in_stream The binary stream to the serialized matrix.
 * @return The unserialized matrix.
 */
template<typename Scalar>
inline Matrix<Scalar> deserialize_binary(std::istream& in_stream) {
	unsigned short scalar_size;
	in_stream.read(reinterpret_cast<char*>(&scalar_size), std::streamsize(sizeof(unsigned short)));
	assert(scalar_size == sizeof(Scalar));
	unsigned rows, cols;
	in_stream.read(reinterpret_cast<char*>(&rows), std::streamsize(sizeof(unsigned)));
	in_stream.read(reinterpret_cast<char*>(&cols), std::streamsize(sizeof(unsigned)));
	Matrix<Scalar> matrix(rows, cols);
	in_stream.read(reinterpret_cast<char*>(matrix.data()),
			std::streamsize(matrix.size() * sizeof(Scalar)));
	return matrix;
}

/**
 * It deserializes a matrix from the binary file at the provided file path.
 *
 * @param file_path The path to the binary file containing the serialized matrix.
 * @return The deserialized matrix.
 */
template<typename Scalar>
inline Matrix<Scalar> deserialize_binary(const std::string& file_path) {
	std::ifstream in_stream(file_path, std::ios::binary);
	assert(in_stream.is_open());
	return deserialize_binary<Scalar>(in_stream);
}

}

#endif /* C_ATTL3_CORE_EIGENPROXY_H_ */


/*
 * GridMapMath.hpp
 *
 *  Created on: Dec 2, 2013
 *      Author: Péter Fankhauser
 *	 Institute: ETH Zurich, Autonomous Systems Lab
 */

#pragma once

#include "grid_map_core/TypeDefs.hpp"
#include "grid_map_core/BufferRegion.hpp"

#include <Eigen/Core>
#include <vector>
#include <map>

namespace grid_map {

/*!
 * Gets the position of a cell specified by its index in the map frame.
 * @param[out] position the position of the center of the cell in the map frame.
 * @param[in] index of the cell.
 * @param[in] mapLength the lengths in x and y direction.
 * @param[in] mapPosition the position of the map.
 * @param[in] resolution the resolution of the map.
 * @param[in] bufferSize the size of the buffer (optional).
 * @param[in] bufferStartIndex the index of the starting point of the circular buffer (optional).
 * @return true if successful, false if index not within range of buffer.
 */
bool getPositionFromIndex(Eigen::Vector2d& position,
                          const Eigen::Array2i& index,
                          const Eigen::Array2d& mapLength,
                          const Eigen::Vector2d& mapPosition,
                          const double& resolution,
                          const Eigen::Array2i& bufferSize,
                          const Eigen::Array2i& bufferStartIndex = Eigen::Array2i::Zero());

/*!
 * Gets the index of the cell which contains a position in the map frame.
 * @param[out] index of the cell.
 * @param[in] position the position in the map frame.
 * @param[in] mapLength the lengths in x and y direction.
 * @param[in] mapPosition the position of the map.
 * @param[in] resolution the resolution of the map.
 * @param[in] bufferSize the size of the buffer (optional).
 * @param[in] bufferStartIndex the index of the starting point of the circular buffer (optional).
 * @return true if successful, false if position outside of map.
 */
bool getIndexFromPosition(Eigen::Array2i& index,
                          const Eigen::Vector2d& position,
                          const Eigen::Array2d& mapLength,
                          const Eigen::Vector2d& mapPosition,
                          const double& resolution,
                          const Eigen::Array2i& bufferSize,
                          const Eigen::Array2i& bufferStartIndex = Eigen::Array2i::Zero());

/*!
 * Checks if position is within the map boundaries.
 * @param[in] position the position which is to be checked.
 * @param[in] mapLength the length of the map.
 * @param[in] mapPosition the position of the map.
 * @return true if position is within map, false otherwise.
 */
bool checkIfPositionWithinMap(const Eigen::Vector2d& position,
                              const Eigen::Array2d& mapLength,
                              const Eigen::Vector2d& mapPosition);

/*!
 * Gets the position of the data structure origin.
 * @param[in] position the position of the map.
 * @param[in] mapLength the map length.
 * @param[out] positionOfOrigin the position of the data structure origin.
 */
void getPositionOfDataStructureOrigin(const Eigen::Vector2d& position,
                                      const Eigen::Array2d& mapLength,
                                      Eigen::Vector2d& positionOfOrigin);

/*!
 * Computes how many cells/indeces the map is moved based on a position shift in
 * the grid map frame. Use this function if you are moving the grid map
 * and want to ensure that the cells match before and after.
 * @param[out] indexShift the corresponding shift of the indices.
 * @param[in] positionShift the desired position shift.
 * @param[in] resolution the resolution of the map.
 * @return true if successful.
 */
bool getIndexShiftFromPositionShift(Eigen::Array2i& indexShift,
                                    const Eigen::Vector2d& positionShift,
                                    const double& resolution);

/*!
 * Computes the corresponding position shift from a index shift. Use this function
 * if you are moving the grid map and want to ensure that the cells match
 * before and after.
 * @param[out] positionShift the corresponding shift in position in the grid map frame.
 * @param[in] indexShift the desired shift of the indeces.
 * @param[in] resolution the resolution of the map.
 * @return true if successful.
 */
bool getPositionShiftFromIndexShift(Eigen::Vector2d& positionShift,
                                    const Eigen::Array2i& indexShift,
                                    const double& resolution);

/*!
 * Checks if index is within range of the buffer.
 * @param[in] index to check.
 * @param[in] bufferSize the size of the buffer.
 * @return true if index is within, and false if index is outside of the buffer.
 */
bool checkIfIndexWithinRange(const Eigen::Array2i& index, const Eigen::Array2i& bufferSize);

/*!
 * Maps an index that runs out of the range of the circular buffer back into allowed the region.
 * This is the 2d version of mapIndexWithinRange(int&, const int&).
 * @param[in/out] index the indeces that will be mapped into the valid region of the buffer.
 * @param[in] bufferSize the size of the buffer.
 */
void mapIndexWithinRange(Eigen::Array2i& index,
                         const Eigen::Array2i& bufferSize);

/*!
 * Maps an index that runs out of the range of the circular buffer back into allowed the region.
 * @param[in/out] index the index that will be mapped into the valid region of the buffer.
 * @param[in] bufferSize the size of the buffer.
 */
void mapIndexWithinRange(int& index, const int& bufferSize);

/*!
 * Limits (cuts off) the position to lie inside the map.
 * @param[in/out] position the position to be limited.
 * @param[in] mapLength the lengths in x and y direction.
 * @param[in] mapPosition the position of the map.
 */
void limitPositionToRange(Eigen::Vector2d& position,
                          const Eigen::Array2d& mapLength,
                          const Eigen::Vector2d& mapPosition);

/*!
 * Provides the alignment transformation from the buffer order (outer/inner storage)
 * and the map frame (x/y-coordinate).
 * @return the alignment transformation.
 */
const Eigen::Matrix2i getBufferOrderToMapFrameAlignment();

/*!
 * Given a map and a desired submap (defined by position and size), this function computes
 * various information about the submap. The returned submap might be smaller than the requested
 * size as it respects the boundaries of the map.
 * @param[out] submapTopLeftIndex the top left index of the returned submap.
 * @param[out] submapBufferSize the buffer size of the returned submap.
 * @param[out] submapPosition the position of the submap (center) in the map frame.
 * @param[out] submapLength the length of the submap.
 * @param[out] requestedIndexInSubmap the index in the submap that corresponds to the requested
 *             position of the submap.
 * @param[in] requestedSubmapPosition the requested submap position (center) in the map frame.
 * @param[in] requestedSubmapLength the requested submap length.
 * @param[in] mapLength the lengths in x and y direction.
 * @param[in] mapPosition the position of the map.
 * @param[in] resolution the resolution of the map.
 * @param[in] bufferSize the buffer size of the map.
 * @param[in] bufferStartIndex the index of the starting point of the circular buffer (optional).
 * @return true if successful.
 */
bool getSubmapInformation(Eigen::Array2i& submapTopLeftIndex,
                          Eigen::Array2i& submapBufferSize,
                          Eigen::Vector2d& submapPosition,
                          Eigen::Array2d& submapLength,
                          Eigen::Array2i& requestedIndexInSubmap,
                          const Eigen::Vector2d& requestedSubmapPosition,
                          const Eigen::Vector2d& requestedSubmapLength,
                          const Eigen::Array2d& mapLength,
                          const Eigen::Vector2d& mapPosition,
                          const double& resolution,
                          const Eigen::Array2i& bufferSize,
                          const Eigen::Array2i& bufferStartIndex = Eigen::Array2i::Zero());

/*!
 * Computes the regions in the circular buffer that make up the data for
 * a requested submap.
 * @param[out] submapBufferRegions the list of buffer regions that make up the submap.
 * @param[in] submapIndex the index (top-left) for the requested submap.
 * @param[in] submapBufferSize the size of the requested submap.
 * @param[in] bufferSize the buffer size of the map.
 * @param[in] bufferStartIndex the index of the starting point of the circular buffer (optional).
 * @return true if successful, false if requested submap is not fully contained in the map.
 */
bool getBufferRegionsForSubmap(std::vector<BufferRegion>& submapBufferRegions,
                               const Index& submapIndex,
                               const Size& submapBufferSize,
                               const Size& bufferSize,
                               const Index& bufferStartIndex = Index::Zero());

/*!
 * Increases the index by one to iterate through the map.
 * Increments either to the neighboring index to the right or to
 * the start of the lower row. Returns false if end of iterations are reached.
 * @param[in/out] index the index in the map that is incremented (corrected for the circular buffer).
 * @param[in] bufferSize the map buffer size.
 * @param[in] bufferStartIndex the map buffer start index.
 * @return true if successfully incremented indeces, false if end of iteration limits are reached.
 */
bool incrementIndex(Eigen::Array2i& index, const Eigen::Array2i& bufferSize,
                    const Eigen::Array2i& bufferStartIndex = Eigen::Array2i::Zero());

/*!
 * Increases the index by one to iterate through the cells of a submap.
 * Increments either to the neighboring index to the right or to
 * the start of the lower row. Returns false if end of iterations are reached.
 *
 * Note: This function does not check if submap actually fits to the map. This needs
 * to be checked before separately.
 *
 * @param[in/out] submapIndex the index in the submap that is incremented.
 * @param[out] index the index in the map that is incremented (corrected for the circular buffer).
 * @param[in] submapTopLefIndex the top left index of the submap.
 * @param[in] submapBufferSize the submap buffer size.
 * @param[in] bufferSize the map buffer size.
 * @param[in] bufferStartIndex the map buffer start index.
 * @return true if successfully incremented indeces, false if end of iteration limits are reached.
 */
bool incrementIndexForSubmap(Eigen::Array2i& submapIndex, Eigen::Array2i& index,
                             const Eigen::Array2i& submapTopLeftIndex,
                             const Eigen::Array2i& submapBufferSize,
                             const Eigen::Array2i& bufferSize,
                             const Eigen::Array2i& bufferStartIndex = Eigen::Array2i::Zero());

/*!
 * Retrieve the index as unwrapped index, i.e., as the corresponding index of a
 * grid map with no circular buffer offset.
 * @param bufferIndex the index in the circular buffer.
 * @param bufferSize the map buffer size.
 * @param bufferStartIndex the map buffer start index.
 * @return the unwrapped index.
 */
Index getIndexFromBufferIndex(const Index& bufferIndex, const Size& bufferSize,
                              const Index& bufferStartIndex);

/*!
 * Returns the 1d index corresponding to the 2d index for either row- or column-major format.
 * Note: Eigen is defaulting to column-major format.
 * @param[in] index the 2d index.
 * @param[in] bufferSize the map buffer size.
 * @param[in] (optional) rowMajor if the 1d index is generated for row-major format.
 * @return the 1d index.
 */
unsigned int get1dIndexFrom2dIndex(const Index& index, const Size& bufferSize,
                                   const bool rowMajor);

/*!
 * Generates a list of indices for a region in the map.
 * @param regionIndex the region top-left index.
 * @param regionSize the region size.
 * @param indices the list of indices of the region.
 */
void getIndicesForRegion(const Index& regionIndex, const Size& regionSize,
                         std::vector<Index> indices);

/*!
 * Generates a list of indices for multiple regions in the map.
 * This method makes sure every index is only once contained in the list.
 * @param regionIndeces the regions' top-left index.
 * @param regionSizes the regions' sizes.
 * @param indices the list of indices of the regions.
 */
void getIndicesForRegions(const std::vector<Index>& regionIndeces, const Size& regionSizes,
                          std::vector<Index> indices);

} // namespace


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file    Rot3.cpp
 * @brief   Rotation, common code between Rotation matrix and Quaternion
 * @author  Alireza Fathi
 * @author  Christian Potthast
 * @author  Frank Dellaert
 * @author  Richard Roberts
 * @author  Varun Agrawal
 */

#include <gtsam/geometry/Rot3.h>
#include <gtsam/geometry/SO3.h>
#include <boost/math/constants/constants.hpp>

#include <cmath>
#include <random>

using namespace std;

namespace gtsam {

/* ************************************************************************* */
void Rot3::print(const std::string& s) const {
  cout << (s.empty() ? "R: " : s + " ");
  gtsam::print(static_cast<Matrix>(matrix()));
}

/* ************************************************************************* */
Rot3 Rot3::Random(std::mt19937& rng) {
  Unit3 axis = Unit3::Random(rng);
  uniform_real_distribution<double> randomAngle(-M_PI, M_PI);
  double angle = randomAngle(rng);
  return AxisAngle(axis, angle);
}

/* ************************************************************************* */
Rot3 Rot3::AlignPair(const Unit3& axis, const Unit3& a_p, const Unit3& b_p) {
  // if a_p is already aligned with b_p, return the identity rotation
  if (std::abs(a_p.dot(b_p)) > 0.999999999) {
    return Rot3();
  }

  // Check axis was not degenerate cross product
  const Vector3 z = axis.unitVector();
  if (z.hasNaN())
    throw std::runtime_error("AlignSinglePair: axis has Nans");

  // Now, calculate rotation that takes b_p to a_p
  const Matrix3 P = I_3x3 - z * z.transpose();  // orthogonal projector
  const Vector3 a_po = P * a_p.unitVector();    // point in a orthogonal to axis
  const Vector3 b_po = P * b_p.unitVector();    // point in b orthogonal to axis
  const Vector3 x = a_po.normalized();          // x-axis in axis-orthogonal plane, along a_p vector
  const Vector3 y = z.cross(x);                 // y-axis in axis-orthogonal plane
  const double u = x.dot(b_po);                 // x-coordinate for b_po
  const double v = y.dot(b_po);                 // y-coordinate for b_po
  double angle = std::atan2(v, u);
  return Rot3::AxisAngle(z, -angle);
}

/* ************************************************************************* */
Rot3 Rot3::AlignTwoPairs(const Unit3& a_p, const Unit3& b_p,  //
                         const Unit3& a_q, const Unit3& b_q) {
  // there are three frames in play:
  // a: the first frame in which p and q are measured
  // b: the second frame in which p and q are measured
  // i: intermediate, after aligning first pair

  // First, find rotation around that aligns a_p and b_p
  Rot3 i_R_b = AlignPair(a_p.cross(b_p), a_p, b_p);

  // Rotate points in frame b to the intermediate frame,
  // in which we expect the point p to be aligned now
  Unit3 i_q = i_R_b * b_q;
  assert(assert_equal(a_p, i_R_b * b_p, 1e-6));

  // Now align second pair: we need to align i_q to a_q
  Rot3 a_R_i = AlignPair(a_p, a_q, i_q);
  assert(assert_equal(a_p, a_R_i * a_p, 1e-6));
  assert(assert_equal(a_q, a_R_i * i_q, 1e-6));

  // The desired rotation is the product of both
  Rot3 a_R_b = a_R_i * i_R_b;
  return a_R_b;
}

/* ************************************************************************* */
bool Rot3::equals(const Rot3 & R, double tol) const {
  return equal_with_abs_tol(matrix(), R.matrix(), tol);
}

/* ************************************************************************* */
Point3 Rot3::operator*(const Point3& p) const {
  return rotate(p);
}

/* ************************************************************************* */
Unit3 Rot3::rotate(const Unit3& p,
    OptionalJacobian<2,3> HR, OptionalJacobian<2,2> Hp) const {
  Matrix32 Dp;
  Unit3 q = Unit3(rotate(p.point3(Hp ? &Dp : 0)));
  if (Hp) *Hp = q.basis().transpose() * matrix() * Dp;
  if (HR) *HR = -q.basis().transpose() * matrix() * p.skew();
  return q;
}

/* ************************************************************************* */
Unit3 Rot3::unrotate(const Unit3& p,
    OptionalJacobian<2,3> HR, OptionalJacobian<2,2> Hp) const {
  Matrix32 Dp;
  Unit3 q = Unit3(unrotate(p.point3(Dp)));
  if (Hp) *Hp = q.basis().transpose() * matrix().transpose () * Dp;
  if (HR) *HR = q.basis().transpose() * q.skew();
  return q;
}

/* ************************************************************************* */
Unit3 Rot3::operator*(const Unit3& p) const {
  return rotate(p);
}

/* ************************************************************************* */
// see doc/math.lyx, SO(3) section
Point3 Rot3::unrotate(const Point3& p, OptionalJacobian<3,3> H1,
    OptionalJacobian<3,3> H2) const {
  const Matrix3& Rt = transpose();
  Point3 q(Rt * p); // q = Rt*p
  const double wx = q.x(), wy = q.y(), wz = q.z();
  if (H1)
    *H1 << 0.0, -wz, +wy, +wz, 0.0, -wx, -wy, +wx, 0.0;
  if (H2)
    *H2 = Rt;
  return q;
}

/* ************************************************************************* */
Point3 Rot3::column(int index) const{
  if(index == 3)
    return r3();
  else if(index == 2)
    return r2();
  else if(index == 1)
    return r1(); // default returns r1
  else
    throw invalid_argument("Argument to Rot3::column must be 1, 2, or 3");
}

/* ************************************************************************* */
Vector3 Rot3::xyz(OptionalJacobian<3, 3> H) const {
  Matrix3 I;Vector3 q;
  if (H) {
    Matrix93 mH;
    const auto m = matrix();
#ifdef GTSAM_USE_QUATERNIONS
    SO3{m}.vec(mH);
#else
    rot_.vec(mH);
#endif

    Matrix39 qHm;
    boost::tie(I, q) = RQ(m, qHm);

    // TODO : Explore whether this expression can be optimized as both
    // qHm and mH are super-sparse
    *H = qHm * mH;
  } else
    boost::tie(I, q) = RQ(matrix());
  return q;
}

/* ************************************************************************* */
Vector3 Rot3::ypr(OptionalJacobian<3, 3> H) const {
  Vector3 q = xyz(H);
  if (H) H->row(0).swap(H->row(2));

  return Vector3(q(2),q(1),q(0));
}

/* ************************************************************************* */
Vector3 Rot3::rpy(OptionalJacobian<3, 3> H) const { return xyz(H); }

/* ************************************************************************* */
double Rot3::roll(OptionalJacobian<1, 3> H) const {
  double r;
  if (H) {
    Matrix3 xyzH;
    r = xyz(xyzH)(0);
    *H = xyzH.row(0);
  } else
    r = xyz()(0);
  return r;
}

/* ************************************************************************* */
double Rot3::pitch(OptionalJacobian<1, 3> H) const {
  double p;
  if (H) {
    Matrix3 xyzH;
    p = xyz(xyzH)(1);
    *H = xyzH.row(1);
  } else
    p = xyz()(1);
  return p;
}

/* ************************************************************************* */
double Rot3::yaw(OptionalJacobian<1, 3> H) const {
  double y;
  if (H) {
    Matrix3 xyzH;
    y = xyz(xyzH)(2);
    *H = xyzH.row(2);
  } else
    y = xyz()(2);
  return y;
}

/* ************************************************************************* */
#ifdef GTSAM_ALLOW_DEPRECATED_SINCE_V42
Vector Rot3::quaternion() const {
  gtsam::Quaternion q = toQuaternion();
  Vector v(4);
  v(0) = q.w();
  v(1) = q.x();
  v(2) = q.y();
  v(3) = q.z();
  return v;
}
#endif

/* ************************************************************************* */
pair<Unit3, double> Rot3::axisAngle() const {
  const Vector3 omega = Rot3::Logmap(*this);
  return std::pair<Unit3, double>(Unit3(omega), omega.norm());
}

/* ************************************************************************* */
Matrix3 Rot3::ExpmapDerivative(const Vector3& x) {
  return SO3::ExpmapDerivative(x);
}

/* ************************************************************************* */
Matrix3 Rot3::LogmapDerivative(const Vector3& x)    {
  return SO3::LogmapDerivative(x);
}

/* ************************************************************************* */
pair<Matrix3, Vector3> RQ(const Matrix3& A, OptionalJacobian<3, 9> H) {
  const double x = -atan2(-A(2, 1), A(2, 2));
  const auto Qx = Rot3::Rx(-x).matrix();
  const Matrix3 B = A * Qx;

  const double y = -atan2(B(2, 0), B(2, 2));
  const auto Qy = Rot3::Ry(-y).matrix();
  const Matrix3 C = B * Qy;

  const double z = -atan2(-C(1, 0), C(1, 1));
  const auto Qz = Rot3::Rz(-z).matrix();
  const Matrix3 R = C * Qz;

  if (H) {
    if (std::abs(y - M_PI / 2) < 1e-2)
      throw std::runtime_error(
          "Rot3::RQ : Derivative undefined at singularity (gimbal lock)");

    auto atan_d1 = [](double y, double x) { return x / (x * x + y * y); };
    auto atan_d2 = [](double y, double x) { return -y / (x * x + y * y); };

    const auto sx = -Qx(2, 1), cx = Qx(1, 1);
    const auto sy = -Qy(0, 2), cy = Qy(0, 0);

    *H = Matrix39::Zero();
    // First, calculate the derivate of x
    (*H)(0, 5) = atan_d1(A(2, 1), A(2, 2));
    (*H)(0, 8) = atan_d2(A(2, 1), A(2, 2));

    // Next, calculate the derivate of y. We have
    // b20 = a20 and b22 = a21 * sx + a22 * cx
    (*H)(1, 2) = -atan_d1(B(2, 0), B(2, 2));
    const auto yHb22 = -atan_d2(B(2, 0), B(2, 2));
    (*H)(1, 5) = yHb22 * sx;
    (*H)(1, 8) = yHb22 * cx;

    // Next, calculate the derivate of z. We have
    // c10 = a10 * cy + a11 * sx * sy + a12 * cx * sy
    // c11 = a11 * cx - a12 * sx
    const auto c10Hx = (A(1, 1) * cx - A(1, 2) * sx) * sy;
    const auto c10Hy = A(1, 2) * cx * cy + A(1, 1) * cy * sx - A(1, 0) * sy;
    Vector9 c10HA = c10Hx * H->row(0) + c10Hy * H->row(1);
    c10HA[1] = cy;
    c10HA[4] = sx * sy;
    c10HA[7] = cx * sy;

    const auto c11Hx = -A(1, 2) * cx - A(1, 1) * sx;
    Vector9 c11HA = c11Hx * H->row(0);
    c11HA[4] = cx;
    c11HA[7] = -sx;

    H->block<1, 9>(2, 0) =
        atan_d1(C(1, 0), C(1, 1)) * c10HA + atan_d2(C(1, 0), C(1, 1)) * c11HA;
  }

  const auto xyz = Vector3(x, y, z);
  return make_pair(R, xyz);
}

/* ************************************************************************* */
ostream &operator<<(ostream &os, const Rot3& R) {
  os << R.matrix().format(matlabFormat());
  return os;
}

/* ************************************************************************* */
Rot3 Rot3::slerp(double t, const Rot3& other) const {
  return interpolate(*this, other, t);
}

/* ************************************************************************* */

} // namespace gtsam