Update diagnostics/drive_base_diagnostics.py

diagnostics/drive_base_diagnostics.py

@@ -0,0 +1,165 @@
+from pybricks.parameters import Direction, Port, Side, Stop
+from pybricks.robotics import DriveBase
+from pybricks.tools import wait, StopWatch
+
+from usys import stdin
+from uselect import poll
+
+class DriveBaseDiagnostics:
+    def __init__(self, hub, motorclass, dbclass):
+        self.hub = hub
+        self.motorclass = motorclass
+        self.dbclass = dbclass
+        self.drive_base = None
+        self.PORT_MAP = {
+            "A": Port.A,
+            "B": Port.B,
+            "C": Port.C,
+            "D": Port.D,
+            "E": Port.E,
+            "F": Port.F,
+            }
+    def initializeDriveBase(self):
+
+        print("DriveBase setup:")
+        print("1 = Load from savefile (paste JSON)")
+        print("2 = Use defaults")
+        print("3 = Enter values manually")
+
+        choice = input("Choose an option: ")
+
+        # Default values
+        WHEEL_DIAMETER = 68.8
+        AXLE_TRACK = 180
+        DRIVE_SETTINGS = (600, 2000, 300, 2000)
+
+        # Motor ports (None until set)
+        leftmotorport = Port.A
+        rightmotorport = Port.B
+
+        # -----------------------------
+        # OPTION 1: LOAD SAVEFILE
+        # -----------------------------
+        if choice == "1":
+            print("Paste JSON:")
+            raw = input("> ")
+
+            # --- wheel ---
+            if "\"wheel\"" in raw:
+                part = raw.split("\"wheel\"")[1]
+                part = part.split(":")[1]
+                part = part.split(",")[0]
+                WHEEL_DIAMETER = float(part)
+
+            # --- axle ---
+            if "\"axle\"" in raw:
+                part = raw.split("\"axle\"")[1]
+                part = part.split(":")[1]
+                part = part.split(",")[0]
+                AXLE_TRACK = float(part)
+
+            # --- settings ---
+            if "\"settings\"" in raw:
+                part = raw.split("\"settings\"")[1]
+                part = part.split("[")[1]
+                part = part.split("]")[0]
+                nums = part.split(",")
+                DRIVE_SETTINGS = (int(nums[0]), int(nums[1]), int(nums[2]), int(nums[3]))
+
+                            # --- left motor port ---
+                if "\"left_port\"" in raw:
+                    part = raw.split("\"left_port\"")[1]
+                    part = part.split("\"")[1]   # first quoted value
+                    leftmotorport = part
+
+                # --- right motor port ---
+                if "\"right_port\"" in raw:
+                    part = raw.split("\"right_port\"")[1]
+                    part = part.split("\"")[1]
+                    rightmotorport = part
+
+
+            print("Loaded config.")
+
+        # -----------------------------
+        # OPTION 3: MANUAL ENTRY
+        # -----------------------------
+        elif choice == "3":
+            WHEEL_DIAMETER = float(input("Wheel diameter: "))
+            AXLE_TRACK = float(input("Axle track: "))
+
+            print("Enter drive settings:")
+            s1 = int(input("Straight speed: "))
+            s2 = int(input("Straight accel: "))
+            s3 = int(input("Turn rate: "))
+            s4 = int(input("Turn accel: "))
+            DRIVE_SETTINGS = (s1, s2, s3, s4)
+
+            # Ask for motor ports HERE (manual only)
+            leftmotorport = input("Left motor port: ")
+            rightmotorport = input("Right motor port: ")
+
+        # -----------------------------
+        # OPTION 2: DEFAULTS
+        # -----------------------------
+        else:
+            pass
+
+        # -----------------------------
+        # CREATE MOTORS
+        # -----------------------------
+        left_motor = self.motorclass(leftmotorport, Direction.COUNTERCLOCKWISE)
+        right_motor = self.motorclass(rightmotorport, Direction.CLOCKWISE)
+
+        # -----------------------------
+        # CREATE DRIVEBASE
+        # -----------------------------
+        self.drive_base = self.dbclass(left_motor, right_motor, WHEEL_DIAMETER, AXLE_TRACK)
+        self.drive_base.settings(*DRIVE_SETTINGS)
+        self.drive_base.use_gyro(True)
+
+        print("DriveBase initialized.")
+        return DRIVE_SETTINGS
+
+
+    
+    def printAll(self):
+        self.driveRobot()
+    def driveRobot(self):
+        drivesettings = self.initializeDriveBase()
+        print(drivesettings)
+        keyboard = poll()
+        keyboard.register(stdin)
+
+        last_key_time = StopWatch()
+        last_key_time.reset()
+
+        while True:
+            key = None
+
+            # Check for keypress immediately
+            if keyboard.poll(0):
+                key = stdin.read(1).upper()
+                last_key_time.reset()
+
+            # Process key
+            if key:
+                if key == "W":
+                    self.drive_base.drive(drivesettings[0], 0)
+                elif key == "A":
+                    self.drive_base.drive(drivesettings[0], -180)
+                elif key == "S":
+                    self.drive_base.drive(-drivesettings[0], 0)
+                elif key == "D":
+                    self.drive_base.drive(drivesettings[0], 180)
+                elif key == "X":
+                    break
+
+            # Auto-stop if no key for 120 ms
+            if last_key_time.time() > 120:
+                self.drive_base.stop()
+
+            # Tiny manual delay to avoid 100% CPU
+            # (much faster than wait())
+            for _ in range(200):
+                pass