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

#include "cmsis_os2.h"
#include "task/user_task.h"

#include "module/arm_oop.hpp"
#include "module/gripper_rm.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 = true,},
    {.can = BSP_CAN_1, .master_id = 0x16,  .can_id = 0x06, .module = MOTOR_DM_J4310, .reverse = false,},
};

static Arm6dof_DHParams_t dh_params[6] = {
  // J1: origin(0, 0, 0.024) rpy(0,0,0)
  {.theta=0.0f, .d=0.024f,     .a=0.000f,    .alpha=0.0f,     .theta_offset=0.0f,      .qmin=-15.7f,   .qmax=15.7f,  .m=2.2629f,  .rc={-0.00718190f,  0.00031034f,  0.06159800f}},
  // J2: origin(-0.001395, 0, 0.1015) rpy(1.5708, 0, -1.5708)
  {.theta=0.0f, .d=0.1015f,    .a=-0.00139f, .alpha=M_PI_2,   .theta_offset=-M_PI_2,   .qmin=-1.57f,   .qmax=1.57f,  .m=0.97482f, .rc={ 0.00316320f, -0.00227070f, -0.22947000f}},
  // J3: origin(0.3265, 0, -0.0071975) rpy(0, 0, 0)
  {.theta=0.0f, .d=-0.00719f,  .a=0.3265f,   .alpha=0.0f,     .theta_offset=0.0f,      .qmin=-1.0f,    .qmax=3.0f,   .m=0.72964f, .rc={ 0.08696300f,  0.00167340f, -0.01780600f}},
  // J4: origin(0.0905, 0.052775, 0.0058025) rpy(1.5708, 0, 3.1416)
  {.theta=0.0f, .d=0.05278f,   .a=0.0905f,   .alpha=M_PI_2,   .theta_offset=M_PI,      .qmin=0.0f,     .qmax=6.3f,   .m=0.54148f, .rc={-0.00005530f,  0.10972717f, -0.00230270f}},
  // J5: origin(0.001627, 0, 0.18467) rpy(1.5708, 0, 0)
  {.theta=0.0f, .d=0.18467f,   .a=0.0016f,   .alpha=M_PI_2,   .theta_offset=0.0f,      .qmin=-1.9f,    .qmax=1.9f,   .m=0.21817f, .rc={ 0.05654200f,  0.00102680f, -0.00131060f}},
  // J6: origin(0.10487, 0.0013347, 0) rpy(1.5708, 0, 1.5708)
  {.theta=0.0f, .d=0.00133f,   .a=0.10487f,  .alpha=M_PI_2,   .theta_offset=M_PI_2,    .qmin=0.0f,     .qmax=6.3f,   .m=0.56255f, .rc={ 0.00001277f,  0.10159000f, -0.00000780f}},
};

