/**
 * @file arm_main.cpp
 * @brief 机械臂主任务 - C++版本
 */

#include "task/user_task.h"

#include "module/arm_oop.hpp"
#include "module/motor_base.hpp"
#include "module/joint.hpp"
#include "bsp/can.h"
#include "device/motor_lz.h"
#include "device/motor_dm.h"
using namespace arm;

// ============================================================================
// 电机和关节配置
// ============================================================================

// LZ电机参数（关节1-3）
static MOTOR_LZ_Param_t lz_params[3] = {
    { .can = BSP_CAN_1, .motor_id = 127, .host_id = 0xff, .module = MOTOR_LZ_RSO3, .reverse = true,   .mode=MOTOR_LZ_MODE_MOTION},
    { .can = BSP_CAN_1, .motor_id = 126, .host_id = 0xff, .module = MOTOR_LZ_RSO3, .reverse = true,   .mode=MOTOR_LZ_MODE_MOTION},
    { .can = BSP_CAN_1, .motor_id = 125, .host_id = 0xff, .module = MOTOR_LZ_RSO3, .reverse = false,  .mode=MOTOR_LZ_MODE_MOTION},
};

// DM电机参数（关节4-6）
static MOTOR_DM_Param_t dm_params[3] = {
    {.can = BSP_CAN_1, .master_id = 0x14,  .can_id = 0x04, .module = MOTOR_DM_J4310, .reverse = false,},
    {.can = BSP_CAN_1, .master_id = 0x15,  .can_id = 0x05, .module = MOTOR_DM_J4310, .reverse = false,},
    {.can = BSP_CAN_1, .master_id = 0x16,  .can_id = 0x06, .module = MOTOR_DM_J4310, .reverse = false,},
};

// DH参数 (单位: m, rad, kg)
// rc[3]: 质心在 DH 连杆坐标系下的坐标 (m)
// 来源：URDF <inertial><origin xyz> 经完整RPY旋转矩阵逆变换到 DH 坐标系
// 变换公式: rc_dh = R_rpy^(-1) * rc_urdf，其中R_rpy来自URDF joint的rpy参数
// 数据来源：fix2026224/arm/urdf/arm.urdf (2026-02-24 重建模更新)
//
// MIT控制参数kp/kd调优建议：
// - kp：位置刚度，值越大响应越快但易振荡。大关节（J2/J3）可适当降低避免抖动
// - kd：阻尼系数，值越大越平稳但响应慢。根据实际测试调整
// - 设置为0则使用默认值（LZ: kp=10,kd=0.5; DM: kp=50,kd=3）
static Arm6dof_DHParams_t dh_params[6] = {
  // J1: 腰部旋转（LZ电机，垂直轴，负载小）  j1 rpy=(0,0,0) → rc_dh=rc_urdf
  {.theta=0.0f, .d=0.0405f,    .a=0.0014f,   .alpha=-M_PI_2,  .theta_offset=0.0f,     .qmin=-15.7f,   .qmax=15.7f,  .m=2.3045f,  .rc={ 0.00669710f, -0.00030098f,  0.04424100f},  .kp=10.0f,  .kd=0.0f},
  // J2: 大臂俯仰（LZ电机，负载最大，降低kp避免抖动）  j2 rpy=(-90,0,-90)
  {.theta=0.0f, .d=0.0f,       .a=0.388f,    .alpha=0.0f,     .theta_offset=-M_PI_2,  .qmin=-1.57f,   .qmax=1.57f,  .m=0.8903f,  .rc={ 0.26586000f, -0.00044116f,  0.00000094f},  .kp=10.0f,  .kd=0.0f},
  // J3: 小臂（LZ电机，负载中等）  j3 rpy=(0,0,-90)
  {.theta=0.0f, .d=0.002795f,  .a=0.047775f, .alpha=-M_PI_2,  .theta_offset=-M_PI_2,  .qmin=-1.0f,    .qmax=3.0f,   .m=0.72964f, .rc={ 0.08696299f,  0.00167372f, -0.02500400f},  .kp=20.0f,  .kd=0.0f},
  // J4: 腕部旋转（DM电机，roll轴，0到2π循环）  j4 rpy=(-90,0,-90)
  {.theta=0.0f, .d=0.241f,     .a=0.0f,      .alpha=M_PI_2,   .theta_offset=M_PI,     .qmin=0.0f,     .qmax=6.3f,  .m=0.60215f, .rc={ 0.00207070f, -0.14360000f,  0.00004920f},  .kp=10.0f,  .kd=0.0f},
  // J5: 腕部俯仰（DM电机，负载小，可提高响应）  j5 rpy=(90,0,-180)
  {.theta=0.0f, .d=0.0f,       .a=0.1065f,   .alpha=-M_PI_2,  .theta_offset=0.0f,     .qmin=-1.9f,   .qmax=1.9f,  .m=0.21817f, .rc={ 0.00001750f,  0.00297341f,  0.05816899f},  .kp=15.0f,  .kd=0.0f},
  //  {.theta=0.0f, .d=0.0f,       .a=0.1065f,   .alpha=-M_PI_2,  .theta_offset=M_PI_2,   .qmin=-1.9f,   .qmax=1.9f,  .m=0.21817f, .rc={ 0.00001750f,  0.00297341f,  0.05816899f},  .kp=15.0f,  .kd=0.0f},
  // J6: 末端Roll（DM电机，最小负载，roll轴，0到2π循环）  j6 rpy=(90,0,180)
  {.theta=0.0f, .d=0.0040025f, .a=0.0f,      .alpha=0.0f,     .theta_offset=M_PI,     .qmin=0.0f,     .qmax=6.3f,   .m=0.57513f, .rc={-0.00002134f,  0.09993500f,  0.00053860f},  .kp=5.0f,  .kd=0.0f},
};

