187 lines
6.7 KiB
Python
187 lines
6.7 KiB
Python
|
from pybricks.hubs import PrimeHub
|
||
|
|
from pybricks.pupdevices import Motor
|
||
|
|
from pybricks.parameters import Port, Stop
|
||
|
|
from pybricks.tools import wait, StopWatch, multitask
|
||
|
|
from umath import pi
|
||
|
|
|
||
|
|
class Attachment:
|
||
|
|
def __init__(self, port, start_angle=0):
|
||
|
|
self.motor = Motor(port)
|
||
|
|
self.start_angle = start_angle
|
||
|
|
|
||
|
|
def move_attachment(self, degrees, speed):
|
||
|
|
self.motor.reset_angle(0)
|
||
|
|
target_angle = degrees
|
||
|
|
tolerance = 2
|
||
|
|
movement_timer = StopWatch()
|
||
|
|
movement_timeout = 3000 # 3 seconds timeout
|
||
|
|
|
||
|
|
stuck_threshold = 50 # ms without significant movement
|
||
|
|
last_angle = 0
|
||
|
|
stuck_timer = StopWatch()
|
||
|
|
|
||
|
|
while movement_timer.time() < movement_timeout:
|
||
|
|
current_angle = self.motor.angle()
|
||
|
|
error = target_angle - current_angle
|
||
|
|
|
||
|
|
if abs(error) <= tolerance:
|
||
|
|
self.motor.stop()
|
||
|
|
|
||
|
|
# Check if motor is stuck
|
||
|
|
if abs(current_angle - last_angle) < 1: # Less than 1 degree movement
|
||
|
|
if stuck_timer.time() > stuck_threshold:
|
||
|
|
print("Motor appears stuck")
|
||
|
|
self.motor.stop()
|
||
|
|
else:
|
||
|
|
stuck_timer.reset()
|
||
|
|
last_angle = current_angle
|
||
|
|
|
||
|
|
if error > 0:
|
||
|
|
self.motor.run(speed)
|
||
|
|
else:
|
||
|
|
self.motor.run(-speed)
|
||
|
|
wait(10) # Consistent timing
|
||
|
|
|
||
|
|
self.motor.stop()
|
||
|
|
print("Movement timeout reached")
|
||
|
|
|
||
|
|
def reset_attachment(self):
|
||
|
|
self.motor.reset_angle(0)
|
||
|
|
return self.move_attachment(self.start_angle, 100)
|
||
|
|
|
||
|
|
class Robot:
|
||
|
|
def __init__(self, left_port, right_port, wheel_diameter):
|
||
|
|
# Initialize hub and motors
|
||
|
|
self.hub = PrimeHub()
|
||
|
|
try:
|
||
|
|
self.left_motor = Motor(left_port)
|
||
|
|
self.right_motor = Motor(right_port)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Failed to initialize motors: {e}")
|
||
|
|
|
||
|
|
# Reset and calibrate IMU
|
||
|
|
try:
|
||
|
|
self.hub.imu.reset_heading(0)
|
||
|
|
wait(1000) # Give IMU time to calibrate
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Failed to initialize IMU: {e}")
|
||
|
|
|
||
|
|
# Robot specs
|
||
|
|
self.wheel_diameter = wheel_diameter
|
||
|
|
self.wheel_circumference = pi * self.wheel_diameter
|
||
|
|
|
||
|
|
# Attachments list
|
||
|
|
self.attachments = {}
|
||
|
|
|
||
|
|
def add_attachment(self, name, port, start_angle=0):
|
||
|
|
try:
|
||
|
|
self.attachments[name] = Attachment(port, start_angle)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Failed to add attachment '{name}': {e}")
|
||
|
|
|
||
|
|
def move_attachment(self, attachment_name, degrees, speed):
|
||
|
|
if attachment_name in self.attachments:
|
||
|
|
return self.attachments[attachment_name].move_attachment(degrees, speed)
|
||
|
|
else:
|
||
|
|
print(f"Attachment '{attachment_name}' not found")
|
||
|
|
|
||
|
|
def reset_attachment(self, attachment_name):
|
||
|
|
if attachment_name in self.attachments:
|
||
|
|
return self.attachments[attachment_name].reset_attachment()
|
||
|
|
else:
|
||
|
|
print(f"Attachment '{attachment_name}' not found")
|
||
|
|
|
||
|
|
def straight(self, distance, speed):
|
||
|
|
if not (-200 <= speed <= 200):
|
||
|
|
print(f"Invalid speed: {speed}. Must be between -200 and 200.")
|
||
|
|
|
||
|
|
target_heading = self.hub.imu.heading()
|
||
|
|
# Reset distance tracking
|
||
|
|
self.left_motor.reset_angle(0)
|
||
|
|
self.right_motor.reset_angle(0)
|
||
|
|
|
||
|
|
# Calculate target distance in motor degrees
|
||
|
|
target_degrees = abs(distance) / self.wheel_circumference * 360
|
||
|
|
|
||
|
|
while True:
|
||
|
|
# Check current distance traveled (assuming both motors should be positive for forward)
|
||
|
|
left_angle = abs(self.left_motor.angle())
|
||
|
|
right_angle = abs(self.right_motor.angle())
|
||
|
|
average_angle = (left_angle + right_angle) / 2
|
||
|
|
|
||
|
|
# Stop if it reached the target
|
||
|
|
if average_angle >= target_degrees:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Get heading error for correction
|
||
|
|
current_heading = self.hub.imu.heading()
|
||
|
|
heading_error = target_heading - current_heading
|
||
|
|
|
||
|
|
# Handle wraparound at 0°/360°
|
||
|
|
if heading_error > 180:
|
||
|
|
heading_error -= 360
|
||
|
|
elif heading_error < -180:
|
||
|
|
heading_error += 360
|
||
|
|
|
||
|
|
# Calculate motor speeds with correction
|
||
|
|
direction = 1 if distance > 0 else -1
|
||
|
|
correction = heading_error * 1.5 # Reduced gain for stability
|
||
|
|
left_speed = direction * (speed + correction)
|
||
|
|
right_speed = direction * (speed - correction)
|
||
|
|
|
||
|
|
# Limit speeds to prevent excessive values
|
||
|
|
left_speed = max(-200, min(200, left_speed))
|
||
|
|
right_speed = max(-200, min(200, right_speed))
|
||
|
|
|
||
|
|
# Apply speeds to motors
|
||
|
|
self.left_motor.run(left_speed)
|
||
|
|
self.right_motor.run(right_speed)
|
||
|
|
wait(10) # Consistent timing
|
||
|
|
|
||
|
|
# Stop motors
|
||
|
|
self.left_motor.stop()
|
||
|
|
self.right_motor.stop()
|
||
|
|
wait(50) # Allow motors to fully stop
|
||
|
|
|
||
|
|
def turn(self, theta, speed):
|
||
|
|
if not (-200 <= speed <= 200):
|
||
|
|
print(f"Invalid speed: {speed}. Must be between -200 and 200.")
|
||
|
|
|
||
|
|
start_angle = self.hub.imu.heading()
|
||
|
|
target_angle = start_angle + theta
|
||
|
|
|
||
|
|
# Normalize target angle to -180 to 180 range
|
||
|
|
while target_angle > 180:
|
||
|
|
target_angle -= 360
|
||
|
|
while target_angle < -180:
|
||
|
|
target_angle += 360
|
||
|
|
|
||
|
|
while True:
|
||
|
|
current_angle = self.hub.imu.heading()
|
||
|
|
# Calculate angle error
|
||
|
|
angle_error = target_angle - current_angle
|
||
|
|
|
||
|
|
# Handle wraparound for shortest path
|
||
|
|
if angle_error > 180:
|
||
|
|
angle_error -= 360
|
||
|
|
elif angle_error < -180:
|
||
|
|
angle_error += 360
|
||
|
|
|
||
|
|
# Stop if we're close enough
|
||
|
|
if abs(angle_error) < 2:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Positive angle_error means we need to turn clockwise (right)
|
||
|
|
if angle_error > 0:
|
||
|
|
# Turn right: left motor forward, right motor backward
|
||
|
|
self.left_motor.run(speed)
|
||
|
|
self.right_motor.run(-speed)
|
||
|
|
else:
|
||
|
|
# Turn left: left motor backward, right motor forward
|
||
|
|
self.left_motor.run(-speed)
|
||
|
|
self.right_motor.run(speed)
|
||
|
|
wait(10) # Consistent timing
|
||
|
|
|
||
|
|
# Stop both motors
|
||
|
|
self.left_motor.stop()
|
||
|
|
self.right_motor.stop()
|