#include "realman_hw_test_common.hpp"

#include <algorithm>
#include <array>
#include <cmath>
#include <sstream>

namespace {

constexpr float kForceThresholdN = 3.0F;
constexpr float kForceReleaseThresholdN = 1.2F;
constexpr float kDangerForceLimitN = 60.0F;
constexpr float kRetreatStepMeters = 0.008F;
constexpr float kMaxTotalRetreatMeters = 0.032F;
constexpr int kMaxRetreatSteps = 4;
constexpr int kMotionSpeedPercent = 5;
constexpr int kWaitForForceSeconds = 12;
constexpr int kPollIntervalMs = 100;
constexpr int kWorkFrame = 0;
constexpr int kNoForceSkipReturnCode = 77;

// If the robot moves toward the push instead of retreating, change this to -1.0F.
constexpr float kRetreatDirectionSign = 1.0F;

struct Vec3 {
    float x = 0.0F;
    float y = 0.0F;
    float z = 0.0F;
};

float norm(const Vec3 &value) {
    return std::sqrt(value.x * value.x + value.y * value.y + value.z * value.z);
}

Vec3 normalize(const Vec3 &value) {
    const float length = norm(value);
    if (length <= 1e-6F) {
        return {};
    }
    return {value.x / length, value.y / length, value.z / length};
}

std::string formatVec3(const Vec3 &value, const char *unit) {
    std::ostringstream out;
    out << std::fixed << std::setprecision(4)
        << "(" << value.x << ", " << value.y << ", " << value.z << ") " << unit;
    return out.str();
}

bool readWorkFrameForce(rm_robot_handle *handle, Vec3 *force, realman_hw_test::TestSuite *suite) {
    rm_force_data_t data = {};
    const int result = rm_get_force_data(handle, &data);
    const bool ok = result == 0;
    if (suite != nullptr) {
        suite->expect(ok, "rm_get_force_data", "error=" + std::to_string(result));
    }
    if (!ok) {
        return false;
    }

    *force = {
        data.work_zero_force_data[0],
        data.work_zero_force_data[1],
        data.work_zero_force_data[2],
    };
    return true;
}

bool waitForExternalForce(rm_robot_handle *handle,
                          Vec3 *detected_force,
                          realman_hw_test::TestSuite *suite) {
    const int iterations = (kWaitForForceSeconds * 1000) / kPollIntervalMs;
    for (int i = 0; i < iterations; ++i) {
        Vec3 force = {};
        if (!readWorkFrameForce(handle, &force, nullptr)) {
            suite->expect(false, "wait_force_read", "rm_get_force_data failed");
            return false;
        }

        const float force_norm = norm(force);
        std::cout << "Waiting force, work-frame force=" << formatVec3(force, "N")
                  << ", norm=" << std::fixed << std::setprecision(3) << force_norm << " N\n";

        if (force_norm > kDangerForceLimitN) {
            suite->expect(false, "force_under_danger_limit", "force_norm_N=" + std::to_string(force_norm));
            return false;
        }
        if (force_norm >= kForceThresholdN) {
            *detected_force = force;
            suite->expect(true, "force_trigger_detected", "force_norm_N=" + std::to_string(force_norm));
            return true;
        }

        realman_hw_test::sleepMs(kPollIntervalMs);
    }

    suite->info("force_trigger", "not detected within " + std::to_string(kWaitForForceSeconds) + " seconds");
    return false;
}

rm_pose_t makeOffsetPose(const Vec3 &direction, float step_meters) {
    rm_pose_t offset = {};
    offset.position.x = kRetreatDirectionSign * direction.x * step_meters;
    offset.position.y = kRetreatDirectionSign * direction.y * step_meters;
    offset.position.z = kRetreatDirectionSign * direction.z * step_meters;
    offset.euler.rx = 0.0F;
    offset.euler.ry = 0.0F;
    offset.euler.rz = 0.0F;
    offset.quaternion.w = 0.0F;
    offset.quaternion.x = 0.0F;
    offset.quaternion.y = 0.0F;
    offset.quaternion.z = 0.0F;
    return offset;
}

std::string offsetDetail(const rm_pose_t &offset) {
    std::ostringstream out;
    out << std::fixed << std::setprecision(4)
        << "dx=" << offset.position.x
        << ", dy=" << offset.position.y
        << ", dz=" << offset.position.z << " m";
    return out.str();
}

} // namespace

