# --- VICKRAM's CODE ---
# To do - make another version of this using DriveBase
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, sin, cos, radians
import asyncio

# Keep these---------------------------------------------------
class Tracker:
    def __init__(self, starting_pos_x, starting_pos_y, starting_angle):
        self.position = [starting_pos_x, starting_pos_y]
        self.angle = starting_angle
    def update(self, straight_distance, delta_angle):
        delta_angle = radians(delta_angle) # Convert to radians
        self.position[0] += straight_distance * cos(delta_angle) # Calculate x coordinate
        self.position[1] += straight_distance * sin(delta_angle) # Calculate y coordinate
        self.angle += delta_angle

    def get_state():
        return self.position, self. angle

class Attachment:
    def __init__(self, port, start_angle=0):
        self.motor = Motor(port)
        self.start_angle = start_angle
    def move(self, degrees, speed):
        self.motor.reset_angle(0)
        target_angle = degrees
        tolerance = 2
        # Movement with timeout
        movement_timer = StopWatch()
        movement_timeout = 3000  # 3 seconds timeout
        while movement_timer.time() < movement_timeout:
            current_angle = self.motor.angle()
            error = target_angle - current_angle
            if abs(error) <= tolerance:
                self.motor.stop()
                break
            if error > 0:
                self.motor.run(speed)
            else:
                self.motor.run(-speed)
            wait(5)
        self.motor.stop()
        self.motor.reset_angle(0)

    def reset(self):
        self.motor.reset_angle(0)
        self.move(self.start_angle)

# Initialize hub and motors
hub = PrimeHub()
left_motor = Motor(Port.C)
right_motor = Motor(Port.A)
hub.imu.reset_heading(0)
tracker = Tracker(0, 0, 0)

# Make sure to measure robot
def straight(distance, speed):  # Distance in millimeters
    target_heading = hub.imu.heading()
    # Reset distance tracking
    left_motor.reset_angle(0)
    right_motor.reset_angle(0)
    # Calculate target distance in motor degrees
    wheel_circumference = pi * 62.4
    target_degrees = abs(distance) / wheel_circumference * 360
    while True:
        # Check current distance traveled - Fixed: Make left_motor abs negative
        left_angle = -abs(left_motor.angle())
        right_angle = abs(right_motor.angle())
        average_angle = (left_angle + right_angle) / 2
        # Stop if it reached the target
        if abs(average_angle) >= target_degrees:
            break
        # Get heading error for correction
        current_heading = hub.imu.heading()
        heading_error = target_heading - current_heading
        # Handle wraparound at 0°/360°
        if heading_error > 180:
            heading_error -= 360
        if heading_error < -180:
            heading_error += 360
        # Calculate motor speeds with correction
        direction = 1 if distance > 0 else -1
        correction = heading_error * 2.0
        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
        left_motor.run(left_speed)
        right_motor.run(right_speed)
        wait(10)
    # Stop motors
    left_motor.stop()
    right_motor.stop()
    wait(20)
    tracker.update(distance, 0)
    print(tracker.get_position())

def turn(theta, speed):  # Negative value is left and positive is right
    start_angle = hub.imu.heading()
    target_angle = theta + start_angle
    # Normalize target angle to -180 to 180 range
    if target_angle > 180:
        target_angle -= 360
    if target_angle < -180:
        target_angle += 360
    while True:
        current_angle = hub.imu.heading()
        # Calculate angle error (how much we still need to turn)
        angle_error = target_angle - current_angle
        # Handle wraparound for shortest path
        if angle_error > 180:
            angle_error -= 360
        if angle_error < -180:
            angle_error += 360
        # Stop if we're close enough (within 2 degrees)
        if abs(angle_error) < 2:
            break
        # Determine turn direction and speed based on error
        if angle_error > 0:
            # Turn right
            left_motor.run(-speed)
            right_motor.run(-speed)
        else:
            # Turn left
            left_motor.run(speed)
            right_motor.run(speed)
        wait(20)
    # Stop both motors
    left_motor.stop()
    right_motor.stop()
    trcker.update(0, theta)
    print(tracker.get_position())
# Setup attachments here-------------------------------
attachment1 = Attachment(Port.D)
attachment2 = Attachment(Port.B)
# Run code goes here-----------------------------------
async def main():
#------------------------------------------------------
run_task(main())