from smbus2 import SMBus
import time
import signal
import sys

# ======================
# Ctrl+C 安全退出
# ======================
def signal_handler(sig, frame):
    print("\n停止电机...")
    stop()
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)

# ======================
# PCA9685 初始化
# ======================
bus = SMBus(1)
addr = 0x60

bus.write_byte_data(addr, 0x00, 0x00)
bus.write_byte_data(addr, 0xFE, 60)

# ======================
# 电机定义
# ======================
MOTOR = {
    "M1": (0, 1),
    "M2": (2, 3),
    "M3": (4, 5),
    "M4": (6, 7),
}

DIR = {
    "M1": 1,
    "M2": 1,
    "M3": 1,
    "M4": 1,
}

# ======================
# PWM输出
# ======================
def set_pwm(ch, value):
    value = max(0, min(4095, value))
    bus.write_byte_data(addr, 0x06 + 4 * ch, 0)
    bus.write_byte_data(addr, 0x07 + 4 * ch, 0)
    bus.write_byte_data(addr, 0x08 + 4 * ch, value & 0xFF)
    bus.write_byte_data(addr, 0x09 + 4 * ch, value >> 8)

# ======================
# 单电机控制
# ======================
def motor_drive(name, speed):
    in1, in2 = MOTOR[name]

    speed *= DIR[name]
    pwm = int(abs(speed) * 4095)

    if speed > 0:
        set_pwm(in1, pwm)
        set_pwm(in2, 0)
    elif speed < 0:
        set_pwm(in1, 0)
        set_pwm(in2, pwm)
    else:
        # 刹车
        set_pwm(in1, 4095)
        set_pwm(in2, 4095)

# ======================
# 基础移动
# ======================
def forward(speed=0.6):
    motor_drive("M1", speed)
    motor_drive("M2", speed)
    motor_drive("M3", -speed)
    motor_drive("M4", -speed)

def backward(speed=0.6):
    motor_drive("M1", -speed)
    motor_drive("M2", -speed)
    motor_drive("M3", speed)
    motor_drive("M4", speed)

def move_left(speed=0.6):
    motor_drive("M1", -speed)
    motor_drive("M2", speed)
    motor_drive("M3", speed)
    motor_drive("M4", -speed)

def move_right(speed=0.6):
    motor_drive("M1", speed)
    motor_drive("M2", -speed)
    motor_drive("M3", -speed)
    motor_drive("M4", speed)

# ======================
# ⭐ 旋转
# ======================
def rotate_left(speed=0.6):
    motor_drive("M1", -speed)
    motor_drive("M2", -speed)
    motor_drive("M3", -speed)
    motor_drive("M4", -speed)

def rotate_right(speed=0.6):
    motor_drive("M1", speed)
    motor_drive("M2", speed)
    motor_drive("M3", speed)
    motor_drive("M4", speed)

# ======================
# ⭐ 斜向移动（新增）
# ======================
def move_left_forward(speed=0.6):
    motor_drive("M1", 0)
    motor_drive("M2", speed)
    motor_drive("M3", 0)
    motor_drive("M4", -speed)

def move_right_forward(speed=0.6):
    motor_drive("M1", speed)
    motor_drive("M2", 0)
    motor_drive("M3", -speed)
    motor_drive("M4", 0)

def move_left_backward(speed=0.6):
    motor_drive("M1", -speed)
    motor_drive("M2", 0)
    motor_drive("M3", speed)
    motor_drive("M4", 0)

def move_right_backward(speed=0.6):
    motor_drive("M1", 0)
    motor_drive("M2", -speed)
    motor_drive("M3", 0)
    motor_drive("M4", speed)

# ======================
# 停止
# ======================
def stop():
    motor_drive("M1", 0)
    motor_drive("M2", 0)
    motor_drive("M3", 0)
    motor_drive("M4", 0)

# ======================
# 测试
# ======================
if __name__ == "__main__":

    actions = [
        ("前进", forward),
        ("后退", backward),
        ("左移", move_left),
        ("右移", move_right),
        ("左前", move_left_forward),
        ("右前", move_right_forward),
        ("左后", move_left_backward),
        ("右后", move_right_backward),
        ("左旋", rotate_left),
        ("右旋", rotate_right),
    ]

    for name, func in actions:
        print(name)
        func()
        time.sleep(1)
        stop()
        time.sleep(2)

    print("完成")