#include <iostream>
#include <random>

#include "pose_data_generator.hpp"
#include "trajectory_visualizer.hpp"

#include "optimizer.h"

#include "pose_graph_g2o_lie_algebra.cpp"

using namespace between_factor;

int main() {
    // 生成数据
    PoseDataGenerator generator;
    vector<Sophus::SE3d> gt_poses, odom_poses;
    vector<RelativeConstraint> relative_constraints;
    // 生成1km矩形轨迹，每米一帧
    // 生成回环约束
    vector<LoopClosureConstraint> loop_constraints;

    double side_length = 250.0;  // 250米边长，总长1km
    int num_frames = 1000;       // 1000帧
    
    generator.GenerateRectangularTrajectory(num_frames, side_length, 
                                           gt_poses, odom_poses, relative_constraints, loop_constraints);

    // 添加一些回环约束（连接轨迹的开始和结束部分）
    for(int i = 0; i < 1; ++i) {
            int idx1 = i * 10;
            int idx2 = num_frames - 1 - i * 10;
            if(idx2 > idx1 + 100) {  // 确保有足够的距离
                SE3d relative_pose = gt_poses[idx1].inverse() * gt_poses[idx2];
                SE3d relative_odometry_pose = odom_poses[idx1].inverse() * odom_poses[idx2];

                loop_constraints.emplace_back(idx1, idx2, relative_pose);

                cout << "relative_odometry_pose " << relative_odometry_pose.translation().transpose() << endl;
                cout << "relative_pose " << relative_pose.translation().transpose() << endl;
            }
    }
    cout << "Generated " << loop_constraints.size() << " loop closure constraints" << endl;

    between_factor::Optimizer optimizer;

    // 优化位姿图
    std::map<int, Vector6d> poses;
    std::vector<Constraint> constriants;
    for(std::uint32_t idx = 0; idx < relative_constraints.size(); ++idx) {
        RelativeConstraint relative_constraint = relative_constraints.at(idx);
        Sophus::SE3d last_odometry = odom_poses.at(relative_constraint.id1);
        Sophus::SE3d curr_odometry = odom_poses.at(relative_constraint.id2);
        poses[relative_constraint.id1] = last_odometry.log();
        poses[relative_constraint.id2] = curr_odometry.log();

        Constraint constriant;
        constriant.id_begin = relative_constraint.id1;
        constriant.id_end = relative_constraint.id2;       
        constriant.measurement = relative_constraint.relative_pose;
        constriant.information = relative_constraint.information;

        constriants.emplace_back(constriant);
    }

    for(std::uint32_t idx = 0; idx < loop_constraints.size(); ++idx) {
        LoopClosureConstraint loop_constraint = loop_constraints.at(idx);
        Constraint constriant;
        constriant.id_begin = loop_constraint.id1;
        constriant.id_end = loop_constraint.id2;       
        constriant.measurement = loop_constraint.relative_pose;
        constriant.information = loop_constraint.information;
        constriants.emplace_back(constriant);
    }
    
    // G2oSolve(constriants, &poses);
    optimizer.Solve(constriants, &poses);

    vector<SE3d> optimized_poses;  // 初始化为里程计位姿
    for(const auto& pose : poses) {
        optimized_poses.push_back(SE3d::exp(pose.second));
    }

    std::cout << " poses.size ========  : " << poses.size() << std::endl;
    std::cout << " optimized_poses.size ========  : " << optimized_poses.size() << std::endl;
    std::cout << " constriants.size ========  : " << constriants.size() << std::endl;

    // 可视化
    cout << "\nStarting visualization..." << endl;
    TrajectoryVisualizer visualizer;
    // visualizer.visualize(gt_poses, odom_poses, optimized_poses);
    visualizer.visualize(gt_poses, odom_poses, optimized_poses, loop_constraints);
    return 0;
}