"""
Raspberry Pi 5 + M5Stack StepMotor Driver 模块驱动
=====================================================

本模块用于驱动 M5Stack StepMotor Driver 模块，控制最多三个步进电机，
驱动圆柱在水槽中直线移动。

本模块结合了两种控制方式：
1. I2C 通信（地址 0x27）：用于控制模块的使能/复位/配置
2. 硬件 PWM（GPIO12）：用于产生步进脉冲，通过调整脉冲频率控制电机速度

方向控制通过独立的 GPIO 引脚（5/6/16）实现，三个电机共享同一路 PWM
但方向可以独立控制。

PWM 概念简介：
  PWM（Pulse Width Modulation，脉冲宽度调制）是一种通过快速开关信号来
  模拟不同电压输出的技术。在步进电机控制中，PWM 的频率决定电机转速——
  频率越高，脉冲越快，电机转动越快。

参考文档：
  https://docs.m5stack.com/en/module/Stepmotor%20Driver%20Module13.2%20v1.1

依赖库：smbus2（I2C 通信库）

使用示例：
    from drv_stepmotor import StepMotorDriver
    with StepMotorDriver() as driver:
        driver.set_speeds([0.02, -0.015, 0.01])
        time.sleep(2)

"""
from __future__ import annotations  # 启用 Python 未来的注解特性，允许在类型提示中使用字符串形式的类名（前向引用）

import math  # 导入 math 模块，用于数学计算（如余弦函数 cos() 用于 S 曲线加减速、圆周率 pi）
import threading  # 导入 threading 模块，用于多线程操作（本模块预留，当前未直接使用）
import time  # 导入 time 模块，用于延时和计时（如加减速过程中的 sleep 暂停）
from typing import List, Optional  # 从 typing 模块导入 List（列表类型）和 Optional（可选类型）用于参数类型注解

from gpiozero import DigitalOutputDevice  # 从 gpiozero 库导入 DigitalOutputDevice 类，用于控制树莓派 GPIO 引脚的电平输出（高/低）
from smbus2 import SMBus  # 从 smbus2 库导入 SMBus 类，用于通过 I2C 总线与 M5Stack 步进电机模块通信
from rpi_hardware_pwm import HardwarePWM  # 从 rpi_hardware_pwm 库导入 HardwarePWM 类，用于控制树莓派硬件 PWM 发生器

# ==================== I2C 寄存器地址常量 ====================
# 以下常量对应 M5Stack StepMotor Driver 模块内部寄存器的 I2C 地址，
# 参考自官方 C++ 实现文件 ref_stepmotor.cpp
MODULE_STEPMOTOR_ADDR = 0x27  # 模块的 I2C 从设备地址（7 位地址 0x27），所有 I2C 通信都发往此地址
STEPMOTOR_REG_IO_CFG = 0x03  # IO 配置寄存器地址，用于配置模块引脚的输入/输出功能
STEPMOTOR_REG_MICROSTEP = 0x01  # 微步进模式配置寄存器地址，设置电机的微步细分数
STEPMOTOR_REG_ENABLE = 0x01  # 使能控制寄存器地址（与微步进共用 0x01，通过不同位控制不同功能）
STEPMOTOR_REG_EXTIO = 0x00  # 外部 IO 寄存器地址，用于扩展输入输出
STEPMOTOR_REG_FAULT = 0x04  # 故障状态寄存器地址，读取过流/过热等错误状态
STEPMOTOR_REG_RESET = 0x05  # 复位寄存器地址，用于单独复位各个电机通道


class StepMotorError(RuntimeError):
    """自定义异常类，继承自 RuntimeError，用于步进电机驱动相关的错误处理"""
    pass  # 类体为空，仅作为自定义异常类型的标识，不需要额外实现


