#!/usr/bin/env python3

import asyncio
import argparse

from mavsdk import System

def arg_parser():
    parser = argparse.ArgumentParser(
        prog='calibrate.py',
        epilog="See '<command> --help' for more detail")

    parser.add_argument('-g', '--gyroscope', action='store_true')
    parser.add_argument('-a', '--accelerometer', action='store_true')
    parser.add_argument('-m', '--magnetometer', action='store_true')
    parser.add_argument('-lh', '--horizon', action='store_true')
    parser.add_argument('-e', '--esc', action='store_true')
    parser.add_argument('--port', type=str, default="/dev/pixhawk")
    parser.add_argument('--baudrate', type=int, default=2000000)

    args = parser.parse_args()
    print(args)

    if args.gyroscope or args.accelerometer or args.magnetometer or args.horizon or args.esc:
        return args

    parser.print_help()
    return None

async def main(args):

    drone = System()

    serial = "serial://" + str(args.port) + ":" + str(args.baudrate) 

    await drone.connect(system_address=serial)

    print("Connected to a device")

    if (args.gyroscope):
        print("-- Starting gyroscope calibration")
        async for progress_data in drone.calibration.calibrate_gyro():
            print(progress_data)
        print("-- Gyroscope calibration finished")

    if (args.accelerometer):
        print("-- Starting accelerometer calibration")
        async for progress_data in drone.calibration.calibrate_accelerometer():
            print(progress_data)
        print("-- Accelerometer calibration finished")

    if (args.magnetometer):
        print("-- Starting magnetometer calibration")
        async for progress_data in drone.calibration.calibrate_magnetometer():
            print(progress_data)
        print("-- Magnetometer calibration finished")

    if (args.horizon):
        print("-- Starting board level horizon calibration")
        async for progress_data in drone.calibration.calibrate_level_horizon():
            print(progress_data)
        print("-- Level horizon calibration finished")

    if (args.esc):
        print("-- Starting ESC calibration")
        print("-- To be done --")
        print("-- ESC calibration finished")


if __name__ == "__main__":
    # Run the asyncio loop
    args = arg_parser()
    if args is not None:
        asyncio.run(main(args)) 
