#include <stdexcept>
#include "arm_define.h"
#include "trapezoidal_trajectory_kinematics.hpp"
#include <cstring>
#include <cmath>

extern double truncate_to_4_decimal(double num);

TrapezoidalTrajectoryKinematics::TrapezoidalTrajectoryKinematics(RM_Service *rmService)
    : max_velocity_(0.4), max_acceleration_(0.1), rmService_(rmService)
{
    max_accelerate_time_ = 4;
    max_accelerate_angle_ = max_velocity_ * max_accelerate_time_;
    min_angle_for_pose_stepping_max_speed_ = ARM_FOLLOWING_CHECKING_STEP * max_velocity_ * 2;

    rm_robot_arm_model_e Mode = RM_MODEL_RM_65_E;
    rm_force_type_e Type = RM_MODEL_RM_B_E;
    rm_algo_init_sys_data(Mode, Type);
}

TrapezoidalTrajectoryKinematics::TrapezoidalTrajectoryKinematics()
    : max_velocity_(0.4), max_acceleration_(0.1), rmService_(nullptr)
{
    max_accelerate_time_ = 4;
    max_accelerate_angle_ = max_velocity_ * max_accelerate_time_;
    min_angle_for_pose_stepping_max_speed_ = ARM_FOLLOWING_CHECKING_STEP * max_velocity_ * 2;

    rmService_ = new RM_Service();
    int result = rmService_->rm_init(RM_TRIPLE_MODE_E);
    if (result != 0) {
        throw std::runtime_error("Failed to init rm.");
        return;
    }
    rm_robot_arm_model_e Mode = RM_MODEL_RM_65_E;
    rm_force_type_e Type = RM_MODEL_RM_B_E;
    rm_algo_init_sys_data(Mode, Type);
}

TrapezoidalTrajectoryKinematics::~TrapezoidalTrajectoryKinematics() {
    delete rmService_;
}

void TrapezoidalTrajectoryKinematics::getAngleForCheckingInner(TrapezoidalTrajectory *traj, int timeStep, float *aimAngles)
{
    if (timeStep >= traj->accelerate_time_step * 2 + traj->uniform_time_step 
            && traj->uniform_time_step != STEPPING_UNIFORM_TIME_STAMP) { // 分支预判优化
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] = traj->endAngles[i];
        }
        return;
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        if (timeStep < traj->accelerate_time_step) {
            aimAngles[i] = traj->startAngles[i] + 0.5 * traj->now_acceleration[i] * (timeStep + 1) * timeStep;
        } else if (traj->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP){
            aimAngles[i] = traj->startAngles[i] + 0.5 * traj->now_acceleration[i] * (traj->accelerate_time_step + 1) * (traj->accelerate_time_step);
            aimAngles[i] += traj->aim_velocity[i] * (timeStep - traj->accelerate_time_step);
        }else if (timeStep < traj->accelerate_time_step + traj->uniform_time_step) {
            aimAngles[i] = traj->endAngles[i] - 0.5 * traj->now_acceleration[i] * (traj->accelerate_time_step - 1) * (traj->accelerate_time_step);
            aimAngles[i] -= traj->aim_velocity[i] * (traj->uniform_time_step - (timeStep - traj->accelerate_time_step));
        } else if (timeStep <= 2 * traj->accelerate_time_step + traj->uniform_time_step) {
            int decelerate_time_step = 2 * traj->accelerate_time_step + traj->uniform_time_step - timeStep;
            aimAngles[i] = traj->endAngles[i] - 0.5 * traj->now_acceleration[i] * (decelerate_time_step - 1) * decelerate_time_step;
        }
    }
}

void TrapezoidalTrajectoryKinematics::getAngleForCheckingOffset(TrapezoidalTrajectory *traj, int timeStepOffset, float *aimAngles)
{
    int timeStep = traj->current_time_step + timeStepOffset;
    getAngleForCheckingInner(traj, timeStep, aimAngles);
}

