import serial
import struct
import time
from enum import Enum

class ChecksumType(Enum):
    """校验方式枚举"""
    FIXED_0x6B = 0x6B  # 固定校验字节
    XOR = 1            # XOR校验
    CRC8 = 2           # CRC-8校验（需额外实现）

class Emm42Driver:
    """Emm42_V5.0闭环步进电机驱动库（串口控制）"""
    def __init__(self, port, baudrate=115200, timeout=0.1, device_id=1, checksum_type=ChecksumType.FIXED_0x6B):
        """
        初始化串口通信
        :param port: 串口号（如"COM3"或"/dev/ttyUSB0"）
        :param baudrate: 波特率（默认115200）
        :param timeout: 读取超时时间（秒）
        :param device_id: 电机地址（1-255，默认1）
        :param checksum_type: 校验方式（默认0x6B固定校验）
        """
        self.ser = serial.Serial(port, baudrate, timeout=timeout)
        self.device_id = device_id
        self.checksum_type = checksum_type
        if not self.ser.is_open:
            self.ser.open()

    def close(self):
        """关闭串口"""
        if self.ser.is_open:
            self.ser.close()

    def _calculate_checksum(self, data):
        """计算校验和"""
        if self.checksum_type == ChecksumType.FIXED_0x6B:
            return 0x6B
        elif self.checksum_type == ChecksumType.XOR:
            checksum = 0
            for b in data:
                checksum ^= b
            return checksum & 0xFF
        elif self.checksum_type == ChecksumType.CRC8:
            # 如需使用CRC8，需补充CRC8计算逻辑
            raise NotImplementedError("CRC8校验未实现")
        return 0x6B

    def _send_command(self, func_code, params=None):
        """
        发送命令帧
        :param func_code: 功能码（参考协议命令列表）
        :param params: 参数列表（字节数组）
        :return: 发送的字节串
        """
        params = params or []
        # 构建命令帧：地址 + 功能码 + 参数 + 校验和
        frame = [self.device_id, func_code] + params
        checksum = self._calculate_checksum(frame)
        frame.append(checksum)
        # 发送数据
        self.ser.write(bytearray(frame))
        return bytearray(frame)

    def _read_response(self, expected_func=None, timeout=0.1):
        """
        读取电机返回数据
        :param expected_func: 期望的功能码（用于验证）
        :param timeout: 超时时间（秒）
        :return: 解析后的响应字典（含地址、功能码、参数），失败返回None
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            if self.ser.in_waiting < 3:  # 最小响应长度：地址(1) + 功能码(1) + 校验(1)
                time.sleep(0.001)
                continue
            # 读取地址和功能码
            addr = self.ser.read(1)[0]
            func = self.ser.read(1)[0]
            # 验证地址和功能码
            if addr != self.device_id:
                continue
            if expected_func is not None and func != expected_func:
                continue
            # 读取参数和校验和（参数长度需根据具体命令判断，此处简化处理）
            params = []
            while self.ser.in_waiting > 1:  # 保留最后1字节作为校验和
                params.append(self.ser.read(1)[0])
            checksum = self.ser.read(1)[0]
            # 验证校验和
            verify_data = [addr, func] + params
            if self._calculate_checksum(verify_data) != checksum:
                continue
            return {
                "addr": addr,
                "func": func,
                "params": bytes(params)
            }
        return None

    # ------------------------------ 基础控制指令 ------------------------------
    def enable_motor(self, enable=True, sync=False):
        """
        电机使能控制
        :param enable: True=使能，False=禁用
        :param sync: 多机同步标志（False=不启用，True=启用）
        :return: 命令状态（0x02=成功，其他=失败）
        """
        # 命令格式：地址 + 0xF3 + 0xAB + 使能状态 + 同步标志 + 校验
        params = [0xAB, 0x01 if enable else 0x00, 0x01 if sync else 0x00]
        self._send_command(0xF3, params)
        resp = self._read_response(expected_func=0xF3)
        return resp["params"][0] if resp else None

    def stop_immediately(self, sync=False):
        """立即停止电机"""
        # 命令格式：地址 + 0xFE + 0x98 + 同步标志 + 校验
        params = [0x98, 0x01 if sync else 0x00]
        self._send_command(0xFE, params)
        resp = self._read_response(expected_func=0xFE)
        return resp["params"][0] if resp else None

    # ------------------------------ 速度模式控制 ------------------------------
    def set_speed_mode(self, direction, speed_rpm, acceleration=0, sync=False):
        """
        速度模式控制
        :param direction: 方向（0=CW顺时针，1=CCW逆时针）
        :param speed_rpm: 转速（0.1-3000 RPM，S_Vel_IS使能时支持0.1精度）
        :param acceleration: 加速度档位（0-255，0=无加减速）
        :param sync: 多机同步标志
        :return: 命令状态（0x02=成功）
        """
        # 速度转换为16位无符号整数（0x0000-0xFFFF）
        speed_hex = struct.pack(">H", int(speed_rpm * 10)) if acceleration else struct.pack(">H", int(speed_rpm))
        params = [
            direction,
            speed_hex[0], speed_hex[1],  # 速度高低位
            acceleration & 0xFF,         # 加速度档位
            0x01 if sync else 0x00       # 同步标志
        ]
        self._send_command(0xF6, params)
        resp = self._read_response(expected_func=0xF6)
        return resp["params"][0] if resp else None

    # ------------------------------ 位置模式控制 ------------------------------
    def set_position_mode(self, direction, speed_rpm, pulses, is_absolute=False, acceleration=0, sync=False):
        """
        位置模式控制
        :param direction: 方向（0=CW，1=CCW）
        :param speed_rpm: 转速（0.1-3000 RPM）
        :param pulses: 脉冲数（根据细分计算，16细分下3200脉冲=1圈）
        :param is_absolute: 是否绝对位置（False=相对，True=绝对）
        :param acceleration: 加速度档位（0-255）
        :param sync: 多机同步标志
        :return: 命令状态（0x02=成功）
        """
        # 速度和脉冲数转换为16/32位无符号整数
        speed_hex = struct.pack(">H", int(speed_rpm))
        pulses_hex = struct.pack(">I", pulses)  # 脉冲数为32位
        params = [
            direction,
            speed_hex[0], speed_hex[1],        # 速度高低位
            acceleration & 0xFF,               # 加速度档位
            pulses_hex[0], pulses_hex[1], pulses_hex[2], pulses_hex[3],  # 脉冲数
            0x01 if is_absolute else 0x00,     # 绝对/相对标志
            0x01 if sync else 0x00             # 同步标志
        ]
        self._send_command(0xFD, params)
        resp = self._read_response(expected_func=0xFD)
        return resp["params"][0] if resp else None

    # ------------------------------ 原点回零指令 ------------------------------
    def trigger_homing(self, mode, sync=False):
        """
        触发回零操作
        :param mode: 回零模式（0=单圈就近，1=单圈方向，2=多圈无限位，3=多圈有限位）
        :param sync: 多机同步标志
        :return: 命令状态（0x02=成功）
        """
        # 命令格式：地址 + 0x9A + 回零模式 + 同步标志 + 校验
        params = [mode & 0xFF, 0x01 if sync else 0x00]
        self._send_command(0x9A, params)
        resp = self._read_response(expected_func=0x9A)
        return resp["params"][0] if resp else None

    # ------------------------------ 状态读取指令 ------------------------------
    def read_bus_voltage(self):
        """读取总线电压（mV）"""
        self._send_command(0x24)  # 功能码0x24对应读取总线电压
        resp = self._read_response(expected_func=0x24)
        if resp and len(resp["params"]) == 2:
            voltage = struct.unpack(">H", resp["params"])[0]  # 16位无符号整数
            return voltage / 1000.0  # 转换为V
        return None

    def read_realtime_speed(self):
        """读取实时转速（RPM）"""
        self._send_command(0x35)  # 功能码0x35对应读取转速
        resp = self._read_response(expected_func=0x35)
        if resp and len(resp["params"]) == 3:
            sign = -1 if resp["params"][0] == 0x01 else 1  # 0x01=负，0x00=正
            speed = struct.unpack(">H", resp["params"][1:3])[0]
            return sign * speed
        return None

    def read_position(self):
        """读取实时位置（脉冲数）"""
        self._send_command(0x36)  # 功能码0x36对应读取位置
        resp = self._read_response(expected_func=0x36)
        if resp and len(resp["params"]) == 5:
            sign = -1 if resp["params"][0] == 0x01 else 1
            position = struct.unpack(">I", resp["params"][1:5])[0]
            return sign * position
        return None

    # ------------------------------ 配置参数指令 ------------------------------
    def set_microstep(self, microstep, save=False):
        """
        设置细分（1-256）
        :param microstep: 细分值（0=256，1=1，2=2，...，255=255）
        :param save: 是否保存（True=掉电保存）
        :return: 命令状态（0x02=成功）
        """
        # 命令格式：地址 + 0x84 + 0x8A + 存储标志 + 细分值 + 校验
        params = [0x8A, 0x01 if save else 0x00, microstep & 0xFF]
        self._send_command(0x84, params)
        resp = self._read_response(expected_func=0x84)
        return resp["params"][0] if resp else None

    def restore_factory_settings(self):
        """恢复出厂设置（需重新上电校准）"""
        self._send_command(0x0F, [0x5F])  # 功能码0x0F + 参数0x5F
        resp = self._read_response(expected_func=0x0F)
        return resp["params"][0] if resp else None
    





#使用示例
if __name__ == "__main__":
    # 初始化电机驱动（替换为实际串口号）
    motor = Emm42Driver(port="COM3", device_id=1, baudrate=115200)

    try:
        # 使能电机
        if motor.enable_motor(enable=True) == 0x02:
            print("电机使能成功")

        # 速度模式控制：逆时针旋转，1000RPM，加速度档位5
        motor.set_speed_mode(direction=1, speed_rpm=1000, acceleration=5)
        time.sleep(2)  # 运行2秒

        # 立即停止
        motor.stop_immediately()
        time.sleep(1)

        # 位置模式控制：顺时针旋转，500RPM，相对位置3200脉冲（16细分下1圈）
        motor.set_position_mode(
            direction=0,
            speed_rpm=500,
            pulses=3200,
            is_absolute=False,
            acceleration=3
        )
        time.sleep(2)

        # 读取状态
        voltage = motor.read_bus_voltage()
        speed = motor.read_realtime_speed()
        position = motor.read_position()
        print(f"总线电压：{voltage}V，实时转速：{speed}RPM，当前位置：{position}脉冲")

        # 触发单圈就近回零
        motor.trigger_homing(mode=0)
        time.sleep(3)

    finally:
        # 禁用电机并关闭串口
        motor.enable_motor(enable=False)
        motor.close()