import rp2
from machine import Pin
import time

print("hello")

class StepperPIO:
    """
    Stepper motor driver using the Raspberry Pi Pico's PIO.
    Offloads step generation to PIO for precise timing.
    """

    @rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_RIGHT, autopull=True)
    def stepper_pio():
        pull()            # Get num_steps from FIFO
        mov(x, osr)       # Store num_steps in X

        pull()            # Get direction bit (1 or 0)
        out(pins, 1)      # Set the DIR pin before stepping

        label("loop")
        jmp(x_dec, "do_step")  # Loop until X = 0
        jmp("done")            # Exit when done

        label("do_step")
        nop().side(1)     # STEP pulse HIGH
        nop().side(0)     # STEP pulse LOW

        


        jmp("loop")       # Repeat for num_steps

        label("done")
        nop()             # Do nothing, waiting for next command

    def __init__(self, sm_id, step_pin, dir_pin, freq=5000):
        """
        Initialize the stepper motor PIO driver.

        :param sm_id: State Machine ID (0-3)
        :param step_pin: GPIO pin for STEP
        :param dir_pin: GPIO pin for DIR
        :param freq: Frequency of steps (Hz)
        """
        self.sm = rp2.StateMachine(
            sm_id, StepperPIO.stepper_pio,
            freq=freq, sideset_base=Pin(step_pin), out_base=Pin(dir_pin)
        )
        self.sm.active(1)  # Start the state machine

    def move(self, num_steps, direction):
        """
        Move the stepper motor a specified number of steps.

        :param num_steps: Number of steps to move
        :param direction: 1 for CW, 0 for CCW
        """
        self.sm.put(num_steps)  # Send step count FIRST
        self.sm.put(direction)  # Send direction AFTER

    def stop(self):
        """ Stop the state machine. """
        self.sm.active(0)


print("hello")
# Example Usage
stepper = StepperPIO(sm_id=0, step_pin=0, dir_pin=1, freq=2000)
stepper.move(200, 1)  # Move 20000 steps clockwise
stepper.move(200, 0)   # Move 2000 steps counterclockwise
stepper.move(200, 1)   # Move 2000 steps counterclockwise
stepper.move(200, 0)   # Move 2000 steps counterclockwise


print("done")
while True:
    print("ok")
    time.sleep(1)