int TrapezoidalTrajectoryKinematics::updateTrajectory(TrapezoidalTrajectory *trajectory, float *aimAngles)
{
    trajectory->current_time_step++;
    if (trajectory->goalType == GOAL_TYPE_POSE_STEPPING && trajectory->current_time_step >= trajectory->accelerate_time_step + trajectory->uniform_time_step) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            int decelerate_time_step = trajectory->accelerate_time_step;
            aimAngles[i] = trajectory->endAngles[i] - 0.5 * trajectory->now_acceleration[i] * (decelerate_time_step + 1) * decelerate_time_step;
        }
        return TRAJECTORY_COMPLETE;
    } else if (trajectory->current_time_step >= trajectory->accelerate_time_step * 2 + trajectory->uniform_time_step
            && trajectory->uniform_time_step != STEPPING_UNIFORM_TIME_STAMP) {
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            aimAngles[i] = trajectory->endAngles[i];
        }
        return TRAJECTORY_COMPLETE;
    }
    getAngleForCheckingInner(trajectory, trajectory->current_time_step, aimAngles);
    return TRAJECTORY_ONGOING;
}

int TrapezoidalTrajectoryKinematics::computeMovingTrajectory(TrapezoidalTrajectory *newTrajectory, const float *start, const float *end, int armId)
{
    newTrajectory->goalType = GOAL_TYPE_MOVING;
    newTrajectory->current_time_step = 0;
    float max_angle = 0.0;
    float angle_diff[USED_ARM_DOF];
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        newTrajectory->endAngles[i] = end[i];
        angle_diff[i] = end[i] - start[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
        }
    }
    if (max_angle < 0.01) {
        newTrajectory->accelerate_time_step = 0;
        newTrajectory->uniform_time_step = 0;
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            newTrajectory->aim_velocity[i] = 0;
            newTrajectory->now_acceleration[i] = 0;
            newTrajectory->endAngles[i] = start[i];
        }
        return OK;
    }
    float aim_velocity = max_velocity_;
    if (max_angle < max_accelerate_angle_) {
        newTrajectory->accelerate_time_step = static_cast<int>(sqrt(max_angle / max_acceleration_)) + 1;
        aim_velocity = max_angle / newTrajectory->accelerate_time_step;
        newTrajectory->uniform_time_step = 0;
    } else {
        newTrajectory->accelerate_time_step = max_accelerate_time_;
        newTrajectory->uniform_time_step = static_cast<int>((max_angle - max_accelerate_angle_) / max_velocity_);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->aim_velocity[i] = aim_velocity * angle_diff[i] / max_angle;
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / newTrajectory->accelerate_time_step;
    }
    return OK;
}

int TrapezoidalTrajectoryKinematics::computeSteppingTrajectory(TrapezoidalTrajectory *newTrajectory, const float *start, const float *speed, int armId)
{
    newTrajectory->goalType = GOAL_TYPE_STEPPING;
    newTrajectory->current_time_step = 0;
    float maxSpeed = 0.0;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        if (fabs(speed[i]) > max_velocity_) {
            newTrajectory->aim_velocity[i] = (speed[i] > 0 ? 1 : -1) * max_velocity_;
        } else {
            newTrajectory->aim_velocity[i] = speed[i];
        }
        if (fabs(newTrajectory->aim_velocity[i]) > maxSpeed) {
            maxSpeed = fabs(newTrajectory->aim_velocity[i]);
        }
    }
    newTrajectory->uniform_time_step = -1;
    newTrajectory->accelerate_time_step = static_cast<int>(maxSpeed / max_acceleration_) + 1;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / newTrajectory->accelerate_time_step;
        newTrajectory->endAngles[i] = 0;
    }
    return OK;
}

int TrapezoidalTrajectoryKinematics::computePoseSteppingTrajectory(TrapezoidalTrajectory *newTrajectory, const float *start, const float *pose, int armId)
{
    newTrajectory->goalType = GOAL_TYPE_POSE_STEPPING;
    newTrajectory->current_time_step = 0;
    float end[ARM_DOF] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    for (int i = 0; i < USED_ARM_DOF; i++) {
        end[i] = start[i];
    }
    int result = PoseSteppingFromNowAnglesAndDirectPose(end, pose);
    if (result != 0){
        return ARM_AIM_CANNOT_REACH; // Inverse kinematics failed
    }

    float max_angle = 0.0;
    float angle_diff[USED_ARM_DOF];
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->startAngles[i] = start[i];
        newTrajectory->endAngles[i] = end[i];
        angle_diff[i] = end[i] - start[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
        }
    }
    if (max_angle < 0.01) {
        newTrajectory->accelerate_time_step = 0;
        newTrajectory->uniform_time_step = 0;
        for (int i = 0; i < USED_ARM_DOF; ++i) {
            newTrajectory->aim_velocity[i] = 0;
            newTrajectory->now_acceleration[i] = 0;
            newTrajectory->endAngles[i] = start[i];
        }
        return OK;
    }
    float aim_velocity = max_velocity_;
    newTrajectory->accelerate_time_step = max_accelerate_time_;
    if (max_angle < min_angle_for_pose_stepping_max_speed_) {
        aim_velocity = max_angle / (ARM_FOLLOWING_CHECKING_STEP * 2);
        newTrajectory->uniform_time_step = ARM_FOLLOWING_CHECKING_STEP * 2 - max_accelerate_time_;
    } else {
        newTrajectory->uniform_time_step = static_cast<int>((max_angle - (max_accelerate_angle_ - 1) * 0.5) / max_velocity_);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->aim_velocity[i] = aim_velocity * angle_diff[i] / max_angle;
        newTrajectory->now_acceleration[i] = newTrajectory->aim_velocity[i] / max_accelerate_time_;
    }
    return OK;
}

