#include <iostream>
#include <fstream>
#include <memory>
#include <vector>
#include <string>
#include <iomanip>

#include "utils/dynamics/robot.h"

constexpr int DOF = 14;
using RobotT = cmvr::dyn::Robot<DOF>;
using StateT = cmvr::dyn::State<DOF>;

static const char* kUrdfPath = "/home/xtkuang/projects/cmvr-es/config/robot_description/hc_description/dual_arm.urdf";

int main(int argc, char** argv) {
    if (argc < 10) {
        std::cerr << "用法: solve_fk <base_line> <target_link> <q1> <q2> <q3> <q4> <q5> <q6> <q7> [-o <filename>]\n";
        return 1;
    }

    bool write_file = false;
    std::string output_file;

    if (argc == 12 && std::string(argv[10]) == "-o") {
        write_file = true;
        output_file = argv[11];
    }

    std::string arm = argv[1];
    std::vector<std::string> link_names = {
        "PELVIS_S",
        "L_SHOULDER_P_S", "L_SHOULDER_R_S", "L_SHOULDER_Y_S", "L_ELBOW_R_S", "L_WRIST_P_S", "L_WRIST_Y_S", "L_WRIST_R_S",
        "R_SHOULDER_P_S", "R_SHOULDER_R_S", "R_SHOULDER_Y_S", "R_ELBOW_R_S", "R_WRIST_P_S", "R_WRIST_Y_S", "R_WRIST_R_S", "R_FINGER_TIP"
    };
    std::vector<std::string> joint_names = {
        "L_SHOULDER_P", "L_SHOULDER_R", "L_SHOULDER_Y", "L_ELBOW_R", "L_WRIST_P", "L_WRIST_Y", "L_WRIST_R",
        "R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y", "R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"
    };

    std::string base_link = argv[1];
    std::string target_link = argv[2];
    bool find_base_link = false;
    bool find_target_link = false;
    for (const auto& link_name : link_names) {
        if (base_link == link_name) {
            find_base_link = true;
        }
        if (target_link == link_name) {
            find_target_link = true;
        }
    }
    if (!find_base_link) {
        std::cerr << "base link 无效" << std::endl;
        return -1;
    }
    if (!find_target_link) {
        std::cerr << "target link 无效" << std::endl;
        return -1;
    }

    Eigen::Vector<double, DOF> q;
    q.setZero();
    for (int i = 0; i < 7; ++i) {
        if (target_link[0] == 'L') {
            q[i] = std::stod(argv[i+3]);
        }
        else if (target_link[0] == 'R') {
            q[i+7] = std::stod(argv[i+3]);
        }
    }

    auto rcfg = cmvr::dyn::LoadRobotFromURDF(kUrdfPath, "PELVIS_S");
    auto robot = std::make_shared<RobotT>(rcfg);
    auto state = robot->MakeState(link_names, joint_names);
    state->SetQ(q);
    robot->ComputeForwardKinematics(state);

    auto base_idx = robot->GetLinkIdx(base_link);
    auto ee_idx   = robot->GetLinkIdx(target_link);
    Eigen::Matrix4d T = robot->GetTransformation(state, base_idx, ee_idx);

    std::cout << std::fixed << std::setprecision(10);
    std::cout << "变换矩阵T:\n" << T << "\n";

    if (write_file) {
        std::ofstream fout(output_file);
        if (!fout) {
            std::cerr << "无法创建输出文件 " << output_file << "\n";
            return 1;
        }
        fout << std::fixed << std::setprecision(6);
        fout << T(0,0) << "," << T(0,1) << "," << T(0,2) << "," << T(0,3) << ","
             << T(1,0) << "," << T(1,1) << "," << T(1,2) << "," << T(1,3) << ","
             << T(2,0) << "," << T(2,1) << "," << T(2,2) << "," << T(2,3) << "\n";
        fout.close();
        std::cout << "结果已保存到 " << output_file << "\n";
    }

    return 0;
}