#include <mutex>
#include <thread>
#include <queue>
#include <map>
#include <vector>
#include <iostream>
#include <chrono>

// ROS
#include <ros/ros.h>
#include <tf/transform_datatypes.h>
#include <tf/transform_broadcaster.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include <geometry_msgs/PoseStamped.h>

// PCL
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/common/eigen.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/point_types.h>

// Eigen
#include <eigen3/Eigen/Dense>

// GTSAM
#include <gtsam/nonlinear/ISAM2.h>
#include <gtsam/nonlinear/ISAM2Params.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/Marginals.h>
#include <gtsam/nonlinear/LevenbergMarquardtOptimizer.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/geometry/Pose3.h>
#include <gtsam/geometry/Rot3.h>
#include <gtsam/linear/NoiseModel.h>

using namespace gtsam;

/**
 * A dedicated class storing position + orientation + time.
 * Public fields allow direct use in PCL macros.
 */
class PosePoint
{
public:
  // Required by PCL macros
  float  x       = 0.f;
  float  y       = 0.f;
  float  z       = 0.f;
  float  intensity = 0.f;
  float  roll    = 0.f;
  float  pitch   = 0.f;
  float  yaw     = 0.f;
  double time    = 0.0;

  PosePoint() = default;

  // Named constructor for direct (x,y,z,roll,pitch,yaw,intensity,time).
  static PosePoint fromXYZRPY(float x, float y, float z,
                              float roll, float pitch, float yaw,
                              float intensity, double time)
  {
    PosePoint pp;
    pp.x         = x;
    pp.y         = y;
    pp.z         = z;
    pp.roll      = roll;
    pp.pitch     = pitch;
    pp.yaw       = yaw;
    pp.intensity = intensity;
    pp.time      = time;
    return pp;
  }

  // Named constructor from a gtsam::Pose3 plus extra fields.
  static PosePoint fromPose3(const Pose3 &pose, float intensity, double time)
  {
    PosePoint pp;
    pp.x         = pose.translation().x();
    pp.y         = pose.translation().y();
    pp.z         = pose.translation().z();
    pp.roll      = pose.rotation().roll();
    pp.pitch     = pose.rotation().pitch();
    pp.yaw       = pose.rotation().yaw();
    pp.intensity = intensity;
    pp.time      = time;
    return pp;
  }

  // Convert to GTSAM Pose3.
  Pose3 toGtsamPose3() const
  {
    return Pose3(Rot3::RzRyRx(double(roll), double(pitch), double(yaw)),
                 Point3(double(x), double(y), double(z)));
  }
};

// Register with PCL so it can be stored in a point cloud.
POINT_CLOUD_REGISTER_POINT_STRUCT(
    PosePoint,
    (float, x, x)
    (float, y, y)
    (float, z, z)
    (float, intensity, intensity)
    (float, roll, roll)
    (float, pitch, pitch)
    (float, yaw, yaw)
    (double, time, time)
)

// ----------------------------------------------------------------------------
// Free functions that do NOT reference any instance data
inline Pose3 trans2gtsamPose(const float transformIn[6])
{
    return Pose3(Rot3::RzRyRx(transformIn[0], transformIn[1], transformIn[2]),
                 Point3(transformIn[3], transformIn[4], transformIn[5]));
}

inline Eigen::Affine3f pclPointToAffine3f(const PosePoint &p)
{
    return pcl::getTransformation(p.x, p.y, p.z, p.roll, p.pitch, p.yaw);
}

// ----------------------------------------------------------------------------
// The refactored Example class
class Example
{
public:
    Example(ros::NodeHandle &nh)
    {
        // iSAM2 parameters
        parameters_.relinearizeThreshold = 0.1;
        parameters_.relinearizeSkip      = 1;
        iSam_ = std::make_unique<ISAM2>(parameters_);

        // Publisher
        pathPub_ = nh.advertise<nav_msgs::Path>("odom_path", 10, true);
    }

