/**
 * @file arm_raise.cpp
 * @brief SAFE standalone test: raise the G1 RIGHT arm via the arm_sdk interface.
 *
 * Based on unitree_sdk2 example g1_arm7_sdk_dds_example.cpp, hardened for safety:
 *   - aborts if rt/lowstate is not received (never commands from bogus zeros)
 *   - ENTER-gated phases (engage -> raise -> lower)
 *   - Ctrl+C ramps the blend weight to 0 (gracefully hands the arm back)
 *
 * Moves ONLY the right arm; left arm + waist are held at their measured pose.
 *
 *   ./arm_raise [networkInterface]      (default: eth0)
 *
 * !!! THE ARM PHYSICALLY MOVES. Clear the space, stop any other arm controller,
 *     keep the E-stop in hand. !!!
 */

#include <array>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <csignal>
#include <cstring>
#include <iostream>
#include <thread>

#include <unitree/idl/hg/LowCmd_.hpp>
#include <unitree/idl/hg/LowState_.hpp>
#include <unitree/robot/channel/channel_publisher.hpp>
#include <unitree/robot/channel/channel_subscriber.hpp>

static const std::string kTopicArmSDK = "rt/arm_sdk";
static const std::string kTopicState  = "rt/lowstate";

enum JointIndex {
    kLeftHipPitch, kLeftHipRoll, kLeftHipYaw, kLeftKnee, kLeftAnkle, kLeftAnkleRoll,
    kRightHipPitch, kRightHipRoll, kRightHipYaw, kRightKnee, kRightAnkle, kRightAnkleRoll,
    kWaistYaw, kWaistRoll, kWaistPitch,
    kLeftShoulderPitch, kLeftShoulderRoll, kLeftShoulderYaw, kLeftElbow,
    kLeftWristRoll, kLeftWristPitch, kLeftWristYaw,
    kRightShoulderPitch, kRightShoulderRoll, kRightShoulderYaw, kRightElbow,
    kRightWristRoll, kRightWristPitch, kRightWristYaw,
    kNotUsedJoint, kNotUsedJoint1, kNotUsedJoint2, kNotUsedJoint3, kNotUsedJoint4, kNotUsedJoint5
};

static std::atomic<bool> g_stop{false};
static void onSigint(int) { g_stop = true; }