void TrapezoidalTrajectoryKinematics::stopTrajectory(TrapezoidalTrajectory *traj)
{
    int formerPoseSteppingUniformTimeStep = traj->uniform_time_step;
    if (traj->current_time_step < traj->accelerate_time_step) {
        traj->accelerate_time_step = traj->current_time_step;
        traj->uniform_time_step = 0;
    } else if (traj->uniform_time_step == STEPPING_UNIFORM_TIME_STAMP ||
        traj->current_time_step < traj->accelerate_time_step + traj->uniform_time_step) {
        traj->uniform_time_step = traj->current_time_step - traj->accelerate_time_step;
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        if (traj->goalType == GOAL_TYPE_POSE_STEPPING) {
            traj->now_acceleration[i] = traj->aim_velocity[i] / traj->accelerate_time_step;
            traj->endAngles[i] = traj->endAngles[i]
                - traj->aim_velocity[i] * (formerPoseSteppingUniformTimeStep - traj->uniform_time_step);
        } else if (traj->current_time_step <= traj->accelerate_time_step + traj->uniform_time_step) {
            traj->endAngles[i] = traj->startAngles[i] + traj->aim_velocity[i] * (traj->uniform_time_step + traj->accelerate_time_step);
        }
    }
    if (traj->goalType == GOAL_TYPE_POSE_STEPPING) {
        traj->goalType = GOAL_TYPE_MOVING;
    }
}

int TrapezoidalTrajectoryKinematics::PoseSteppingFromNowAnglesAndDirectPose(float *nowAngles, const float *direct)
{
    rm_pose_t pose = rmService_->rm_algo_forward_kinematics(NULL, nowAngles);

    rm_inverse_kinematics_params_t inverse_params;
    memcpy(inverse_params.q_in, nowAngles, ARM_DOF * sizeof(float));
    inverse_params.q_pose.position.x = direct[POSE_POSITION_X] + truncate_to_4_decimal(pose.position.x);
    inverse_params.q_pose.position.y = direct[POSE_POSITION_Y] + truncate_to_4_decimal(pose.position.y);
    inverse_params.q_pose.position.z = direct[POSE_POSITION_Z] + truncate_to_4_decimal(pose.position.z);
    inverse_params.q_pose.quaternion.x = direct[POSE_QUATERNION_X] + truncate_to_4_decimal(pose.quaternion.x);
    inverse_params.q_pose.quaternion.y = direct[POSE_QUATERNION_Y] + truncate_to_4_decimal(pose.quaternion.y);
    inverse_params.q_pose.quaternion.z = direct[POSE_QUATERNION_Z] + truncate_to_4_decimal(pose.quaternion.z);
    inverse_params.q_pose.quaternion.w = direct[POSE_QUATERNION_W] + truncate_to_4_decimal(pose.quaternion.w);
    inverse_params.flag = 0;

    int result = rmService_->rm_algo_inverse_kinematics(NULL, inverse_params, nowAngles);
    if (result != 0){
        return ARM_AIM_CANNOT_REACH; // Inverse kinematics failed
    }
    return 0;
}