    // Called from main to push odometry messages into our queue
    void laserOdometryHandler(const nav_msgs::Odometry::ConstPtr &laserOdometry)
    {
        mBuf_.lock();
        odometryBuf_.push(laserOdometry);
        mBuf_.unlock();

        timeLaserInfoStamp_ = laserOdometry->header.stamp;
        timeLaserInfoCur_   = laserOdometry->header.stamp.toSec();
    }

    // Main processing loop: update, add factors, correct
    void runThread()
    {
        while (ros::ok())
        {
            updateInitGuess();
            saveKeyFrameAndFactor();
            correctPoses();
            std::this_thread::sleep_for(std::chrono::milliseconds(2));
        }
    }

    // Loop closure detection on a separate thread
    void loopClosureThread()
    {
        ros::Rate rate(1.0);
        while (ros::ok())
        {
            rate.sleep();
            performLoopClosure();
        }
    }

private:
    // 1) Update from odometry buffer
    void updateInitGuess()
    {
        if (!odometryBuf_.empty())
        {
            mBuf_.lock();
            auto odom = odometryBuf_.front();
            odometryBuf_.pop();
            while (!odometryBuf_.empty())
            {
                odometryBuf_.pop();  // drop to keep it real-time
                ROS_INFO("Dropped an extra odometry frame");
            }
            mBuf_.unlock();

            // Convert to transformTobeMapped_
            Eigen::Quaterniond q_odom(
                odom->pose.pose.orientation.w,
                odom->pose.pose.orientation.x,
                odom->pose.pose.orientation.y,
                odom->pose.pose.orientation.z);
            Eigen::Vector3d t_odom(
                odom->pose.pose.position.x,
                odom->pose.pose.position.y,
                odom->pose.pose.position.z);
            odomToTransform(q_odom, t_odom);
        }
    }

    // 2) Decide if a new keyframe is needed
    bool saveFrame()
    {
        static bool first = true;
        if (first)
        {
            // first keyframe
            // Use a named constructor for clarity:
            PosePoint initial = PosePoint::fromXYZRPY(
              transformTobeMapped_[3], transformTobeMapped_[4], transformTobeMapped_[5],
              transformTobeMapped_[0], transformTobeMapped_[1], transformTobeMapped_[2],
              0.0 /*intensity*/, 0.0 /*time*/);

            cloudKeyPoses6D_->push_back(initial);
            first = false;
            return true;
        }

        // Compare transform of last keyframe vs. current
        PosePoint &lastKf = cloudKeyPoses6D_->back();
        Eigen::Affine3f transStart = pclPointToAffine3f(lastKf);
        Eigen::Affine3f transFinal = pcl::getTransformation(
            transformTobeMapped_[3], transformTobeMapped_[4], transformTobeMapped_[5],
            transformTobeMapped_[0], transformTobeMapped_[1], transformTobeMapped_[2]);
        Eigen::Affine3f transBetween = transStart.inverse() * transFinal;

        float x, y, z, roll, pitch, yaw;
        pcl::getTranslationAndEulerAngles(transBetween, x, y, z, roll, pitch, yaw);

        float distThresh  = 1.0f; 
        float angleThresh = 0.2f; 

        if (fabs(roll)  < angleThresh &&
            fabs(pitch) < angleThresh &&
            fabs(yaw)   < angleThresh &&
            sqrtf(x*x + y*y + z*z) < distThresh)
        {
            return false;  // not enough motion to be a keyframe
        }
        return true;
    }

