// cartesian_controller_test.cpp
// ---------------------------------------------------------------------------
// Refactored test for CartesianController with trajectory timing & logging
//   * Left arm: circle in Y‑Z plane (normal‑X), r = 0.05m
//   * Right arm: square in Y‑Z plane (normal‑X), side = 0.05m
//   * Prints joint angles (q[DOF]) and EE positions (x,y,z) every step
// ---------------------------------------------------------------------------

#include <chrono>
#include <iostream>
#include <memory>
#include <vector>
#include <cmath>
#include <iomanip>

#include "utils/dynamics/robot.h"
#include "utils/solver/qp_solver.h"

// ---------- CONFIG ---------------------------------------------------------
constexpr int DOF = 14;                                    // robot DOF
using RobotT      = cmvr::dyn::Robot<DOF>;
using StateT      = cmvr::dyn::State<DOF>;
using ControllerT = cmvr::ctrl::CartesianController<DOF>;

static const char* kUrdfPath = "/home/xtkuang/projects/cmvr-es/config/robot_description/hc_description/dual_arm.urdf";
static const char* kBaseLink = "PELVIS_S";
static const char* kLeftEE   = "L_WRIST_R_S";
static const char* kRightEE  = "R_WRIST_R_S";

// ---------- TRAJECTORY GENERATORS -----------------------------------------
std::vector<Eigen::Matrix4d> generateCircle(const Eigen::Matrix4d& T_center,
                                            double radius, int n_points)
{
    std::vector<Eigen::Matrix4d> poses; poses.reserve(n_points);
    const double x  = T_center(0,3);
    const double y0 = T_center(1,3);
    const double z0 = T_center(2,3);
    const Eigen::Matrix3d R = T_center.block<3,3>(0,0);
    for (int i = 0; i < n_points; ++i) {
        double th = 2.0 * M_PI * i / n_points;
        double y = y0 + radius * std::cos(th);
        double z = z0 + radius * std::sin(th);
        Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
        T.block<3,3>(0,0) = R;
        T(0,3) = x; T(1,3) = y; T(2,3) = z;
        poses.push_back(T);
    }
    return poses;
}

std::vector<Eigen::Matrix4d> generateSquare(const Eigen::Matrix4d& T_center,
                                            double side, int n_points)
{
    std::vector<Eigen::Matrix4d> poses; poses.reserve(n_points);
    const double half = side / 2.0;
    const double x  = T_center(0,3);
    const double y0 = T_center(1,3);
    const double z0 = T_center(2,3);
    const Eigen::Matrix3d R = T_center.block<3,3>(0,0);

    for (int i = 0; i < n_points; ++i) {
        double u = static_cast<double>(i) / n_points;       // 0‑1
        double seg = u * 4.0;                               // 4 edges
        double y, z;
        if (seg < 1.0) {               // −Y → +Z
            y = -half;                 z = -half + seg*side;
        } else if (seg < 2.0) {        // +Z → +Y
            y = -half + (seg-1)*side;  z =  half;
        } else if (seg < 3.0) {        // +Y → −Z
            y =  half;                 z =  half - (seg-2)*side;
        } else {                       // −Z → −Y
            y =  half - (seg-3)*side;  z = -half;
        }
        Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
        T.block<3,3>(0,0) = R;
        T(0,3) = x; T(1,3) = y0 + y; T(2,3) = z0 + z;
        poses.push_back(T);
    }
    return poses;
}

std::vector<Eigen::Matrix4d> generateLine(const Eigen::Matrix4d& T_center, double length, int n_points)
{
    std::vector<Eigen::Matrix4d> poses; poses.reserve(n_points);
    double delta_x = length / (double)n_points;
    for (int i = 0; i < n_points; ++i) {
        const Eigen::Matrix3d R = T_center.block<3,3>(0,0);
        Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
        T.block<3,3>(0,0) = R;
        T(0,3) = T_center(0,3) + delta_x * (i+1);
        T(1,3) = T_center(1,3);
        T(2,3) = T_center(2,3);
        poses.push_back(T);
    }
    return poses;
}

// ---------- CSV HELPERS ----------------------------------------------------
// 写入表头
template<int DOF>
void writeCsvHeader(std::ofstream& csv)
{
    csv << std::fixed << std::setprecision(6);
    csv << "idx";
    for (int d = 0; d < DOF; ++d) csv << ", q" << d;
    csv << ", pLx, pLy, pLz, pRx, pRy, pRz\n";
}

// 写入一行数据
template<int DOF>
void writeCsvRow(std::ofstream& csv,
                 int idx,
                 const Eigen::Vector<double, DOF>& q,
                 const Eigen::Vector3d& pL,
                 const Eigen::Vector3d& pR)
{
    csv << idx;
    for (int d = 0; d < DOF; ++d)
        csv << ", " << q[d];
    csv << ',' << pL.x() << ',' << pL.y() << ',' << pL.z();
    csv << ',' << pR.x() << ',' << pR.y() << ',' << pR.z();
    csv << '\n';
}

