from machine import Pin
import time
import sys
import uselect  # We use uselect for better polling

class StepperControl:
    # UPDATED: Added m0, m1, m2 to __init__
    def __init__(self, step_pin, dir_pin, m0_pin, m1_pin, m2_pin, initialMicroStep):
        self.step_pin = Pin(step_pin, Pin.OUT)
        self.dir_pin = Pin(dir_pin, Pin.OUT)
        
        # Microstepping Control Pins
        self.m0 = Pin(m0_pin, Pin.OUT)
        self.m1 = Pin(m1_pin, Pin.OUT)
        self.m2 = Pin(m2_pin, Pin.OUT)
        
        # Default Parameters
        self.base_steps_per_rev = 200 # Native NEMA17 steps (1.8 deg)
        self.steps_per_rev = 200      # Current steps (changes with microstepping)
        self.current_step_pos = 0
        self.rpm = 60
        self.direction = 1  # 1 = CW, 0 = CCW
        
        # Hardware Init
        self.dir_pin.value(self.direction)
        self.step_pin.value(0)
        
        # Initialize at Full Step (1/1) by default
        self.set_microstep(initialMicroStep) 
        
        print(f"System Initialized. Default RPM: {self.rpm}")

    def _get_delay(self):
        """Calculates step delay. Adds safety for very high speeds."""
        if self.rpm <= 0:
            return 0.01
        # Delay logic relies on accurate self.steps_per_rev 
        delay = 60 / (self.rpm * self.steps_per_rev)
        return delay

    # --- NEW: Set Microstepping Resolution ---
    def set_microstep(self, mode):
        """
        Sets DRV8825 M0, M1, M2 pins based on requested resolution.
        Updates self.steps_per_rev so RPM math remains correct.
        """
        # Dictionary mapping: 'Mode' -> (M0, M1, M2, Multiplier)
        modes = {
            '1':    (0, 0, 0, 1),
            '1/2':  (1, 0, 0, 2),
            '1/4':  (0, 1, 0, 4),
            '1/8':  (1, 1, 0, 8),
            '1/16': (0, 0, 1, 16),
            '1/32': (1, 0, 1, 32)
        }
        
        if mode in modes:
            vals = modes[mode]
            # Set GPIO pins
            self.m0.value(vals[0])
            self.m1.value(vals[1])
            self.m2.value(vals[2])
            
            # Update step count (e.g., 200 * 32 = 6400 steps/rev)
            self.steps_per_rev = self.base_steps_per_rev * vals[3]
            print(f"Microstepping set to {mode}. Steps per Rev: {self.steps_per_rev}")
            print(f"Microstepp as angle{360/self.steps_per_rev}")
        else:
            print("Invalid mode. Options: 1, 1/2, 1/4, 1/8, 1/16, 1/32")

    # --- Req 1: Set RPM and Start Rotating ---
    def set_rpm_and_run(self):
        try:
            val = eval(input("Enter RPM to start: "))
            self.rpm = val
            print(f"RPM set to {self.rpm}. Starting Motor...")
            self.run_continuous() 
        except ValueError:
            print("Invalid number.")

    # --- Req 2: Go to Angle ---
    def go_to_angle(self):
        try:
            target_angle = eval(input("Enter Target Angle: "))
        except ValueError:
            print("Invalid Angle.")
            return

        print(f"Moving to {target_angle} degrees...")
        # Calculate steps based on CURRENT microstepping resolution
        target_step = int((target_angle / 360) * self.steps_per_rev)
        diff = target_step - self.current_step_pos
        
        if diff == 0:
            print("Already at target.")
            return

        # Temporary direction for this move
        move_dir = 1 if diff > 0 else 0
        self.dir_pin.value(move_dir)
        
        delay = self._get_delay()
        
        for _ in range(abs(diff)):
            self.step_pin.value(1)
            time.sleep(delay / 2)
            self.step_pin.value(0)
            time.sleep(delay)
            self.current_step_pos += (1 if move_dir else -1)
            
        print(f"Reached {target_angle} deg.")
        # Restore global direction
        self.dir_pin.value(self.direction)

    # --- Req 3: Toggle Direction Logic ---
    def toggle_direction_setting(self):
        self.direction = not self.direction
        self.dir_pin.value(self.direction)
        print(f"Dir: {'CW' if self.direction else 'CCW'}")

    # --- Req 4: Stop ---
    def stop(self):
        self.step_pin.value(0)
        print("Stopped.")

    # --- Req 5: Restart ---
    def restart(self):
        print("Restarting... Returning to 0.")
        
        # Return to 0
        target = 0
        diff = target - self.current_step_pos
        move_dir = 1 if diff > 0 else 0
        self.dir_pin.value(move_dir)
        
        delay = 0.005 
        for _ in range(abs(diff)):
            self.step_pin.value(1)
            time.sleep(delay / 2)
            self.step_pin.value(0)
            time.sleep(delay)
            
        self.current_step_pos = 0
        self.rpm = 60
        self.direction = 1
        self.dir_pin.value(1)
        # Reset to full step on restart (optional, but safer)
        self.set_microstep('1')
        print("System Reset.")

    # --- The Continuous Loop ---
    def run_continuous(self):
        print("\n[RUNNING] Press 'd' to swap direction, 's' to stop (in serial monitor).")
        
        poll_obj = uselect.poll()
        poll_obj.register(sys.stdin, uselect.POLLIN)
        
        self.dir_pin.value(self.direction)
        
        while True:
            # 1. Step the motor
            delay = self._get_delay()
            self.step_pin.value(1)
            time.sleep(delay / 2)
            self.step_pin.value(0)
            time.sleep(delay)
            
            # Update Position
            if self.direction:
                self.current_step_pos += 1
            else:
                self.current_step_pos -= 1

            # 2. Check for Input
            while poll_obj.poll(0):
                ch = sys.stdin.read(1)
                if not ch:
                    break
                ch = ch.lower()
                
                if ch == 'd':
                    self.toggle_direction_setting()
                    print("---------------------------------------")
                    print("[RUNNING] Press 'd' to swap direction, 's' to stop.")
                
                elif ch == 's':
                    self.stop()
                    return 

# --- Main Program ---
# UPDATED: Added m0, m1, m2 pin definitions
motor = StepperControl(step_pin=3, dir_pin=2, m0_pin=4, m1_pin=5, m2_pin=6 , initialMicroStep = "1/32")

while True:
    print("\n=== MAIN MENU ===")
    print("1: Set RPM & Run")
    print("2: Go to Angle")
    print("3: Toggle Direction (Set Only)")
    print("4: Restart")
    print("5: Set Microstepping") # New Option
    
    cmd = input("Command: ").strip()
    
    if cmd == '1':
        motor.set_rpm_and_run()
    elif cmd == '2':
        motor.go_to_angle()
    elif cmd == '3':
        motor.toggle_direction_setting()
    elif cmd == '4':
        motor.restart()
    elif cmd == '5': # New Option Logic
        print("Available Microsteps: 1, 1/2, 1/4, 1/8, 1/16, 1/32")
        choice = input("Enter resolution: ")
        motor.set_microstep(choice)
    else:
        print("Unknown command.")