    // 3) Add odometry factor to the factor graph
    void addOdomFactor()
    {
        // If empty => first factor => Prior
        if (cloudKeyPoses3D_->points.empty())
        {
            auto priorNoise = noiseModel::Diagonal::Variances(
                (Vector(6) << 1e-2, 1e-2, M_PI*M_PI, 1e4, 1e4, 1e4).finished());

            tSAMgraph_.add(PriorFactor<Pose3>(
                0, trans2gtsamPose(transformTobeMapped_), priorNoise));
            initialEstimate_.insert(0, trans2gtsamPose(transformTobeMapped_));
        }
        else
        {
            // Odom factor between last pose and current
            auto odometryNoise = noiseModel::Diagonal::Variances(
                (Vector(6) << 1e-6, 1e-6, 1e-6, 1e-4, 1e-4, 1e-4).finished());

            // previous
            Eigen::Vector3d euler_last = lastQuat_.toRotationMatrix().eulerAngles(2,1,0);
            Pose3 poseFrom(Rot3::RzRyRx(euler_last[0], euler_last[1], euler_last[2]),
                           Point3(lastTrans_[0], lastTrans_[1], lastTrans_[2]));

            // current
            Eigen::Vector3d euler_curr = currQuat_.toRotationMatrix().eulerAngles(2,1,0);
            Pose3 poseTo(Rot3::RzRyRx(euler_curr[0], euler_curr[1], euler_curr[2]),
                         Point3(currTrans_[0], currTrans_[1], currTrans_[2]));

            tSAMgraph_.add(BetweenFactor<Pose3>(
                cloudKeyPoses3D_->size()-1,
                cloudKeyPoses3D_->size(),
                poseFrom.between(poseTo),
                odometryNoise));

            initialEstimate_.insert(cloudKeyPoses3D_->size(), poseTo);
        }

        // Store current odometry as "last"
        recordLastFrame();
    }

    // 4) Add any loop factors that were discovered
    void addLoopFactor()
    {
        if (loopIndexQueue_.empty()) return;

        for (size_t i = 0; i < loopIndexQueue_.size(); i++)
        {
            int idxFrom = loopIndexQueue_[i].first;
            int idxTo   = loopIndexQueue_[i].second;
            Pose3 poseBetween = loopPoseQueue_[i];
            auto noise = noiseModel::Diagonal::Sigmas(
                (Vector(6) << Vector3::Constant(1e-6), Vector3::Constant(0.03)).finished());

            tSAMgraph_.add(BetweenFactor<Pose3>(idxFrom, idxTo, poseBetween, noise));
        }
        loopIndexQueue_.clear();
        loopPoseQueue_.clear();
        loopIsClosed_ = true;
    }

    // 5) Combine the logic for saving a keyframe + building factor graph + optimizing
    void saveKeyFrameAndFactor()
    {
        if (!saveFrame())
            return;

        addOdomFactor();
        addLoopFactor();

        // iSAM2 optimization
        iSam_->update(tSAMgraph_, initialEstimate_);
        iSam_->update();

        // If loop closure was added, re-run a few times
        if (loopIsClosed_)
        {
            for(int i = 0; i < 5; i++)
                iSam_->update();
        }
        // Clear for next iteration
        tSAMgraph_.resize(0);
        initialEstimate_.clear();

        // Get the newest pose
        iSamCurrentEstimate_ = iSam_->calculateEstimate();
        Pose3 latestEstimate = iSamCurrentEstimate_.at<Pose3>(iSamCurrentEstimate_.size() - 1);

        // Convert to a PosePoint
        PosePoint thisPose6D = PosePoint::fromPose3(
            latestEstimate,
            static_cast<float>(cloudKeyPoses3D_->size()), // intensity as index
            timeLaserInfoCur_);

        // For the 3D key-poses (PCL uses e.g. pcl::PointXYZ)
        pcl::PointXYZ p3d;
        p3d.x         = thisPose6D.x;
        p3d.y         = thisPose6D.y;
        p3d.z         = thisPose6D.z;
        p3d.data[3]   = 1.0f;   // homogeneous coordinate
        p3d.intensity = thisPose6D.intensity;

        cloudKeyPoses3D_->push_back(p3d);
        cloudKeyPoses6D_->push_back(thisPose6D);

        // Publish path
        updatePath(thisPose6D);
        globalPath_.header.stamp    = timeLaserInfoStamp_;
        globalPath_.header.frame_id = "camera_init";
        pathPub_.publish(globalPath_);
    }

