import socket
import json
import numpy as np
import math

intersection_angle_Rad=4.76* (math.pi / 180)  #填充数据

# 连接机器人控制器
def connectETController(ip, port=8055):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.connect((ip, port))
        return (True, sock)
    except Exception as e:
        sock.close()
        return (False,)


# 发送命令并接收响应
def sendCMD(sock, cmd, params=None, id=1):
    if not params:
        params = []
    else:
        params = json.dumps(params)
    sendStr = '{"method":"%s","params":%s,"jsonrpc":"2.0","id":%d}' % (cmd, params, id) + "\n"
    try:
        sock.sendall(bytes(sendStr, "utf-8"))
        ret = sock.recv(1024)
        jdata = json.loads(str(ret, "utf-8"))
        if "result" in jdata.keys():
            return (True, json.loads(jdata["result"]), jdata["id"])
        elif "error" in jdata.keys():
            return (False, jdata["error"], jdata["id"])
        else:
            return (False, None, None)
    except Exception as e:
        return (False, None, None)


if __name__ == "__main__":
    # 机器人IP地址
    robot_ip = "192.168.1.100"
    result = connectETController(robot_ip)
    if len(result) == 2:
        conSuc, sock = result
        if conSuc:
            print("成功连接到机器人")
            Angle1 = None
            PS1 = None
            Angle2 = None
            PS2 = None

            while True:
                user_input = input("请输入Y1（开始位置）或Y2（结束位置）或Q（退出）：").strip().upper()

                if user_input == 'Q':
                    break

                # 读取当前的关节位置和末端位姿
                suc, joint_pos, id = sendCMD(sock, "get_joint_pos")
                if suc:
                    print("当前关节位置:", joint_pos)
                else:
                    print("无法获取关节位置")

                suc, tcp_pose, id = sendCMD(sock, "get_tcp_pose")
                if suc:
                    position = tcp_pose[:3]
                    # 将后三个弧度制姿态值转换为角度值
                    orientation_rad = tcp_pose[3:]
                    orientation_deg = np.rad2deg(orientation_rad)
                    tcp_pose_deg = position + orientation_deg.tolist()
                    print("当前末端位姿（位置 + 姿态角度值）:", tcp_pose_deg)
                    print("当前末端位姿（原始数据）:", tcp_pose)
                else:
                    print("无法获取末端位姿")

                if user_input == 'Y1':
                    Angle1 = joint_pos
                    PS1 = tcp_pose
                    print("开始位置数据已记录")
                elif user_input == 'Y2':
                    Angle2 = joint_pos
                    PS2 = tcp_pose
                    print("结束位置数据已记录")
                else:
                    continue

            # 输出记录的数据
            if Angle1 is not None:
                print("Start_Pose 的关节位置 (Angle1):", Angle1)
                print("Start_Pose 的末端位姿 (PS1):", PS1)
            if Angle2 is not None:
                print("End_Pose 的关节位置 (Angle2):", Angle2)
                print("End_Pose 的末端位姿 (PS2):", PS2)

            try:
                # 下面根据PS1的位姿，PS3的位置，插值一条直线
                # 设定时间间隔

                # 计算两点之间的距离
                PS3_Init = np.concatenate((PS2[:2], [PS1[2]], PS1[3:]))
                distance = np.linalg.norm(PS1 - PS3_Init)
                Z_Value=distance/math.cos(intersection_angle_Rad)*math.sin(intersection_angle_Rad)

                PS3 = np.concatenate((PS2[:2], [PS1[2] + Z_Value], PS1[3:]))
                PS4 = np.concatenate((PS2[:2], [PS1[2]], PS1[3:]))
                PS5 = np.concatenate((PS1[:2], [PS1[2] + Z_Value], PS1[3:]))

                reference_pos = Angle2
                target_pose = PS3.tolist()
                suc, angle, id = sendCMD(sock, "inverseKinematic", {"targetPose": target_pose, "referencePos": reference_pos})
                if suc:
                    angle3 = angle
                else:
                    print("逆运动学3计算失败")

                # 第三个位姿对应的关节位置
                reference_pos = Angle2
                target_pose = PS4.tolist()
                suc, angle, id = sendCMD(sock, "inverseKinematic", {"targetPose": target_pose, "referencePos": reference_pos})
                if suc:
                    angle4 = angle
                else:
                    print("逆运动学4计算失败")

                # 第四个位姿对应的关节位置
                reference_pos = Angle1
                target_pose = PS5.tolist()
                suc, angle, id = sendCMD(sock, "inverseKinematic", {"targetPose": target_pose, "referencePos": reference_pos})
                if suc:
                    angle5 = angle
                else:
                    print("逆运动学5计算失败")

                # 存储数据到txt文件
                data = {
                    "angle1": Angle1,
                    "angle2": Angle2,
                    "angle3": angle3 if 'angle3' in locals() else None,
                    "angle4": angle4 if 'angle4' in locals() else None,
                    "angle5": angle5 if 'angle5' in locals() else None,
                    "PS1": PS1,
                    "PS2": PS2,
                    "PS3": PS3.tolist() if 'PS3' in locals() else None,
                    "PS4": PS4.tolist() if 'PS4' in locals() else None,
                    "PS5": PS5.tolist() if 'PS5' in locals() else None
                }

                with open('data.txt', 'w') as f:
                    json.dump(data, f)

            except Exception as e:
                print(f"执行过程中出现错误: {e}")
            finally:
                sock.close()
        else:
            print("无法连接到机器人")
    else:
        print("无法连接到机器人")