int main() {
    using namespace realman_hw_test;

    TestSuite suite("realman_hw_force_retreat");
    RobotSession session(&suite);
    if (!session.connect()) {
        return suite.finish();
    }

    rm_robot_info_t info = {};
    if (!readRobotInfo(session.handle(), &info, &suite)) {
        return suite.finish();
    }
    suite.expect(info.force_type == RM_MODEL_RM_SF_E || info.force_type == RM_MODEL_RM_ISF_E,
                 "six_axis_force_supported",
                 std::string("actual=") + forceTypeName(info.force_type));

    suite.expect(kForceThresholdN > kForceReleaseThresholdN, "force_threshold_gt_release");
    suite.expect(kRetreatStepMeters > 0.0F && kRetreatStepMeters <= 0.010F,
                 "retreat_step_safe",
                 "step_m=" + std::to_string(kRetreatStepMeters));
    suite.expect(kMaxTotalRetreatMeters <= 0.040F,
                 "max_total_retreat_safe",
                 "max_total_m=" + std::to_string(kMaxTotalRetreatMeters));
    suite.expect(kMotionSpeedPercent >= 1 && kMotionSpeedPercent <= 10,
                 "retreat_speed_safe",
                 "speed_percent=" + std::to_string(kMotionSpeedPercent));
    suite.expect(kRetreatDirectionSign == 1.0F || kRetreatDirectionSign == -1.0F,
                 "retreat_direction_sign_valid",
                 "sign=" + std::to_string(kRetreatDirectionSign));

    Vec3 initial_force = {};
    std::cout << "\nManual force retreat test.\n"
              << "Apply a gentle force to the end effector within "
              << kWaitForForceSeconds << " seconds. Threshold: "
              << kForceThresholdN << " N. Keep E-stop reachable.\n";
    const bool triggered = waitForExternalForce(session.handle(), &initial_force, &suite);
    if (!triggered) {
        const int finish_code = suite.finish();
        return finish_code == 0 ? kNoForceSkipReturnCode : finish_code;
    }

    const Vec3 initial_direction = normalize(initial_force);
    suite.info("initial_force_work_frame", formatVec3(initial_force, "N"));
    suite.info("retreat_direction_work_frame", formatVec3(initial_direction, "unit"));

    float total_retreat = 0.0F;
    bool movement_ok = true;
    Vec3 latest_force = initial_force;

    for (int step = 0; step < kMaxRetreatSteps; ++step) {
        const float force_norm = norm(latest_force);
        if (force_norm <= kForceReleaseThresholdN) {
            suite.info("force_released", "force_norm_N=" + std::to_string(force_norm));
            break;
        }
        if (force_norm > kDangerForceLimitN) {
            suite.expect(false, "force_under_danger_limit", "force_norm_N=" + std::to_string(force_norm));
            movement_ok = false;
            break;
        }
        if (total_retreat + kRetreatStepMeters > kMaxTotalRetreatMeters + 1e-6F) {
            suite.info("max_retreat_reached", "total_m=" + std::to_string(total_retreat));
            break;
        }

        const Vec3 retreat_direction = normalize(latest_force);
        const rm_pose_t offset = makeOffsetPose(retreat_direction, kRetreatStepMeters);
        suite.info("retreat_step_" + std::to_string(step + 1), offsetDetail(offset));

        const int move_result = rm_movel_offset(session.handle(), offset, kMotionSpeedPercent, 0, 0, kWorkFrame, 1);
        const bool move_ok = move_result == 0;
        suite.expect(move_ok, "rm_movel_offset_step_" + std::to_string(step + 1),
                     "error=" + std::to_string(move_result));
        if (!move_ok) {
            rm_set_arm_stop(session.handle());
            movement_ok = false;
            break;
        }

        total_retreat += kRetreatStepMeters;
        realman_hw_test::sleepMs(250);

        if (!readWorkFrameForce(session.handle(), &latest_force, &suite)) {
            movement_ok = false;
            break;
        }
        suite.info("force_after_step_" + std::to_string(step + 1), formatVec3(latest_force, "N"));
    }

    suite.info("total_retreat_m", std::to_string(total_retreat));
    suite.expect(total_retreat <= kMaxTotalRetreatMeters + 1e-6F,
                 "total_retreat_limit",
                 "total_m=" + std::to_string(total_retreat));
    suite.expect(movement_ok, "force_retreat_sequence_ok");

    return suite.finish();
}