    // 6) Correct all older poses after loop closures
    void correctPoses()
    {
        if (cloudKeyPoses3D_->empty()) return;

        if (loopIsClosed_)
        {
            iSamCurrentEstimate_ = iSam_->calculateEstimate();
            int numPoses = iSamCurrentEstimate_.size();
            globalPath_.poses.clear();

            for (int i = 0; i < numPoses; i++)
            {
                Pose3 est = iSamCurrentEstimate_.at<Pose3>(i);

                cloudKeyPoses3D_->points[i].x = est.translation().x();
                cloudKeyPoses3D_->points[i].y = est.translation().y();
                cloudKeyPoses3D_->points[i].z = est.translation().z();

                PosePoint &pose6D = cloudKeyPoses6D_->points[i];
                pose6D.x    = est.translation().x();
                pose6D.y    = est.translation().y();
                pose6D.z    = est.translation().z();
                pose6D.roll  = est.rotation().roll();
                pose6D.pitch = est.rotation().pitch();
                pose6D.yaw   = est.rotation().yaw();

                updatePath(pose6D);
            }
            loopIsClosed_ = false;
        }
    }

    // 7) Detect loop closures. (Simple approach.)
    void performLoopClosure()
    {
        if (cloudKeyPoses3D_->points.empty()) return;

        // Make copies so we can do kdtree search safely
        *cloudKeyPoses3Dcopy_ = *cloudKeyPoses3D_;
        *cloudKeyPoses6Dcopy_ = *cloudKeyPoses6D_;

        int loopKeyCur, loopKeyPre;
        if (!detectLoopClosureDistance(&loopKeyCur, &loopKeyPre))
            return;

        Pose3 poseFrom = cloudKeyPoses6Dcopy_->points[loopKeyCur].toGtsamPose3();
        Pose3 poseTo   = cloudKeyPoses6Dcopy_->points[loopKeyPre].toGtsamPose3();

        loopIndexQueue_.push_back(std::make_pair(loopKeyCur, loopKeyPre));
        loopPoseQueue_.push_back(poseFrom.between(poseTo));
        loopIndexContainer_[loopKeyCur] = loopKeyPre;
    }

    // 7a) The distance-based check for loop closure
    bool detectLoopClosureDistance(int *latestID, int *closestID)
    {
        int loopKeyCur = cloudKeyPoses3Dcopy_->size() - 1;
        int loopKeyPre = -1;
        float searchRadius = 10.0f; // must be far in time for a loop

        if (loopIndexContainer_.find(loopKeyCur) != loopIndexContainer_.end())
            return false;

        // Kdtree search
        std::vector<int>   pointSearchInd;
        std::vector<float> pointSearchSqDis;
        kdtreeHistoryKeyPoses_->setInputCloud(cloudKeyPoses3Dcopy_);
        kdtreeHistoryKeyPoses_->radiusSearch(
            cloudKeyPoses3Dcopy_->back(),
            searchRadius,
            pointSearchInd,
            pointSearchSqDis,
            5);

        for (auto &id : pointSearchInd)
        {
            double poseTime = cloudKeyPoses6Dcopy_->points[id].time;
            if (fabs(poseTime - timeLaserInfoCur_) > searchRadius)
            {
                loopKeyPre = id;
                break;
            }
        }

        if (loopKeyPre == -1 || loopKeyCur == loopKeyPre) return false;

        *latestID  = loopKeyCur;
        *closestID = loopKeyPre;
        return true;
    }

