import math
import time
import sys
from collections import deque
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import matplotlib.font_manager as fm

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

from clients._path_setup import ensure_paths

ensure_paths()

from google.protobuf import timestamp_pb2

from clients.base_client import RobotClientBase

from cmvr.api import common_pb2
from cmvr.api import arm_command_pb2 as pb
from cmvr.api import arm_service_pb2_grpc as rpc


class ServoJClient(RobotClientBase):
    """Client to send ServoJ commands to the robot"""

    def send(self, joint_list, vel=1.0, device_id="right_arm"):
        """
        joint_list: list of dicts, e.g.,
            [
                {"joint_name": "L_SHOULDER_P", "rad": 0.0},
                ...
            ]
        """
        target = pb.JointPositionCommand(position=[j["rad"] for j in joint_list])

        req = pb.ServoJ.Request(
            header=self._create_header(device_id),
            target=target,
        )

        try:
            start_time = time.time()
            resp = self.stub.servoJ(req, timeout=10)
            rpc_time = time.time() - start_time
            success = getattr(resp.header, "success", None)
            error_msg = getattr(resp.header, "error_message", "")
            return success, error_msg, rpc_time
        except Exception as e:
            return False, str(e), 0.0

    def _create_header(self, device_id):
        header = common_pb2.CommandHeader.Request()
        header.device_id = device_id
        ts = timestamp_pb2.Timestamp()
        ts.GetCurrentTime()
        header.timestamp.CopyFrom(ts)
        return header


def deg2rad(degrees):
    return degrees * math.pi / 180.0


class DelayPlotter:
    """实时延迟波形绘制器 - 优化版，减少闪烁"""

    def __init__(self, max_points=300):
        self.max_points = max_points
        self.rpc_delays = deque(maxlen=max_points)
        self.running = True
        self.last_update = time.time()
        self.update_interval = 0.2  # 每200ms更新一次图表，减少重绘频率

        # 创建图形 - 使用双缓冲减少闪烁
        plt.ioff()  # 关闭交互模式
        self.fig, self.ax = plt.subplots(figsize=(12, 6))
        self.fig.suptitle('gRPC Call Delay Monitor', fontsize=14)

        # RPC延迟图
        self.ax.set_ylabel('Delay (ms)')
        self.ax.set_xlabel('Sample Point')
        self.ax.set_title('gRPC Call Delay Over Time')
        self.ax.grid(True, alpha=0.3)
        self.line, = self.ax.plot([], [], 'b-', linewidth=1)
        self.ax.axhline(y=5.0, color='r', linestyle='--', alpha=0.5)

        # 设置固定的Y轴范围，避免自动缩放导致的闪烁
        self.ax.set_ylim(0, 20)

        # 设置固定的X轴范围
        self.ax.set_xlim(0, max_points)

        plt.tight_layout()

        # 启用双缓冲
        self.fig.set_dpi(100)

    def add_data(self, rpc_time):
        """添加新数据点"""
        self.rpc_delays.append(rpc_time * 1000)  # 转换为ms

    def update_plot(self):
        """手动更新图表，不使用动画"""
        if not self.running or len(self.rpc_delays) < 2:
            return

        # 控制更新频率
        current_time = time.time()
        if current_time - self.last_update < self.update_interval:
            return
        self.last_update = current_time

        x_data = list(range(len(self.rpc_delays)))
        y_data = list(self.rpc_delays)

        # 更新数据
        self.line.set_data(x_data, y_data)

        # 更新X轴范围，但保持平滑
        if len(x_data) > 0:
            self.ax.set_xlim(max(0, len(x_data) - self.max_points), len(x_data))

        # 重绘
        self.fig.canvas.draw_idle()  # 使用draw_idle代替draw，减少阻塞
        self.fig.canvas.flush_events()

    def start(self):
        """启动绘图"""
        plt.ion()  # 开启交互模式
        plt.show(block=False)

    def stop(self):
        """停止绘图"""
        self.running = False
        plt.close(self.fig)