static JointControlParams joint_params[6] = {
    {.qmin=-15.7f, .qmax=15.7f, .kp=10.0f, .kd=0.0f},
    {.qmin=-1.57f, .qmax=1.57f, .kp=10.0f, .kd=0.0f},
    {.qmin=-1.0f,  .qmax=3.0f,  .kp=20.0f, .kd=0.0f},
    {.qmin=0.0f,   .qmax=6.3f,  .kp=3.0f,  .kd=0.0f},
    {.qmin=-1.9f,  .qmax=1.9f,  .kp=5.0f,  .kd=0.0f},
    {.qmin=0.0f,   .qmax=6.3f,  .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;

Arm_CMD_t arm_cmd;  // 当前机械臂控制命令（含 target_pose / set_target_as_current）

static MOTOR_RM_Param_t MakeGripperRMParam() {
  MOTOR_RM_Param_t param{};
  param.can = BSP_CAN_1;
  param.id = 0x201;
  param.module = MOTOR_M3508;
  param.reverse = false;
  param.gear = true;
  return param;
}

// ============================================================================
// 测试用调试变量（可在调试器Watch窗口直接观察和修改）
// ============================================================================
// 阶段控制：在调试器中修改此值切换测试阶段
//   0 = 仅计算重力补偿力矩，全部电机松弛，观察 gravity_torques_dbg 是否合理
//   1 = 仅关节6（Roll轴）输出纯力矩，其余松弛（Roll轴重力补偿接近0，最安全）
//   2 = 全部六轴输出（完整重力补偿，GRAVITY_COMP模式）
//   3 = 笛卡尔空间控制（旧版满秩库代数雅可比算法）
//   4 = 2.5D 解析降维控制（专为比赛设计的无奇异点稳定版，推荐！）
uint8_t test_stage = 5;
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_scales[6] = {1.0f, 2.0f, 1.5f, 2.0f, 2.6f, 1.0f};
// 轨迹进度观察 [0.0~1.0]，1.0=已到达目标；调试器中可观察运动是否平滑完成
float traj_progress_dbg = 0.0f;
// 轨迹速度设置：可在调试器中修改这两个值，重启case4生效
float traj_lin_vel = 0.03f;   // 末端线速度 (m/s)，默认 150mm/s
float traj_ang_vel = 0.2f;    // 末端角速度 (rad/s)，默认 ~57°/s

uint8_t control_frame = 1;  // 选择在何种坐标系下控制 1= 世界系，2=工具系，3=航向系
// ============================================================================
// IK测试调试变量
// ============================================================================
// ik_test_enable：设置为1时每帧调用一次IK测试；设置为0停止
//   - 输入：target_pose（目标位姿）+ current_angles（当前角度作为初始猜测）
//   - 结果：ik_test_result（IK解算出的关节角度），ik_test_ret（返回码：0=成功 -1=失败）
int ik_test_enable = 0;
Arm6dof_JointAngles_t ik_test_result = {0};  // IK解算结果（rad），调试器Watch观察
int ik_test_ret = 0;                          // IK返回码：0=成功，-1=无解或未初始化
int ik_test_ret_analytical = 0;               // 解析IK原始返回码
int ik_test_ret_numerical = 0;                // 数值IK原始返回码（仅回退时更新）
int ik_test_solver_dbg = 0;                   // 0=解析成功, 1=数值回退成功, -1=全部失败

// cmd task 中的机械臂命令指针（cmd.cpp 中定义），通过 extern 访问以便写回初始位姿
extern Arm_CMD_t *cmd_for_arm;

/**
 * @brief 将当前世界系末端位姿转换到指定控制坐标系下的 target_pose，
 *        使 MoveCartesian*(result) 的目标恰好等于当前 cur_world（无初始跳变）
 * @param frame  1=世界系  2=工具系  3=航向系
 * @param cur    当前末端位姿（世界系，由 FK 计算得到）
 */
static Arm6dof_Pose_t SyncTargetToFrame(uint8_t frame, const Arm6dof_Pose_t& cur) {
    Arm6dof_Pose_t target;
    switch (frame) {
        case 2: {
            // 工具系：p_tool = R_cur^T * p_world，使 R_cur*p_tool = p_world
            // 姿态置零：R_tool = I，使 R_world = R_cur * I = R_cur（保持当前姿态）
            float rpy[3] = {cur.yaw, cur.pitch, cur.roll};
            Matrixf<3,3> RT = robotics::rpy2r(Matrixf<3,1>(rpy)).trans();
            float p[3] = {cur.x, cur.y, cur.z};
            Matrixf<3,1> pt = RT * Matrixf<3,1>(p);
            target.x = pt[0][0]; target.y = pt[1][0]; target.z = pt[2][0];
            target.yaw = 0.0f; target.pitch = 0.0f; target.roll = 0.0f;
            break;
        }
        case 3: {
            // 航向系：p_heading = Rz(-yaw)*[x,y]，z 直接赋值
            // 姿态：R_heading = Rz(-yaw)*R_cur，使 Rz(yaw)*R_heading = R_cur
            float cy = cosf(cur.yaw), sy = sinf(cur.yaw);
            target.x =  cy * cur.x + sy * cur.y;
            target.y = -sy * cur.x + cy * cur.y;
            target.z = cur.z;
            float rpy_neg[3] = {-cur.yaw, 0.0f, 0.0f};
            float rpy_c[3]   = { cur.yaw, cur.pitch, cur.roll};
            Matrixf<3,3> Rn  = robotics::rpy2r(Matrixf<3,1>(rpy_neg))
                             * robotics::rpy2r(Matrixf<3,1>(rpy_c));
            Matrixf<3,1> rn  = robotics::r2rpy(Rn);
            target.yaw = rn[0][0]; target.pitch = rn[1][0]; target.roll = rn[2][0];
            break;
        }
        default:
            // 世界系：直接复制
            target = cur;
            break;
    }
    return target;
}

static void SyncCommandTargetFromCurrent() {
  if (!robot_arm) {
    return;
  }

  const Arm6dof_Pose_t sync = SyncTargetToFrame(control_frame, robot_arm->GetEndPose());
  Arm6dof_JointAngles_t sync_joints = arm_cmd.target_joints;
  for (int i = 0; i < 6; ++i) {
    if (joints[i]) {
      sync_joints.q[i] = joints[i]->GetCurrentAngle();
    }
  }

  arm_cmd.target_pose = sync;
  arm_cmd.target_joints = sync_joints;

  // 同步上游CMD中的目标值，避免下一帧被旧目标覆盖导致跳变。
  // 注意：这里只同步目标量，不修改夹爪toggle等一次性控制位。
  if (cmd_for_arm) {
    cmd_for_arm->target_pose = sync;
    cmd_for_arm->target_joints = sync_joints;
  }
}

static void HandleSetZeroRequest() {
  if (!setzero) {
    return;
  }

  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;
}

static void RelaxAllMotors() {
  for (int i = 0; i < 6; ++i) {
    motors[i]->Relax();
  }
}

static void ApplyJointTorques(const float torques[6]) {
  for (int i = 0; i < 6; ++i) {
    joints[i]->TorqueControl(torques[i]);
  }
}

static void MoveBySelectedFrame(const Arm6dof_Pose_t& target) {
  switch (control_frame) {
    case 2:
      robot_arm->MoveCartesianTool(target);
      break;
    case 3:
      robot_arm->MoveCartesianHeading(target);
      break;
    default:
      robot_arm->MoveCartesian(target);
      break;
  }
}

static void HandleCartesianModeError() {
  if (robot_arm->GetState() == MotionState::ERROR) {
    robot_arm->ResetError();
    SyncCommandTargetFromCurrent();
  }
}

static void RunCartesianControl(ControlMode mode) {
  robot_arm->SetLinVelLimit(traj_lin_vel);
  robot_arm->SetAngVelLimit(traj_ang_vel);
  robot_arm->SetMode(mode);
  MoveBySelectedFrame(arm_cmd.target_pose);
  robot_arm->Control();
  HandleCartesianModeError();
  ik_from_fk_result = robot_arm->GetIkAngles();
  traj_progress_dbg = robot_arm->GetTrajProgress();
}

static void RunJointControlFromCmd() {
  float target_joints[6];
  for (int i = 0; i < 6; ++i) {
    target_joints[i] = arm_cmd.target_joints.q[i];
  }

  robot_arm->SetMode(ControlMode::JOINT_POSITION);
  if (robot_arm->MoveJoint(target_joints) != 0) {
    // 超限时回收错误并把命令吸附到当前角，避免持续触发错误
    robot_arm->ResetError();
    SyncCommandTargetFromCurrent();
  }
  robot_arm->Control();
  ik_from_fk_result = robot_arm->GetIkAngles();
  traj_progress_dbg = 0.0f;
}

static void RunTestStageControl() {
  switch (test_stage) {
    case 0: {
      // 检测电机正电流旋转方向与urdf规定是否一致
      static float rotationDirectionCheck[6] = {0.0f};
      ApplyJointTorques(rotationDirectionCheck);
      break;
    }

    case 1:
      // 阶段1：仅计算，不输出。观察 gravity_torques_dbg[0~5] 是否量级合理
      RelaxAllMotors();
      break;

    case 2:
      // 阶段2：输出重力补偿力矩，验证各轴方向与量级
      ApplyJointTorques(gravity_torques_dbg);
      break;

    case 4:
      // 笛卡尔空间轨迹规划控制：基于数值逆解(雅可比)
      RunCartesianControl(ControlMode::CARTESIAN_POSITION);
      break;

    case 5:
      // 2.5D 解析降维控制（支持关节/笛卡尔两种上层指令）
      if (arm_cmd.ctrl_type == ARM_CTRL_CUSTOM_JOINT) {
        RunJointControlFromCmd();
      } else {
        RunCartesianControl(ControlMode::CARTESIAN_ANALYTICAL);
      }
      break;

    case 3:
    default:
      // 默认：全部六轴 GRAVITY_COMP 模式（位置保持 + 重力补偿前馈）
      robot_arm->SetMode(ControlMode::GRAVITY_COMP);
      robot_arm->Control();
      break;
  }
}

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解析器
  Gripper_Init(MakeGripperRMParam());
  
  // 初始化运动学模型
  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], joint_params[i], q_offset[i], ARM_MAIN_FREQ);
  }
  // 关节3与关节2存在机械耦合：J3电机读数包含J2的运动，需减去J2角度才是真实J3角度
  joints[2]->SetCoupledJoint(joints[1]);

  robot_arm = new RoboticArm();
  for (int i = 0; i < 6; ++i) {
      robot_arm->AddJoint(i, joints[i]);
  }
  robot_arm->Init();
  
  // 使能重力补偿
  robot_arm->EnableGravityCompensation(true);
  robot_arm->SetGravityCompScales(gravity_comp_scales);
  
  // 设置控制模式（可选以下模式之一）:
  // GRAVITY_COMP: 位置保持 + 重力补偿前馈（位控+补偿）
  // TEACH:        纯重力补偿力矩输出，零刚度（示教拖动）
  // JOINT_POSITION: 关节位置控制 + 重力补偿前馈
  robot_arm->SetMode(ControlMode::GRAVITY_COMP);

  // 读取当前末端位姿，转换到 control_frame 后同步到 cmd，防止上电时大范围跳变
  osDelay(100);
  robot_arm->Update();
  {
    SyncCommandTargetFromCurrent();
  }

  while (1) {
    tick += delay_tick; /* 计算下一个唤醒时刻 */
/*---------------------------零点校准---------------------------*/
    HandleSetZeroRequest();
/*--------------------------------------------------------------*/

    // 每帧允许在调试器中独立修改各轴补偿系数
    robot_arm->SetGravityCompScales(gravity_comp_scales);

    // 更新机械臂状态（读取关节角度、FK、重力补偿计算）
    robot_arm->Update();

    const osStatus_t arm_cmd_rx = osMessageQueueGet(task_runtime.msgq.arm.cmd, &arm_cmd, NULL, 0);
    if (arm_cmd_rx != osOK) {
      // 未收到新命令时清理单帧脉冲，避免重复消费上一帧toggle类信号
      arm_cmd.gripper_toggle = false;
      arm_cmd.set_target_as_current = false;
    }

    // 夹爪两态控制：直接跟随 gripper_close，内部状态机负责稳夹优先策略
    Gripper_Step(arm_cmd.gripper_close);
    
    robot_arm->Enable(arm_cmd.enable);

    // 使能上升沿自动同步：避免使能时 target_pose 与当前实际位姿不一致导致初始跳变
    static bool last_enable = false;
    if (arm_cmd.enable && !last_enable) {
      SyncCommandTargetFromCurrent();
    }
    last_enable = arm_cmd.enable;

    // set_target_as_current：将 target_pose 同步到当前末端位姿（转换到当前 control_frame）
    if (arm_cmd.set_target_as_current) {
      SyncCommandTargetFromCurrent();
    }

    // 控制类型切换沿：同步目标，避免 pose/joint 切换时命令继承旧目标导致突变
    static Arm_CtrlType_t last_ctrl_type = ARM_CTRL_REMOTE_CARTESIAN;
    if (arm_cmd.enable && arm_cmd.ctrl_type != last_ctrl_type) {
      SyncCommandTargetFromCurrent();
    }
    last_ctrl_type = arm_cmd.ctrl_type;

    
    // 每帧同步调试观察变量
    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 };

    RunTestStageControl();

    // IK测试：以当前角度为初始猜测，对 target_pose 求逆运动学
    // 在调试器中将 ik_test_enable 置1启用，观察 ik_test_result 和 ik_test_ret
    // if (ik_test_enable) {
    //   ik_test_ret_analytical = robot_arm->InverseKinematicsAnalytical(&arm_cmd.target_pose,
    //                                                                    &ik_test_result,
    //                                                                    &current_angles);
    //   ik_test_ret_numerical = 0;
    //   ik_test_solver_dbg = -1;

    //   // 解析解失败时回退到数值IK，便于调试观察并区分“无解”与“分支失败”。
    //   if (ik_test_ret_analytical == 0) {
    //     ik_test_ret = 0;
    //     ik_test_solver_dbg = 0;
    //   } else {
    //     ik_test_ret_numerical = Arm6dof_InverseKinematics(&arm_cmd.target_pose,
    //                                                       &current_angles,
    //                                                       &ik_test_result,
    //                                                       0.001f,
    //                                                       50);
    //     if (ik_test_ret_numerical == 0) {
    //       ik_test_ret = 0;
    //       ik_test_solver_dbg = 1;
    //     } else {
    //       ik_test_ret = -1;
    //       ik_test_solver_dbg = -1;
    //     }
    //   }
    // }

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

} // extern "C"