    // Utilities
    void odomToTransform(const Eigen::Quaterniond &q, const Eigen::Vector3d &t)
    {
        // Convert to Euler
        Eigen::Vector3d euler = q.toRotationMatrix().eulerAngles(2,1,0);
        transformTobeMapped_[0] = euler[0];
        transformTobeMapped_[1] = euler[1];
        transformTobeMapped_[2] = euler[2];
        transformTobeMapped_[3] = t.x();
        transformTobeMapped_[4] = t.y();
        transformTobeMapped_[5] = t.z();

        // Save current as "curr"
        currQuat_  = q;
        currTrans_ = t;
    }

    void recordLastFrame()
    {
        lastQuat_  = currQuat_;
        lastTrans_ = currTrans_;
    }

    void updatePath(const PosePoint &pose_in)
    {
        geometry_msgs::PoseStamped pose_stamped;
        pose_stamped.header.stamp    = ros::Time().fromSec(pose_in.time);
        pose_stamped.header.frame_id = "camera_init";
        pose_stamped.pose.position.x = pose_in.x;
        pose_stamped.pose.position.y = pose_in.y;
        pose_stamped.pose.position.z = pose_in.z;
        tf::Quaternion q = tf::createQuaternionFromRPY(pose_in.roll, pose_in.pitch, pose_in.yaw);
        pose_stamped.pose.orientation.x = q.x();
        pose_stamped.pose.orientation.y = q.y();
        pose_stamped.pose.orientation.z = q.z();
        pose_stamped.pose.orientation.w = q.w();

        globalPath_.poses.push_back(pose_stamped);
    }

private:
    // Replaces the old “g” globals
    float transformTobeMapped_[6] = {0.f}; 
    ros::Time timeLaserInfoStamp_;
    double    timeLaserInfoCur_   = 0.0;
    bool      loopIsClosed_       = false;

    NonlinearFactorGraph tSAMgraph_;
    Values              initialEstimate_;
    std::unique_ptr<ISAM2> iSam_;    // was “ISAM2*”
    Values              iSamCurrentEstimate_;
    ISAM2Params         parameters_;

    nav_msgs::Path globalPath_;
    ros::Publisher pathPub_;

    // Keyframe storage
    // We store 3D positions (pcl::PointXYZ) plus the 6D PosePoint
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloudKeyPoses3D_     { new pcl::PointCloud<pcl::PointXYZ>() };
    pcl::PointCloud<PosePoint>::Ptr     cloudKeyPoses6D_     { new pcl::PointCloud<PosePoint>() };
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloudKeyPoses3Dcopy_ { new pcl::PointCloud<pcl::PointXYZ>() };
    pcl::PointCloud<PosePoint>::Ptr     cloudKeyPoses6Dcopy_ { new pcl::PointCloud<PosePoint>() };
    pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr kdtreeHistoryKeyPoses_{ new pcl::KdTreeFLANN<pcl::PointXYZ>() };

    // Loop closure
    std::map<int,int> loopIndexContainer_;
    std::vector<std::pair<int,int>> loopIndexQueue_;
    std::vector<Pose3> loopPoseQueue_;

    // Buffers and locks
    std::mutex mBuf_;
    std::queue<nav_msgs::Odometry::ConstPtr> odometryBuf_;

    // For storing last/current odometry states
    Eigen::Quaterniond lastQuat_ = Eigen::Quaterniond(1,0,0,0);
    Eigen::Vector3d    lastTrans_{0,0,0};
    Eigen::Quaterniond currQuat_ = Eigen::Quaterniond(1,0,0,0);
    Eigen::Vector3d    currTrans_{0,0,0};
};

// ----------------------------------------------------------------------------
// Minimal main
int main(int argc, char **argv)
{
    ros::init(argc, argv, "keyFrame");
    ros::NodeHandle nh;

    Example example(nh);

    // Threads
    std::thread loopThread(&Example::loopClosureThread, &example);
    std::thread runThread(&Example::runThread, &example);

    // Subscriber
    ros::Subscriber subOdom = nh.subscribe<nav_msgs::Odometry>(
        "/odometry", 100, &Example::laserOdometryHandler, &example);

    ros::spin();

    loopThread.join();
    runThread.join();
    return 0;
}