static float q_offset[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
// 全局对象指针
IMotor* motors[6] = {nullptr};       // 多态接口
MotorLZ* motors_lz[3] = {nullptr};   // 用于调试查看
MotorDM* motors_dm[3] = {nullptr};   // 用于调试查看
Joint* joints[6] = {nullptr};
RoboticArm* robot_arm = nullptr;

Arm6dof_Pose_t target_pose = {0};  // 控制输入：目标末端位姿，单位 m/rad（内部计算全链路SI单位）

// ============================================================================
// 测试用调试变量（可在调试器Watch窗口直接观察和修改）
// ============================================================================
// 阶段控制：在调试器中修改此值切换测试阶段
//   0 = 仅计算重力补偿力矩，全部电机松弛，观察 gravity_torques_dbg 是否合理
//   1 = 仅关节6（Roll轴）输出纯力矩，其余松弛（Roll轴重力补偿接近0，最安全）
//   2 = 全部六轴输出（完整重力补偿，GRAVITY_COMP模式）
//   3 = 笛卡尔空间控制：控制末端目标位姿，逆运动学解算各关节角度
uint8_t test_stage = 3;
Arm6dof_JointAngles_t current_angles;  // 调试观察：当前各关节角度 (rad)
// 重力补偿力矩观察数组（对应关节1~6，单位 N·m）
float gravity_torques_dbg[6] = {0.0f};
// 末端位姿观察（mm换算，便于观察，不参与计算）
struct { float x, y, z, roll, pitch, yaw; } end_pose_mm_dbg = {0};
// FK→IK验证：IK解算得到的关节角度
Arm6dof_JointAngles_t ik_from_fk_result;  // IK解算结果

int setzero=0;
// 重力补偿缩放系数（可在调试器中实时修改，用于测试补偿强度）
float gravity_comp_scale = 1.0f;
// 轨迹进度观察 [0.0~1.0]，1.0=已到达目标；调试器中可观察运动是否平滑完成
float traj_progress_dbg = 0.0f;
// 轨迹速度设置：可在调试器中修改这两个值，重启case4生效
float traj_lin_vel = 0.15f;   // 末端线速度 (m/s)，默认 150mm/s
float traj_ang_vel = 1.0f;    // 末端角速度 (rad/s)，默认 ~57°/s
// 同步控制：设置为1时将target_pose同步为当前位姿（世界系，供MoveCartesian使用），同步后自动清零
int sync_target_to_current = 0;
// 航向系同步：设置为1时将target_pose同步为当前位姿的航向系表示（供MoveCartesianHeading使用），同步后自动清零
// 位置：p_h = Rz(-yaw)*p_w
// 姿态：R_h = Rz(-yaw)*R_cur 提取RPY，同步后直接修改 yaw/pitch/roll 即可直观控制末端各轴旋转
int sync_heading_to_current = 0;
// 工具系同步：设置为1时将target_pose同步为当前位姿的工具系表示（供MoveCartesianTool使用），同步后自动清零
// 位置：p_tool = R_cur^T * p_world
// 姿态：roll=pitch=yaw=0 即 R_target=I，机械臂保持当前末端姿态不变
int sync_tool_to_current = 0;

extern "C" {

void Task_arm_main(void* argument) {

  /* 计算任务运行到指定频率需要等待的tick数 */
  const uint32_t delay_tick = osKernelGetTickFreq() / ARM_MAIN_FREQ;
  osDelay(ARM_MAIN_INIT_DELAY); /* 延时一段时间再开启任务 */
  uint32_t tick = osKernelGetTickCount(); /* 控制任务运行频率的计时 */

  BSP_CAN_Init();
  MOTOR_LZ_Init();  // 注册LZ电机ID解析器
  
  // 初始化运动学模型
  Arm6dof_Init(dh_params);
  
  // 创建电机对象（同时保存到基类和派生类指针）
  motors_lz[0] = new MotorLZ(lz_params[0]);
  motors_lz[1] = new MotorLZ(lz_params[1]);
  motors_lz[2] = new MotorLZ(lz_params[2]);
  motors_dm[0] = new MotorDM(dm_params[0]);
  motors_dm[1] = new MotorDM(dm_params[1]);
  motors_dm[2] = new MotorDM(dm_params[2]);
  
  // 填充基类指针数组（用于多态管理）
  motors[0] = motors_lz[0];
  motors[1] = motors_lz[1];
  motors[2] = motors_lz[2];
  motors[3] = motors_dm[0];
  motors[4] = motors_dm[1];
  motors[5] = motors_dm[2];
  for (int i = 0; i < 6; ++i) {
      joints[i] = new Joint(i, motors[i], dh_params[i],  q_offset[i], ARM_MAIN_FREQ);
  }
  robot_arm = new RoboticArm();
  for (int i = 0; i < 6; ++i) {
      robot_arm->AddJoint(i, joints[i]);
  }
  robot_arm->Init();
  robot_arm->Enable(true);
  
  // 使能重力补偿
  robot_arm->EnableGravityCompensation(true);
  robot_arm->SetGravityCompScale(gravity_comp_scale);  // 补偿系数，可微调
  
  // 设置控制模式（可选以下模式之一）:
  // GRAVITY_COMP: 位置保持 + 重力补偿前馈（位控+补偿）
  // TEACH:        纯重力补偿力矩输出，零刚度（示教拖动）
  // JOINT_POSITION: 关节位置控制 + 重力补偿前馈
  robot_arm->SetMode(ControlMode::GRAVITY_COMP);

  // 将 target_pose 初始化为当前末端位姿，防止 stage4 首次进入时进行大范围跳变
  robot_arm->Update();  // 先跨一帧读取当前状态
  target_pose = robot_arm->GetEndPose();
  
  while (1) {
    tick += delay_tick; /* 计算下一个唤醒时刻 */
/*---------------------------零点校准---------------------------*/
    if (setzero) {
      if (setzero <= 3) {
        MOTOR_LZ_SetZero(&lz_params[setzero-1]);
      } else if (setzero > 3&&setzero < 7) {
        MOTOR_DM_SetZero(&dm_params[setzero-4]);
      } else if (setzero == 7) {
        for (int i = 0; i < 3; ++i) {
          MOTOR_LZ_SetZero(&lz_params[i]);
        }
        for (int i = 0; i < 3; ++i) {
          MOTOR_DM_SetZero(&dm_params[i]);
        }
      }
      setzero = 0;
    }
/*--------------------------------------------------------------*/

    // 更新机械臂状态（读取关节角度、FK、重力补偿计算）
    robot_arm->Update();
    
    // 同步target到current（调试时手动触发）
    if (sync_target_to_current) {
      target_pose = robot_arm->GetEndPose();
      sync_target_to_current = 0;  // 自动清零
    }
    // 航向系同步：将target_pose设为当前位姿的航向系表示，供MoveCartesianHeading使用
    // 位置：Rz(-yaw) * p_world → 航向系坐标
    // 姿态：R_heading = Rz(-yaw) * R_cur，提取RPY后写入target_pose
    //       同步后 yaw/pitch/roll 即为末端在航向系下的当前姿态，修改哪个轴就转哪个轴
    if (sync_heading_to_current) {
      const Arm6dof_Pose_t& ep = robot_arm->GetEndPose();
      float cy = cosf(ep.yaw), sy = sinf(ep.yaw);
      // 位置逆变换
      target_pose.x = cy * ep.x + sy * ep.y;
      target_pose.y = -sy * ep.x + cy * ep.y;
      target_pose.z = ep.z;
      // 姿态逆变换：R_h = Rz(-yaw) * R_cur，再提取RPY
      float rpy_cur_arr[3] = {ep.yaw, ep.pitch, ep.roll};
      float rpy_mzy_arr[3] = {-ep.yaw, 0.0f, 0.0f};  // Rz(-yaw)
      Matrixf<3,3> R_h = robotics::rpy2r(Matrixf<3,1>(rpy_mzy_arr))
                       * robotics::rpy2r(Matrixf<3,1>(rpy_cur_arr));
      Matrixf<3,1> rpy_h = robotics::r2rpy(R_h);
      target_pose.yaw   = rpy_h[0][0];
      target_pose.pitch = rpy_h[1][0];
      target_pose.roll  = rpy_h[2][0];
      sync_heading_to_current = 0;  // 自动清零
    }
    // 工具系同步：将target_pose设为当前位姿的工具系表示，供MoveCartesianTool使用
    // 位置：p_tool = R_cur^T * p_world
    // 姿态：roll=pitch=yaw=0 即 R_target=I → R_world=R_cur，保持当前末端姿态不变
    if (sync_tool_to_current) {
      const Arm6dof_Pose_t& ep = robot_arm->GetEndPose();
      float rpy_cur_arr[3] = {ep.yaw, ep.pitch, ep.roll};
      Matrixf<3,3> R = robotics::rpy2r(Matrixf<3,1>(rpy_cur_arr));
      float p_arr[3] = {ep.x, ep.y, ep.z};
      Matrixf<3,1> p_tool = R.trans() * Matrixf<3,1>(p_arr);  // R^T * p_world
      target_pose.x     = p_tool[0][0];
      target_pose.y     = p_tool[1][0];
      target_pose.z     = p_tool[2][0];
      target_pose.yaw   = 0.0f;  // R_target = I，保持姿态
      target_pose.pitch = 0.0f;
      target_pose.roll  = 0.0f;
      sync_tool_to_current = 0;  // 自动清零
    }

    // 每帧同步调试观察变量
    for (int i = 0; i < 6; ++i) {
      current_angles.q[i] = joints[i]->GetCurrentAngle();
      gravity_torques_dbg[i] = robot_arm->GetGravityTorque(i);
    }
    
    // 末端位姿换算为mm，便于调试器观察（不参与任何计算）
    const Arm6dof_Pose_t& ep = robot_arm->GetEndPose();
    end_pose_mm_dbg = { ep.x * 1000.0f, ep.y * 1000.0f, ep.z * 1000.0f,
                        ep.roll, ep.pitch, ep.yaw };

      switch (test_stage) {
        case 0:
          // 阶段1：仅计算，不输出。观察 gravity_torques_dbg[0~5] 是否量级合理
          for (int i = 0; i < 6; ++i) motors[i]->Relax();
          break;

        case 1:
          // 阶段1：测试单个关节力矩输出（用于验证力矩计算）
          // 修改下面的索引来测试不同关节：j2=1, j3=2, j5=4
          motors[0]->Relax();
          motors[1]->Relax();
          motors[2]->Relax();
          motors[3]->Relax();
          motors[4]->Relax();
          motors[5]->Relax();
          // joints[1]->TorqueControl(gravity_torques_dbg[1]);  // j2大臂
          // joints[2]->TorqueControl(gravity_torques_dbg[2]);  // j3小臂
          // joints[4]->TorqueControl(gravity_torques_dbg[4]);  // j5腕部

          // joints[0]->TorqueControl(gravity_torques_dbg[0]); 
          // joints[3]->TorqueControl(gravity_torques_dbg[3]); 
          // joints[5]->TorqueControl(gravity_torques_dbg[5]); 
          break;
        case 2:
        default:
          // 阶段4：全部六轴 GRAVITY_COMP 模式（位置保持 + 重力补偿前馈）
          robot_arm->SetMode(ControlMode::GRAVITY_COMP);
          robot_arm->Control();
          break;

        case 3:
          // 笛卡尔空间轨迹规划控制
          // 修改 target_pose 后，机械臂自动从当前位姿平滑运动到目标位姿
          robot_arm->SetLinVelLimit(traj_lin_vel);
          robot_arm->SetAngVelLimit(traj_ang_vel);
          robot_arm->SetMode(ControlMode::CARTESIAN_POSITION);
          // robot_arm->MoveCartesianTool(target_pose);
          // robot_arm->MoveCartesianHeading(target_pose);
          robot_arm->MoveCartesian(target_pose);
          robot_arm->Control();

          ik_from_fk_result  = robot_arm->GetIkAngles();
          traj_progress_dbg  = robot_arm->GetTrajProgress();
          break;
      }

    osDelayUntil(tick); /* 运行结束，等待下一次唤醒 */
  }
} 

} // extern "C"