if __name__ == "__main__":
    # 初始化客户端
    client = ServoJClient()
    joint_name = 'R_WRIST_R'

    # 固定发送0度位置
    fixed_angle = 0

    # 发送频率
    freq = 200
    dt = 1.0 / freq  # 5ms

    print(f"=== gRPC Delay Test ===")
    print(f"Joint: {joint_name}")
    print(f"Command: fixed angle {fixed_angle}°")
    print(f"Send frequency: {freq}Hz (interval {dt * 1000:.1f}ms)")
    print("-" * 50)
    print("Press Ctrl+C to stop")
    print("-" * 50)

    # 初始化绘图器
    plotter = DelayPlotter(max_points=300)
    plotter.start()

    # 统计
    sent_count = 0
    rpc_times = deque(maxlen=100)
    start_time = time.time()

    # 用于定期打印统计信息
    last_stats_time = time.time()
    stats_interval = 2.0  # 每2秒打印一次统计

    # 用于控制显示更新的频率
    last_display_time = time.time()
    display_interval = 0.05  # 每50ms更新一次控制台显示

    try:
        while True:  # 无限循环
            loop_start = time.time()

            # 发送固定角度的命令
            joint_cmd = [{'joint_name': joint_name, 'rad': deg2rad(fixed_angle)}]
            success, error, rpc_time = client.send(joint_cmd, vel=0.5)

            # 记录
            rpc_times.append(rpc_time)
            sent_count += 1

            # 添加到绘图器
            plotter.add_data(rpc_time)

            # 更新图表（受update_plot内部的频率控制）
            plotter.update_plot()

            # 控制控制台显示更新频率
            current_time = time.time()
            if current_time - last_display_time >= display_interval:
                elapsed = current_time - start_time
                avg_rpc = sum(rpc_times) / len(rpc_times) * 1000
                current_rpc = rpc_time * 1000

                sys.stdout.write(
                    f"\rTime: {elapsed:5.1f}s | "
                    f"Sent: {sent_count:6d} | "
                    f"Current: {current_rpc:6.2f}ms | "
                    f"Avg: {avg_rpc:6.2f}ms   "
                )
                sys.stdout.flush()
                last_display_time = current_time

            # 定期打印详细统计信息（每2秒）
            if current_time - last_stats_time >= stats_interval:
                if rpc_times:
                    delays_ms = [t * 1000 for t in rpc_times]
                    print(f"\n[Stats @ {current_time - start_time:.1f}s] "
                          f"Avg:{np.mean(delays_ms):.2f}ms "
                          f"Max:{np.max(delays_ms):.2f}ms "
                          f"Min:{np.min(delays_ms):.2f}ms "
                          f"P95:{np.percentile(delays_ms, 95):.2f}ms")
                last_stats_time = current_time

            # 控制发送频率
            elapsed = time.time() - loop_start
            if elapsed < dt:
                time.sleep(dt - elapsed)

    except KeyboardInterrupt:
        print("\n\n⏹️ Test stopped by user")

        # 最终统计
        total_time = time.time() - start_time
        print(f"\n=== Final Statistics ===")
        print(f"Test duration: {total_time:.2f}s")
        print(f"Total commands: {sent_count}")
        print(f"Actual frequency: {sent_count / total_time:.1f}Hz")
        if rpc_times:
            delays_ms = [t * 1000 for t in rpc_times]
            print(f"Average RPC: {np.mean(delays_ms):.2f}ms")
            print(f"Maximum RPC: {np.max(delays_ms):.2f}ms")
            print(f"Minimum RPC: {np.min(delays_ms):.2f}ms")
            print(f"StdDev RPC: {np.std(delays_ms):.2f}ms")
            print(f"P95 RPC: {np.percentile(delays_ms, 95):.2f}ms")
            print(f"P99 RPC: {np.percentile(delays_ms, 99):.2f}ms")

        # 等待用户查看图表
        print("\nClose the graph window to exit...")
        try:
            input("Press Enter to close...")  # 等待用户输入再关闭
        except:
            pass

    finally:
        plotter.stop()
        client.close()
        print("Client closed")