// ---------- MAIN TEST ROUTINE ---------------------------------------------
void runTrajectoryTest(int n_points_circle = 200,
                       int n_points_square = 200,
                       double dt = 0.002)
{
    // 1. Load robot ---------------------------------------------------------
    auto rcfg  = cmvr::dyn::LoadRobotFromURDF(kUrdfPath, kBaseLink);
    auto robot = std::make_shared<RobotT>(rcfg);

    std::vector<std::string> link_names = {
            kBaseLink,
            "L_SHOULDER_P_S", "L_SHOULDER_R_S", "L_SHOULDER_Y_S", "L_ELBOW_R_S", "L_WRIST_P_S", "L_WRIST_Y_S", kLeftEE,
            "R_SHOULDER_P_S", "R_SHOULDER_R_S", "R_SHOULDER_Y_S", "R_ELBOW_R_S", "R_WRIST_P_S", "R_WRIST_Y_S", kRightEE
    };
    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"
    };
    auto state = robot->MakeState(link_names, joint_names);

    Eigen::Vector<double, DOF> q_init;
    q_init << -0.1818, -0.573093, -1.15019, -1.82123, 2.14858, 0.436445, 0.0384556,
    0.0734063, 1.16235, 1.72241, 1.81379, -2.70904, 0.0389824, 0.356423;
    state->SetQ(q_init);
    robot->ComputeForwardKinematics(state);

    // EE centers -----------------------------------------------------------
    const auto base_idx  = robot->GetLinkIdx(kBaseLink);
    const auto left_idx  = robot->GetLinkIdx(kLeftEE);
    const auto right_idx = robot->GetLinkIdx(kRightEE);
    const Eigen::Matrix4d T_left_center  = robot->GetTransformation(state, base_idx, left_idx);
    const Eigen::Matrix4d T_right_center = robot->GetTransformation(state, base_idx, right_idx);

    // Trajectories ---------------------------------------------------------
    const auto circle_traj = generateCircle (T_left_center , 0.1, n_points_circle);
    // const auto circle_traj = generateLine(T_left_center, 0.1, 200);
    // const auto square_traj = generateSquare(T_right_center, 0.1, n_points_square);
    const auto square_traj = generateLine(T_right_center, 0.15, n_points_square);
    std::cout << "--- square_traj (" << square_traj.size() << " waypoints) ---\n";
    std::cout << "--- end square_traj ---\n";
    const int N = std::max(circle_traj.size(), square_traj.size());

    // Controller -----------------------------------------------------------
    ControllerT ctrl(robot);

    // Print CSV header -----------------------------------------------------
    std::cout << std::fixed << std::setprecision(6);
    std::cout << "idx";
    for (int d = 0; d < DOF; ++d) std::cout << ", q" << d;
    std::cout << ", pLx, pLy, pLz, pRx, pRy, pRz\n";

    // Write CSV ------------------------------------------------------------
    std::string csv_path = "../../joint_positions.csv";
    std::ofstream csv(csv_path);
    if (!csv) throw std::runtime_error("cannot open " + csv_path);

    writeCsvHeader<DOF>(csv);   // ← 写表头

    // Timing --------------------------------------------------------------
    namespace chrono = std::chrono;
    const auto t_start = chrono::high_resolution_clock::now();

    for (int k = 0; k < N; ++k) {
        const Eigen::Matrix4d& T_left  = circle_traj [k % circle_traj.size()];
        const Eigen::Matrix4d& T_right = square_traj [k % square_traj.size()];

        std::vector<cmvr::ctrl::PoseTarget> targets = {
                {kLeftEE , T_left , 0.5, 1.0},
                {kRightEE, T_right, 0.5, 1.0}
        };

        Eigen::Vector<double, DOF> q_cmd;
        bool ok = ctrl.compute(state, kBaseLink, targets, dt,
                               ControllerT::Mode::Position, q_cmd, 60, 1e-4);
        if (!ok) std::cerr << "[WARN] IK failed at step " << k << '\n';

        state->SetQ(q_cmd);
        robot->ComputeForwardKinematics(state);

        // Current EE positions -------------------------------------------
        const Eigen::Matrix4d T_L_now = robot->GetTransformation(state, base_idx, left_idx);
        const Eigen::Matrix4d T_R_now = robot->GetTransformation(state, base_idx, right_idx);
        const Eigen::Vector3d pL = T_L_now.block<3,1>(0,3);
        const Eigen::Vector3d pR = T_R_now.block<3,1>(0,3);

        // CSV‑style print --------------------------------------------------
        std::cout << k;
        for (int d = 0; d < DOF; ++d) std::cout << ", " << q_cmd[d];
        std::cout << ", " << pL.transpose() << ", " << pR.transpose() << '\n';
        writeCsvRow<DOF>(csv, k, q_cmd, pL, pR);
    }

    const auto t_end = chrono::high_resolution_clock::now();
    const double total_ms = chrono::duration<double, std::milli>(t_end - t_start).count();
    csv.close();

    // Summary -------------------------------------------------------------
    std::cout << "---------- Trajectory Test Summary ----------\n";
    std::cout << "Total points : " << N << "\n";
    std::cout << "Total time   : " << total_ms << " ms\n";
    std::cout << "Average / pt : " << total_ms / N << " ms" << std::endl;
}

// ---------- MAIN ----------------------------------------------------------
int main() {
    try { runTrajectoryTest(); }
    catch (const std::exception &e) {
        std::cout << e.what();
    }
}