from pyb import Pin, Timer

class Motor:
    def __init__(self, pwm_range):
        # 左轮A， 右轮B
        self.io_p_A = Pin("Y5", Pin.OUT_PP)  # 电机A的正极IO口
        self.io_n_A = Pin("Y6", Pin.OUT_PP)
        self.io_p_B = Pin("X11", Pin.OUT_PP)  # 电机B的正极IO口
        self.io_n_B = Pin("X12", Pin.OUT_PP)
        self.io_PWM_A = Pin("Y3", Pin.OUT_PP)   # 电机A、B的PWM输入IO口
        self.io_PWM_B = Pin("Y4", Pin.OUT_PP)
        timer_M = Timer(4, freq=1000)    # PWM输出定时器
        self.ch_PWM_A = timer_M.channel(3, Timer.PWM, pin=self.io_PWM_A)
        self.ch_PWM_B = timer_M.channel(4, Timer.PWM, pin=self.io_PWM_B)
        self.pwm_range = pwm_range  # PWM范围

    def motor_run(self, pwm_A, pwm_B):
        '''
        函数功能：电机驱动
        入口参数：两个电机的PWM值（范围为pwm_range）
        '''
        self.ch_PWM_A.pulse_width_percent(100*(pwm_A/self.pwm_range)) # 给电机PWM占空比，默认pwm_range为100
        self.ch_PWM_B.pulse_width_percent(100*(pwm_B/self.pwm_range))

    def direction_control(self, pwm_A, pwm_B):
        '''
        函数功能：控制方向，并将PWM值转换为正值
        入口参数：两个电机的PWM值（范围为pwm_range）
        返回值 ：两个电机的正PWM值
        '''
        def ahead_m_A():
            self.io_p_A.high(); self.io_n_A.low();
            # self.io_p_A.low(); self.io_n_A.high();
        def back_m_A():
            self.io_p_A.low(); self.io_n_A.high();
            # self.io_p_A.high(); self.io_n_A.low();
        def ahead_m_B():
            self.io_p_B.high(); self.io_n_B.low();
            # self.io_p_B.low(); self.io_n_B.high();
        def back_m_B():
            self.io_p_B.low(); self.io_n_B.high();
            # self.io_p_B.high(); self.io_n_B.low();
        def stop():
            self.io_p_A.low(); self.io_n_A.low();
            self.io_p_B.low(); self.io_n_B.low();
        
        if (pwm_A > 0 and pwm_B >= 0) or (pwm_A >= 0 and pwm_B > 0): # 向前
            ahead_m_A()
            ahead_m_B()
            # print("ahead")
        elif (pwm_A < 0 and pwm_B <= 0) or (pwm_A <= 0 and pwm_B < 0): # 向后
            back_m_A() 
            back_m_B()
            # print("back")
        elif pwm_A < 0 and pwm_B > 0: # 原地左转，只有当没有其他动作时才会原地左转
            back_m_A()
            ahead_m_B()
            # print("left")
        elif pwm_A > 0 and pwm_B < 0: # 原地右转，只有当没有其他动作时才会原地右转
            ahead_m_A()
            back_m_B()
            # print("right")
        else:
            stop()
            # print("stop")

        pwm_A = abs(pwm_A)
        pwm_B = abs(pwm_B)
        return pwm_A, pwm_B