int TrapezoidalTrajectoryKinematics::continuePoseSteppingTrajectory(TrapezoidalTrajectory *newTrajectory, const float *nowAngles, const float *direct)
{
    if (!newTrajectory || newTrajectory->goalType != GOAL_TYPE_POSE_STEPPING) {
        return ARM_GOAL_CANCELLED;
    }
    float newJoints[ARM_DOF] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    for (int i = 0; i < USED_ARM_DOF; i++) {
        newJoints[i] = newTrajectory->endAngles[i];
    }
    int result = PoseSteppingFromNowAnglesAndDirectPose(newJoints, direct);
    if (result != 0){
        return ARM_AIM_CANNOT_REACH; // Inverse kinematics failed
    }

    float angle_diff[USED_ARM_DOF];
    float max_angle = 0.0;
    float max_velocity = 0.0;
    int fastestIndex = 0;
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        newTrajectory->endAngles[i] = newJoints[i];
        angle_diff[i] = newTrajectory->endAngles[i] - nowAngles[i];
        float abs_diff = fabs(angle_diff[i]);
        if (abs_diff > max_angle) {
            max_angle = abs_diff;
            fastestIndex = i;
        }
    }
    max_velocity = fabs(newTrajectory->aim_velocity[fastestIndex]);
    if (max_velocity * (ARM_FOLLOWING_CHECKING_STEP * 2 + (max_accelerate_time_ - 1) * 0.5) >= max_angle) {
        max_velocity = max_angle / (ARM_FOLLOWING_CHECKING_STEP * 2 + (max_accelerate_time_ - 1) * 0.5);
    }
    for (int i = 0; i < USED_ARM_DOF; ++i) {
        if (max_angle < 0.0001) {
            newTrajectory->aim_velocity[i] = 0;
        } else {
            newTrajectory->aim_velocity[i] = max_velocity * angle_diff[i] / max_angle;
        }
    }
    newTrajectory->uniform_time_step += max_angle / max_velocity;
    return 0;
}

int TrapezoidalTrajectoryKinematics::ForwardKinematics(float *jointAngles, float *endEffectorPose)
{
    if (!rmService_)
    {
        return UNKNOWN_ERR; // Error: RM_Service not initialized
    }

    // Call the RM_Service's ForwardKinematics method
    rm_pose_t pose = rmService_->rm_algo_forward_kinematics(NULL, jointAngles);
    endEffectorPose[POSE_POSITION_X] = pose.position.x;
    endEffectorPose[POSE_POSITION_Y] = pose.position.y;
    endEffectorPose[POSE_POSITION_Z] = pose.position.z;
    endEffectorPose[POSE_QUATERNION_X] = pose.quaternion.x;
    endEffectorPose[POSE_QUATERNION_Y] = pose.quaternion.y;
    endEffectorPose[POSE_QUATERNION_Z] = pose.quaternion.z;
    endEffectorPose[POSE_QUATERNION_W] = pose.quaternion.w;
    return 0;
}

int TrapezoidalTrajectoryKinematics::InverseKinematics(const float *endEffectorPose, float *jointAngles)
{
    if (!rmService_)
    {
        return UNKNOWN_ERR; // Error: RM_Service not initialized
    }

    float nowJoints_[ARM_DOF] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    for (int i = 0; i < USED_ARM_DOF; i++) {
        nowJoints_[i] = jointAngles[i];
    }
    rm_inverse_kinematics_params_t inverse_params;
    memcpy(inverse_params.q_in, nowJoints_, sizeof(nowJoints_));
    inverse_params.q_pose.position.x = endEffectorPose[POSE_POSITION_X];
    inverse_params.q_pose.position.y = endEffectorPose[POSE_POSITION_Y];
    inverse_params.q_pose.position.z = endEffectorPose[POSE_POSITION_Z];
    inverse_params.q_pose.quaternion.x = endEffectorPose[POSE_QUATERNION_X];
    inverse_params.q_pose.quaternion.y = endEffectorPose[POSE_QUATERNION_Y];
    inverse_params.q_pose.quaternion.z = endEffectorPose[POSE_QUATERNION_Z];
    inverse_params.q_pose.quaternion.w = endEffectorPose[POSE_QUATERNION_W];
    inverse_params.flag = 0;

    int result = rmService_->rm_algo_inverse_kinematics(NULL, inverse_params, nowJoints_);
    if (result != 0){
        return ARM_AIM_CANNOT_REACH; // Inverse kinematics failed
    }
    memcpy(jointAngles, nowJoints_, sizeof(float) * USED_ARM_DOF);
    return 0;
}