from pybricks.hubs import PrimeHub
from pybricks.pupdevices import Motor, ColorSensor, UltrasonicSensor, ForceSensor
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop, Axis
from pybricks.robotics import DriveBase
from pybricks.tools import wait, StopWatch, run_task, multitask
import umath

hub = PrimeHub()

left_motor = Motor(Port.A, Direction.COUNTERCLOCKWISE)
right_motor = Motor(Port.B, Direction.CLOCKWISE)
left_arm = Motor(Port.C, Direction.CLOCKWISE, [[12,36]], [[12,20,24]])
right_arm = Motor(Port.D, Direction.CLOCKWISE, [[12,36],[12,20,24]])
lazer_ranger = UltrasonicSensor(Port.E)
color_sensor = ColorSensor(Port.F)

WHEEL_DIAMETER = 68.8
AXLE_TRACK = 180
drive_base = DriveBase(left_motor, right_motor, WHEEL_DIAMETER, AXLE_TRACK)
drive_base.settings(600, 500, 300, 200)
drive_base.use_gyro(False)

current = {"x": 0, "y": 0, "a": 0}
target  = {"x": 0, "y": 200, "a": 0}
done = False

def addtoposition(a, r):
    current["x"] += r * umath.sin(umath.radians(a))
    current["y"] += r * umath.cos(umath.radians(a))
    current["a"] = a

async def track_heading():
    while True:
        current["a"] = hub.imu.heading()
        await wait(20)

async def printStats():
    last_dist = 0
    while True:
        if done:
            break
        current_dist = drive_base.distance()
        instant_dist = current_dist - last_dist
        last_dist = current_dist

        current["x"] += instant_dist * umath.sin(umath.radians(current["a"]))
        current["y"] += instant_dist * umath.cos(umath.radians(current["a"]))

        dx = target["x"] - current["x"]
        dy = target["y"] - current["y"]
        remaining = umath.sqrt(dx*dx + dy*dy)
        desired_a = 90 - umath.degrees(umath.atan2(dy, dx))
        print("x:", current["x"], "y:", current["y"], "desired_a:", desired_a, "current_a:", current["a"], "remaining:", remaining)

        await wait(50)

async def pid_turn(target_angle):
    KP = 3
    KI = 0.01
    KD = 1.5
    SETTLE_THRESHOLD = 0.5
    SETTLE_TIME = 500

    integral = 0
    last_error = 0
    settled_timer = StopWatch()
    settled = False

    while True:
        error = target_angle - current["a"]
        while error > 180:  error -= 360
        while error < -180: error += 360

        integral += error
        integral = max(-100, min(100, integral))
        derivative = error - last_error
        last_error = error

        turn_rate = KP * error + KI * integral + KD * derivative
        turn_rate = max(-200, min(200, turn_rate))

        drive_base.drive(0, turn_rate)

        if abs(error) < SETTLE_THRESHOLD:
            if not settled:
                settled = True
                settled_timer.reset()
            elif settled_timer.time() > SETTLE_TIME:
                drive_base.stop()
                break
        else:
            settled = False

        await wait(20)

async def driveTrajectory():
    global done
    ARRIVE_THRESHOLD = 30  # give stage 2 more room to work
    FINAL_THRESHOLD = 7
    SPEED = 400
    SLOW_SPEED = 30
    TURN_GAIN = 3.0

    # stage 1
    while True:
        dx = target["x"] - current["x"]
        dy = target["y"] - current["y"]
        remaining = umath.sqrt(dx*dx + dy*dy)
        if remaining < ARRIVE_THRESHOLD:
            drive_base.stop()
            break
        desired_a = 90 - umath.degrees(umath.atan2(dy, dx))
        error = desired_a - current["a"]
        while error > 180:  error -= 360
        while error < -180: error += 360
        turn_rate = max(-200, min(200, error * TURN_GAIN))
        # in stage 1, scale speed down as it gets closer
        speed = max(150, SPEED * min(1, remaining / 300))
        drive_base.drive(speed, turn_rate)
        await wait(20)

    print("Stage 1 done")

    # stage 2
    while True:
        dx = target["x"] - current["x"]
        dy = target["y"] - current["y"]
        remaining = umath.sqrt(dx*dx + dy*dy)
        if remaining < FINAL_THRESHOLD:
            drive_base.stop()
            break
        a_rad = umath.radians(current["a"])
        forward_x = umath.sin(a_rad)
        forward_y = umath.cos(a_rad)
        signed_dist = dx * forward_x + dy * forward_y
        speed = SLOW_SPEED if signed_dist > 0 else -SLOW_SPEED
        drive_base.drive(speed, 0)
        await wait(20)

    drive_base.stop()
    print("Stage 2 done")

    # stage 3 — PID angle correction
    await pid_turn(target["a"])

    done = True
    print("Done! x:", current["x"], "y:", current["y"], "a:", current["a"])

async def main():
    drive_base.reset()
    await multitask(driveTrajectory(), track_heading(), printStats())

hub.imu.reset_heading(0)
while True:
    if hub.imu.ready():
        break

run_task(main())
print("done")