class StepMotorDriver:
    """步进电机驱动器类，最多控制 3 个步进电机。

    本类封装了 I2C 通信和硬件 PWM 的全部控制逻辑，提供启动、停止、调速、
    方向控制等完整功能。三个步进电机共享同一路硬件 PWM（产生步进脉冲），
    但每个电机的方向由独立的 GPIO 引脚控制。

    重要设计说明：
    - 三个电机共享同一路 PWM：因为所有电机同时以相同频率步进，仅方向可以不同
    - I2C 只用于模块的使能/复位配置，不用于产生步进脉冲
    - 步进脉冲由树莓派硬件 PWM 直接产生，不受 CPU 负载影响

    构造函数参数说明：
    - pwm_pins: 用于产生步进脉冲的 GPIO 引脚编号（BCM 编号），默认为 GPIO12
    - dir_pins: 用于控制方向信号的 GPIO 引脚编号列表，最多 3 个，默认 [5, 6, 16]
    - steps_per_rev: 电机转一圈所需的步数（微步进），默认 200×32×14 = 89600
    - mm_per_rev: 电机转一圈对应的直线移动距离（毫米），默认 π×20 ≈ 62.83 mm
    - ramp_time: 加减速的斜坡时间（秒），默认 20 秒
    """

    def __init__(
        self,
        pwm_pins: int = None,  # PWM 引脚编号（BCM 编号），默认为 12（即 GPIO12）
        dir_pins: Optional[List[int]] = None,  # 方向控制 GPIO 引脚编号列表，默认为 [5, 6, 16]
        steps_per_rev: int = 200*32*14,  # 电机转一圈的总步数（微步进后）
        mm_per_rev: float = math.pi*20,  # 电机转一圈对应的直线移动距离（毫米）
        ramp_time: float = 20,  # 加减速斜坡时间（秒）
    ) -> None:
        # 保存所有配置参数为实例变量，供其他方法访问
        self.pwm_pins = pwm_pins or 12  # 如果未传入 PWM 引脚编号，默认使用 GPIO12
        self.dir_pins = dir_pins or [5, 6, 16]  # 如果未传入方向引脚列表，默认使用 GPIO5、GPIO6、GPIO16
        self.steps_per_rev = steps_per_rev  # 每转总步数 = 200（电机物理步数）× 32（微步数）× 14（减速比）
        self.mm_per_rev = mm_per_rev  # 每转移动距离 = π × 20（驱动轮直径，单位 mm）
        self.max_speed = 0.1  # 最大速度限制（米/秒），防止电机速度过快导致失步或机械冲击
        self.ramp_time = float(ramp_time)  # 加减速斜坡时间（秒），强制转为 float 类型以确保精度
        self.i2c_addr = MODULE_STEPMOTOR_ADDR  # I2C 设备地址，使用预定义的常量 0x27
        self.i2c_bus = SMBus(1)  # 创建 I2C 总线对象，使用 I2C 总线 1（树莓派 5 上 I2C-1 对应引脚 GPIO2/SDA、GPIO3/SCL）

        self.default_dir = [True, True, False]  # 三个电机的默认方向：电机 0 和电机 1 正向，电机 2 反向
        self._pwm = HardwarePWM(pwm_channel=0, hz=1, chip=0)  # 创建硬件 PWM 对象：通道 0，初始 1Hz，PWM 芯片 0
        self._dir_pins = [DigitalOutputDevice(d, initial_value=False) for d in self.dir_pins]  # 为每个方向引脚创建 GPIO 输出对象，初始值为低电平
        self.current_speed = 0.0  # 当前速度（米/秒），初始化为 0 表示电机处于停止状态

        self.set_dir(True)  # 设置三个电机的初始方向为正向


    def _write_byte(self, reg: int, value: int) -> None:
        """向模块的指定 I2C 寄存器写入一个字节的数据。

        这是一个内部方法（以下划线开头），封装了 I2C 寄存器写入操作。

        参数:
            reg: 寄存器地址（如 0x01 使能寄存器、0x05 复位寄存器等）
            value: 要写入的 8 位数据值（0~255），写入前会自动截断为低 8 位

        异常:
            StepMotorError: 如果 I2C 通信失败则抛出此异常
        """
        value &= 0xFF  # 截断高位：只保留低 8 位（0~255），确保写入一个字节
        try:
            self.i2c_bus.write_byte_data(self.i2c_addr, reg, value)  # 通过 I2C 总线向指定寄存器的地址写入数据
        except Exception as e:
            raise StepMotorError(f"I2C write_byte failed: {e}")  # 如果 I2C 操作异常，包装为自定义异常向上抛出

    def _read_byte(self, reg: int) -> int:
        """读取模块指定 I2C 寄存器的一个字节数据。

        这是一个内部方法（以下划线开头），封装了 I2C 寄存器读取操作。

        参数:
            reg: 寄存器地址

        返回:
            int: 读取到的 8 位数据值（0~255）

        异常:
            StepMotorError: 如果 I2C 通信失败则抛出此异常
        """
        try:
            self.i2c_bus.write_byte(self.i2c_addr, reg)  # 先向 I2C 设备发送要读取的寄存器地址（写入操作选择寄存器）
            data = self.i2c_bus.read_byte(self.i2c_addr)  # 再从设备读取该寄存器的当前值（读取操作获取数据）
            return data  # 返回读取到的数据值
        except Exception as e:
            raise StepMotorError(f"I2C read_byte failed: {e}")  # 如果 I2C 操作异常，包装为自定义异常向上抛出
        
    def enable_motor(self):
        """使能电机：允许电机驱动芯片接收脉冲信号，电机可以正常转动。

        通过修改使能寄存器（地址 0x01）的第 4 位（bit 4）来实现：
        bit 4 = 0 表示使能，bit 4 = 1 表示禁止。
        操作流程：先读取寄存器的当前值 → 修改特定位 → 将新值写回寄存器。
        这种"读-改-写"模式确保不影响寄存器中其他位的设置。
        """
        reg_data = self._read_byte(STEPMOTOR_REG_ENABLE)  # 读取使能寄存器（地址 0x01）的当前值
        reg_data &= 0xEF  # 清除第 4 位（bit 4 置 0）：0xEF = 0b11101111，bit 4 = 0 表示使能电机
        self._write_byte(STEPMOTOR_REG_ENABLE, reg_data)  # 将修改后的值写回使能寄存器（地址 0x01）

    def disable_motor(self):
        """禁止电机：阻止电机驱动芯片接收脉冲信号，电机停止转动。

        通过修改使能寄存器（地址 0x01）的第 4 位（bit 4）来实现：
        将其他位清零、仅保留 bit 4 的值。
        注意：原代码使用 reg_data &= 0x10 仅保留 bit 4 并清除其他位，
        这与标准的"将 bit 4 置 1"（reg_data |= 0x10）有所不同，
        但为保持逻辑不变，此处不修改代码。
        """
        reg_data = self._read_byte(STEPMOTOR_REG_ENABLE)  # 读取使能寄存器（地址 0x01）的当前值
        reg_data &= 0x10  # 将除第 4 位之外的所有位清零：0x10 = 0b00010000，仅保留 bit 4
        self._write_byte(STEPMOTOR_REG_ENABLE, reg_data)  # 将修改后的值写回使能寄存器（地址 0x01）

    def reset_motor(self, resmtr: int, enable: bool):
        """复位或释放指定的电机通道。

        通过复位寄存器（地址 0x05）的低 3 位分别控制三个电机：
        bit 0 控制电机 0，bit 1 控制电机 1，bit 2 控制电机 2。
        将某位置 1 表示复位该电机，置 0 表示释放。

        参数:
            resmtr: 电机索引（0、1、2），分别对应三个步进电机通道
            enable: True 表示复位（将对应位设为 1），False 表示释放（将对应位设为 0）
        """
        reg_data = self._read_byte(STEPMOTOR_REG_RESET)  # 读取复位寄存器（地址 0x05）的当前值
        reg_data &= 0x07  # 仅保留低 3 位（bit 0~2），分别对应三个电机的复位状态，高位清零
        if not enable:
            reg_data &= ~(0x01 << resmtr)  # 如果 enable=False，将对应电机的位清零：~(1<<n) 将第 n 位取反后做 AND
        else:
            reg_data |= (0x01 << resmtr)  # 如果 enable=True，将对应电机的位置 1：(1<<n) 左移 n 位得到只有第 n 位为 1 的掩码
        self._write_byte(STEPMOTOR_REG_RESET, reg_data)  # 将修改后的值写回复位寄存器（地址 0x05）

    def close(self) -> None:
        """关闭驱动器：停止所有电机、释放硬件资源。

        按顺序执行：紧急停止 → I2C 禁能 → 所有电机通道释放 → 清理 GPIO/PWM/I2C 对象。
        使用 try/except 确保即使某个清理步骤出错，后续步骤也能继续执行。
        """
        self.emergency_stop()  # 紧急停止：立即切断 PWM 输出，电机瞬间失去脉冲
        self.disable_motor()  # 通过 I2C 设置使能寄存器，禁止电机驱动芯片
        for idx in range(3):  # 循环三个电机通道（索引 0、1、2）
            self.reset_motor(idx, False)  # 释放每个电机通道：将复位寄存器中对应位清 0
        # 关闭 gpiozero 相关的 GPIO 设备
        try:
            self._pwm.close()  # 关闭硬件 PWM 对象，释放 PWM 通道资源
            self.i2c_bus.close()  # 关闭 I2C 总线连接，释放总线资源
            for d in self._dir_pins:  # 遍历所有方向控制引脚对象
                d.off()  # 将方向 GPIO 引脚置为低电平（关闭输出）
                d.close()  # 关闭方向 GPIO 引脚对象，释放资源
        except Exception:
            pass  # 忽略清理过程中的任何异常，确保程序能正常退出

    def set_dir(self, dir: bool = True) -> None:
        """设置三个电机的整体方向。

        注意：三个电机只能同时设置为相同的方向（全部正向或全部反向）。
        每个电机的实际方向由 default_dir（电机默认方向）和传入的 dir 参数
        通过异或（XOR）运算决定，这样可以灵活设置每个电机的正反。

        参数:
            dir: True 表示正向，False 表示反向
                 电机 n 的实际方向 = default_dir[n] XOR dir

        异常:
            StepMotorError: 如果电机正在运行（current_speed ≠ 0）则不能改变方向
        """
        if self.current_speed == 0.0:  # 检查电机是否处于停止状态，运行时禁止改变方向
            for idx, dir_dev in enumerate(self._dir_pins):  # 遍历三个方向控制引脚（索引 0、1、2）
                boolean_dir = self.default_dir[idx] ^ dir  # 计算实际方向：默认方向 XOR 传入方向（异或运算）
                if boolean_dir:
                    dir_dev.on()  # 如果计算结果为 True，将 GPIO 引脚置为高电平（电机正转）
                else:
                    dir_dev.off()  # 如果计算结果为 False，将 GPIO 引脚置为低电平（电机反转）
        else:
            raise StepMotorError("Cannot change direction while motor is running")  # 电机运行时改变方向会损坏驱动器

    def start(self, speed_m_s: float) -> None:
        """启动电机，并使用 S 曲线加速到目标速度。

        S 曲线加速原理：
        使用余弦函数产生平滑的"S"形加速曲线，速度变化率（即加速度）
        在启动和结束时较小，中间较大，从而减少机械冲击和电机失步风险。

        S 曲线加速公式：
            s = (cos(π × (1 - i/100)) + 1) / 2 × speed_m_s
        其中 i 从 0 到 99 逐步增加：
        - i=0  时：cos(π×(1-0)) = cos(π) = -1，s = (-1+1)/2 × speed = 0 × speed = 0（起始速度为 0）
        - i=50 时：cos(π×(1-0.5)) = cos(π/2) = 0，s = (0+1)/2 × speed = 0.5 × speed（中间速度为目标一半）
        - i=99 时：cos(π×(1-0.99)) = cos(0.01π) ≈ 1，s = (1+1)/2 × speed ≈ speed（接近目标速度）
        整个过程中，加速度先增大后减小，形成"S"形曲线，相比线性加速更加平顺。

        频率换算公式：
            freq = s × steps_per_rev / (mm_per_rev / 1000)
        推导过程：
            - freq(Hz) = 每秒步数（脉冲数）
            - 每秒步数 = 速度(m/s) × 步数/米
            - 步数/米 = 步数/转 ÷ 米/转 = steps_per_rev ÷ (mm_per_rev / 1000)
            - steps_per_rev = 200（电机物理步数，1.8°/步）× 32（驱动器微步细分数）× 14（行星减速器减速比）
            - mm_per_rev = π × 20（驱动轮直径 20mm）= 62.83 mm/转

        参数:
            speed_m_s: 目标速度（米/秒），必须为正值

        异常:
            StepMotorError: 如果电机已经在运行中，则不能重复启动
        """
        if self.current_speed == 0.0:  # 仅在电机当前处于停止状态时才能启动
            self._pwm.change_frequency(1)  # 先将 PWM 频率设为最低的 1Hz，避免从 0Hz 直接跳变
            self._pwm.start(100)  # 启动硬件 PWM 输出，占空比设为 100%（模块内部处理步进脉冲的生成）
            for i in range(100):  # 将加速过程分为 100 个步进点，从 i=0 到 i=99 逐步加速
                # S 曲线加速公式：余弦函数值从 cos(π) = -1 逐渐变化到 cos(0.01π) ≈ 1
                # 经 (cos + 1) / 2 映射后，系数从 0 逐渐变化到接近 1
                s = (math.cos(math.pi * (1-i/100)) + 1) / 2 * speed_m_s  # 计算当前步进点的瞬时目标速度（米/秒）
                # 将速度（米/秒）换算为 PWM 频率（Hz）：
                # freq = s × steps_per_rev / (mm_per_rev / 1000)
                # 除以 1000 因为 mm_per_rev 是毫米，需要转换为米
                freq = max(s * self.steps_per_rev / (self.mm_per_rev / 1000), 1)  # 换算频率，max(...,1) 确保频率不低于 1Hz
                self._pwm.change_frequency(freq)  # 更新硬件 PWM 的频率，使电机以对应速度转动
                time.sleep(self.ramp_time / 100)  # 等待 ramp_time/100 秒，控制每步加速的持续时间
            self.current_speed = speed_m_s  # 加速完成后，将当前速度记录为目标速度
        else:
            raise StepMotorError("Motor is already running")  # 如果电机已经运行，拒绝重复启动（必须先 stop）

    def stop(self) -> None:
        """停止电机，使用 S 曲线减速到零。

        S 曲线减速公式：
            s = (cos(π × i/100) + 1) / 2 × current_speed
        其中 i 从 0 到 99 逐步增加：
        - i=0  时：cos(0) = 1，s = (1+1)/2 × speed = 1 × speed（保持当前速度开始减速）
        - i=50 时：cos(π/2) = 0，s = (0+1)/2 × speed = 0.5 × speed（速度降到一半）
        - i=99 时：cos(0.99π) ≈ -1，s = (-1+1)/2 × speed = 0 × speed = 0（减速到零）

        与 start() 中的加速曲线对称，减速曲线也是 S 形，确保停止过程平滑。
        减速时间比加速时间短（使用 ramp_time/4 而不是 ramp_time/100），
        因为减速到零后电机即停止，不需要像加速那样缓慢进入。
        """
        if self.current_speed != 0.0:  # 只有在电机正在运行时才能执行停止操作
            for i in range(100):  # 将减速过程分为 100 个步进点，从 i=0 到 i=99 逐步减速
                # S 曲线减速公式：余弦值从 cos(0) = 1 逐渐变化到 cos(0.99π) ≈ -1
                # 经 (cos + 1) / 2 映射后，系数从 1 逐渐下降到接近 0
                s = (math.cos(math.pi * i/100) + 1) / 2 * self.current_speed  # 计算当前步进点的瞬时速度（米/秒）
                freq = max(s * self.steps_per_rev / (self.mm_per_rev / 1000), 1)  # 将速度换算为 PWM 频率，最小 1Hz
                self._pwm.change_frequency(freq)  # 更新硬件 PWM 频率，降低电机转速
                time.sleep(self.ramp_time / 4 / 100)  # 等待 ramp_time/400 秒，比加速步进快 4 倍
            self._pwm.stop()  # 减速完成后，完全停止硬件 PWM 输出
            self._pwm.change_frequency(1)  # 将 PWM 频率复位到 1Hz 默认值，为下一次启动做准备
            self.current_speed = 0.0  # 将当前速度设为 0，标记电机已停止
        else:
            raise StepMotorError("Motor is already stopped")  # 如果电机已经停止，拒绝重复执行停止操作
        
    def change_speed(self, speed_m_s: float) -> None:
        """在电机运行过程中，平滑地将速度切换到新目标值。

        速度切换 S 曲线公式：
            s = (cos(π × i/100) + 1) / 2 × (current_speed - speed_m_s) + speed_m_s
        公式推导：
            将减速公式推广到任意目标速度：
            - 令 delta = (current_speed - speed_m_s)，即当前速度与目标速度的差值
            - 减速公式变为：s = 余弦因子 × delta + speed_m_s
            - i=0  时：因子=1，s = 1×delta + speed_m_s = current_speed（保持当前速度）
            - i=99 时：因子≈0，s = 0×delta + speed_m_s = speed_m_s（达到目标速度）
        传入的速度可以是加速（目标 > 当前）或减速（目标 < 当前），
        S 曲线都能平滑过渡。

        参数:
            speed_m_s: 新的目标速度（米/秒），可高于或低于当前速度

        异常:
            StepMotorError: 如果电机未运行，调用此方法会抛出异常，提示使用 start()
        """
        if self.current_speed != 0.0:  # 只有在电机正在运行时才能切换速度
            for i in range(100):  # 分 100 步进行速度过渡，从 i=0 到 i=99
                # 速度切换公式：当前的 S 曲线因子从 1 过渡到 0，
                # 乘以速度差后加到目标速度上，实现从当前速度到目标速度的平滑过渡
                s = (math.cos(math.pi * i/100) + 1) / 2 * (self.current_speed - speed_m_s) + speed_m_s
                freq = max(s * self.steps_per_rev / (self.mm_per_rev / 1000), 1)  # 将速度换算为 PWM 频率，最小 1Hz
                self._pwm.change_frequency(freq)  # 更新硬件 PWM 频率，改变电机转速
                time.sleep(self.ramp_time / 4 / 100)  # 等待 ramp_time/400 秒，控制每步过渡的持续时间
            self.current_speed = speed_m_s  # 速度过渡完成后，将当前速度更新为新目标速度
        else:
            raise StepMotorError("Motor is not running. Use start() to start the motor.")  # 电机停止时不能用 change_speed

    def emergency_stop(self) -> None:
        """紧急停止：立即切断 PWM 输出，不经过 S 曲线减速。

        与 stop() 方法的区别：
        - stop()：使用 S 曲线平滑减速（约需 ramp_time/4 秒）
        - emergency_stop()：直接切断 PWM 输出，无减速过程（立即停止）

        适用于紧急情况下的快速停机，但会产生较大的机械冲击。
        """
        self._pwm.stop()  # 立即停止硬件 PWM 输出，电机瞬间失去步进脉冲
        self._pwm.change_frequency(1)  # 将 PWM 频率复位到 1Hz 默认值，为下一次启动做准备
        self.current_speed = 0.0  # 将当前速度设为 0，标记电机已停止
    
    def __enter__(self):
        """进入上下文管理器：当使用 'with StepMotorDriver() as driver:' 语法时自动调用

        返回:
            StepMotorDriver: 返回自身实例，供 with 语句块中的变量引用
        """
        return self  # 返回当前 StepMotorDriver 实例
    
    def __exit__(self, exc_type, exc, tb):
        """退出上下文管理器：当 with 语句块执行完毕时自动调用

        参数:
            exc_type: 异常类型（如果语句块中发生了异常）
            exc: 异常对象
            tb: 异常回溯信息
        """
        self.close()  # 自动调用 close() 清理所有资源，确保即使发生异常也能释放硬件