import sys
import time
from google.protobuf import timestamp_pb2

import matplotlib.pyplot as plt  # 新增

sys.path.append("../generated")

from generated.cmvr.api import humanoid_robot_service_pb2_grpc as rpc
from generated.cmvr.api import humanoid_robot_command_pb2 as pb
from clients.base_client import RobotClientBase

# 固定关节顺序
JOINT_ORDER = [
    "L_SHOULDER_P", "L_SHOULDER_R", "L_SHOULDER_Y",
    "L_ELBOW_R", "L_WRIST_P", "L_WRIST_Y", "L_WRIST_R",
    "R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
    "R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R",
    "HEAD_P", "HEAD_Y", "HEAD_R","WAIST_Y","WAIST_R"
]


class GetJointStateClient(RobotClientBase):
    """Continuous client to fetch robot joint states"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 时间轴
        self.time_history = []
        self._start_time = None
        # 记录每个关节的角度、速度
        self.pos_history = {name: [] for name in JOINT_ORDER}
        self.vel_history = {name: [] for name in JOINT_ORDER}

    def send(self, interval=0.5, device_id="hc01"):
        """Continuously fetch joint states and print in fixed order"""
        try:
            if self._start_time is None:
                self._start_time = time.time()

            while True:
                # Construct request
                req = pb.JointRequest()
                req.header.device_id = device_id
                timestamp = timestamp_pb2.Timestamp()
                timestamp.GetCurrentTime()
                req.header.timestamp.CopyFrom(timestamp)

                # Call RPC with timeout
                try:
                    resp = self.stub.getJointState(req, timeout=10)
                except Exception as e:
                    print(f"RPC call failed: {e}")
                    time.sleep(interval)
                    continue

                # ====== 1. 同时取 position 和 velocity ======
                joint_pos_dict = {}
                joint_vel_dict = {}
                for state in resp.state:
                    # 这里就用你说的三元 zip
                    for name, pos, vel in zip(state.name, state.position, state.velocity):
                        joint_pos_dict[name] = round(pos, 12)
                        joint_vel_dict[name] = round(vel, 12)

                # ====== 2. 记录到历史数据里，用于之后画图 ======
                t = time.time() - self._start_time
                self.time_history.append(t)
                for name in JOINT_ORDER:
                    self.pos_history[name].append(joint_pos_dict.get(name, 0.0))
                    self.vel_history[name].append(joint_vel_dict.get(name, 0.0))

                # 保持原来的打印逻辑（打印角度）
                joint_list = [{"joint_name": name, "rad": joint_pos_dict.get(name, 0.0)}
                              for name in JOINT_ORDER]

                # Print timestamp
                print(f"[{time.strftime('%H:%M:%S')}]")
                # Print each joint, last one without comma
                for i, j in enumerate(joint_list):
                    if i < len(joint_list) - 1:
                        print(f"{j},")
                    else:
                        print(f"{j}")

                time.sleep(interval)

        except KeyboardInterrupt:
            print("\nStopped fetching joint states.")
            # 例子：退出时画几个关节的角度 & 速度曲线
    #         try:
    #             # 你可以根据需要改关节列表
    #             self.plot_pos_and_vel([ "R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
    # "R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"])
    #         except Exception as e:
    #             print(f"绘图失败: {e}")

    # ====== 3. 绘制速度和角度（同一张图，两行子图） ======
    def plot_pos_and_vel(self, joint_names):
        """
        在一张图中绘制多个关节的角度和速度：
        - 上 subplot：角度（rad）
        - 下 subplot：速度（rad/s）

        joint_names: str 或 [str, ...]
        """
        if isinstance(joint_names, str):
            joint_names = [joint_names]

        # 检查关节名合法性
        invalid = [n for n in joint_names if n not in self.pos_history]
        if invalid:
            print(f"未知关节名: {invalid}")
            print(f"可选关节: {list(self.pos_history.keys())}")
            return

        if not self.time_history:
            print("还没有采集到任何数据，无法绘图。")
            return

        t = self.time_history

        plt.figure(figsize=(10, 6))

        # 上：角度
        plt.subplot(2, 1, 1)
        for name in joint_names:
            plt.plot(t, self.pos_history[name], label=name)
        plt.ylabel("Position (rad)")
        plt.title("Joint Position")
        plt.grid(True)
        plt.legend()

        # 下：速度
        plt.subplot(2, 1, 2)
        for name in joint_names:
            plt.plot(t, self.vel_history[name], label=name)
        plt.xlabel("Time (s)")
        plt.ylabel("Velocity (rad/s)")
        plt.title("Joint Velocity")
        plt.grid(True)
        plt.legend()

        plt.tight_layout()
        plt.show()


if __name__ == "__main__":
    client = GetJointStateClient()
    try:
        client.send(interval=0.005)
    finally:
        client.close()