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),   # 前左（Y轴）
    "M2": (2, 3),   # 前右（Y轴）
    "M3": (4, 5),   # 后左（X轴）
    "M4": (6, 7),   # 后右（X轴）
}

# ⚠️ 需要你微调方向（如果反了就改 ±1）
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)

# ======================
# ⭐ 全向轮核心控制（重点）
# vx: 左右 (-1~1)
# vy: 前后 (-1~1)
# w : 旋转 (-1~1)
# ======================
def move(vx, vy, w):
    m1 = vy + w
    m2 = vy - w
    m3 = vx + w
    m4 = vx - w

    # 归一化（防止超限）
    max_val = max(abs(m1), abs(m2), abs(m3), abs(m4), 1)

    motor_drive("M1", m1 / max_val)
    motor_drive("M2", m2 / max_val)
    motor_drive("M3", m3 / max_val)
    motor_drive("M4", m4 / max_val)

# ======================
# 基础动作（封装）
# ======================
def forward(speed=0.6):
    move(0, speed, 0)

def backward(speed=0.6):
    move(0, -speed, 0)

def left(speed=0.6):
    move(-speed, 0, 0)

def right(speed=0.6):
    move(speed, 0, 0)

def rotate_left(speed=0.6):
    move(0, 0, speed)

def rotate_right(speed=0.6):
    move(0, 0, -speed)

def stop():
    motor_drive("M1", 0)
    motor_drive("M2", 0)
    motor_drive("M3", 0)
    motor_drive("M4", 0)

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

    print("前进")
    forward()
    time.sleep(1)

    print("后退")
    backward()
    time.sleep(1)

    print("左移")
    left()
    time.sleep(1)

    print("右移")
    right()
    time.sleep(1)

    print("左旋转")
    rotate_left()
    time.sleep(1)

    print("停止")
    stop()