int main(int argc, char **argv)
{
    const std::string iface = argc > 1 ? argv[1] : "eth0";
    std::signal(SIGINT, onSigint);

    unitree::robot::ChannelFactory::Instance()->Init(0, iface);

    unitree::robot::ChannelPublisherPtr<unitree_hg::msg::dds_::LowCmd_> pub;
    pub.reset(new unitree::robot::ChannelPublisher<unitree_hg::msg::dds_::LowCmd_>(kTopicArmSDK));
    pub->InitChannel();

    static unitree_hg::msg::dds_::LowState_ state;
    std::atomic<bool> got{false};
    unitree::robot::ChannelSubscriberPtr<unitree_hg::msg::dds_::LowState_> sub;
    sub.reset(new unitree::robot::ChannelSubscriber<unitree_hg::msg::dds_::LowState_>(kTopicState));
    sub->InitChannel([&](const void *m) {
        std::memcpy(&state, m, sizeof(state));
        got = true;
    }, 1);

    // SAFETY 1: require fresh state before touching the arm.
    std::cout << "Waiting for " << kTopicState << " on '" << iface << "' ..." << std::endl;
    for (int i = 0; i < 300 && !got; i++) std::this_thread::sleep_for(std::chrono::milliseconds(10));
    if (!got) {
        std::cerr << "ERROR: no " << kTopicState << " received. Is the robot controller up, and is '"
                  << iface << "' the interface to the robot? Aborting (no arm command sent)." << std::endl;
        return 1;
    }

    const std::array<JointIndex, 17> arm = {
        kLeftShoulderPitch, kLeftShoulderRoll, kLeftShoulderYaw, kLeftElbow,
        kLeftWristRoll, kLeftWristPitch, kLeftWristYaw,
        kRightShoulderPitch, kRightShoulderRoll, kRightShoulderYaw, kRightElbow,
        kRightWristRoll, kRightWristPitch, kRightWristYaw,
        kWaistYaw, kWaistRoll, kWaistPitch};
    const int RIGHT0 = 7; // right arm occupies arm[7..13]

    const float kp = 60.f, kd = 1.5f, dt = 0.02f;
    const float max_delta = 0.5f * dt; // 0.5 rad/s velocity limit
    const auto  sleep_t = std::chrono::milliseconds(20);

    std::array<float, 17> start{};
    for (int i = 0; i < 17; i++) start[i] = state.motor_state().at(arm[i]).q();

    // Right-arm raise target [ShPitch, ShRoll, ShYaw, Elbow, WrRoll, WrPitch, WrYaw].
    // Conservative greeting raise (arm out + forearm up). Tune here.
    const std::array<float, 7> rraise = {0.0f, -1.0f, 0.0f, 1.4f, 0.0f, 0.0f, 0.0f};
    std::array<float, 17> target = start;
    for (int j = 0; j < 7; j++) target[RIGHT0 + j] = rraise[j];

    std::array<float, 17> des = start;
    float weight = 0.f;
    unitree_hg::msg::dds_::LowCmd_ msg;

    auto publish = [&](float w) {
        msg.motor_cmd().at(kNotUsedJoint).q(w);
        for (int i = 0; i < 17; i++) {
            auto &mc = msg.motor_cmd().at(arm[i]);
            mc.q(des[i]); mc.dq(0.f); mc.kp(kp); mc.kd(kd); mc.tau(0.f);
        }
        pub->Write(msg);
    };

    auto release = [&]() {
        std::cout << "Releasing arm (weight -> 0) ..." << std::endl;
        while (weight > 0.f) {
            weight = std::max(0.f, weight - 0.015f);
            publish(weight);
            std::this_thread::sleep_for(sleep_t);
        }
        publish(0.f);
    };

    // Move 'des' toward 'goal' (velocity limited). Returns false if interrupted.
    auto move_to = [&](const std::array<float, 17> &goal, int hold_steps) -> bool {
        for (int step = 0; step < 600; step++) {
            if (g_stop) return false;
            float maxerr = 0.f;
            for (int i = 0; i < 17; i++) {
                float d = std::clamp(goal[i] - des[i], -max_delta, max_delta);
                des[i] += d;
                maxerr = std::max(maxerr, std::abs(goal[i] - des[i]));
            }
            publish(1.f);
            std::this_thread::sleep_for(sleep_t);
            if (maxerr < 1e-3f) { // arrived -> hold
                for (int h = 0; h < hold_steps && !g_stop; h++) {
                    publish(1.f);
                    std::this_thread::sleep_for(sleep_t);
                }
                return !g_stop;
            }
        }
        return !g_stop;
    };

    std::cout << "\n*** G1 RIGHT-ARM RAISE TEST ***\n"
              << "The right arm WILL move. Ensure clear space + E-stop ready.\n"
              << "Press ENTER to engage the arm (no motion yet) ... " << std::flush;
    std::cin.get();
    if (g_stop) { release(); return 0; }

    // Phase 1: engage — ramp weight to 1 while commanding the measured pose (no motion).
    while (weight < 1.f && !g_stop) {
        weight = std::min(1.f, weight + 0.015f);
        publish(weight);
        std::this_thread::sleep_for(sleep_t);
    }
    if (g_stop) { release(); return 0; }

    std::cout << "Engaged. Press ENTER to RAISE the right arm ... " << std::flush;
    std::cin.get();
    if (g_stop || !move_to(target, 100)) { release(); return 0; }

    std::cout << "Raised. Press ENTER to LOWER the arm ... " << std::flush;
    std::cin.get();
    move_to(start, 25);

    release();
    std::cout << "Done." << std::endl;
    return 0;
}