from pybricks import version
from pybricks.hubs import InventorHub
from pybricks.parameters import Icon
from pybricks.tools import run_task, wait, StopWatch, multitask

hub = InventorHub()

print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
print("Testing the stationary drift of the Prime Hub.")
print("Pybricks version:", version)
print("Pybricks IMU settings:", hub.imu.settings())
print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")


stop_watch = StopWatch()


async def calibrate_imu():
    hub.imu.settings(3, 2500)
    while not hub.imu.stationary():
        await wait(100)


async def lock_recalibration():
    hub.imu.settings(0, 0)
    await wait(1)


async def heading_and_drift_print_loop():
    prev_heading = hub.imu.heading()
    print("Time[min], Heading[deg], Drift [deg/min]")
    while True:
        await wait(60 * 1000)
        heading = hub.imu.heading()
        drift = (heading - prev_heading)
        prev_heading = heading
        print("{:2d}, {:4.1f}, {:2.1f}".format(
            int(stop_watch.time() / (60 * 1000)), heading, drift))


async def led_impact_test():

    print("### Testing LED impact on drift START: 5 minutes without LED - 5 minutes with LED")

    async def led_on_after_for(delay_min=5, for_min=5):
        await wait(delay_min * 60 * 1000)
        hub.display.icon(Icon.FULL)
        await wait(for_min * 60 * 1000)

    stop_watch.reset()
    display = hub.display.off()
    await calibrate_imu()
    await lock_recalibration()
    hub.imu.reset_heading(0)
    await multitask(
        heading_and_drift_print_loop(),
        led_on_after_for(5, 5),
        race=True),

    print("### Testing LED impact on drift DONE")

run_task(led_impact_test())