from dronekit import connect, VehicleMode, LocationGlobalRelative
import time
import math

# Connect to the vehicle
vehicle = connect('127.0.0.1:14550', wait_ready=True)  # Replace with your connection string

def arm_and_takeoff(target_altitude):
    print("Arming motors...")
    while not vehicle.is_armable:
        print("Waiting for vehicle to initialize...")
        time.sleep(1)

    vehicle.mode = VehicleMode("GUIDED")
    vehicle.armed = True

    while not vehicle.armed:
        print("Waiting for arming...")
        time.sleep(1)

    print("Taking off!")
    vehicle.simple_takeoff(target_altitude)

    while True:
        print(f"Altitude: {vehicle.location.global_relative_frame.alt:.2f} m")
        if vehicle.location.global_relative_frame.alt >= target_altitude * 0.95:
            print("Reached target altitude")
            break
        time.sleep(1)

def send_ned_velocity(velocity_x, velocity_y, velocity_z, duration):
    """
    Move vehicle in the NED direction.
    velocity_x: velocity in m/s along the North axis
    velocity_y: velocity in m/s along the East axis
    velocity_z: velocity in m/s along the Down axis (positive is down)
    duration: time in seconds to execute the command
    """
    msg = vehicle.message_factory.set_position_target_local_ned_encode(
        0, 0, 0,  # Target system, target component, target coordinate frame
        mavutil.mavlink.MAV_FRAME_LOCAL_NED,  # Frame
        0b0000111111000111,  # Type mask
        0, 0, 0,  # Position (not used)
        velocity_x, velocity_y, velocity_z,  # Velocity
        0, 0, 0,  # Acceleration
        0, 0)  # Yaw, yaw rate
    for _ in range(int(duration * 10)):
        vehicle.send_mavlink(msg)
        time.sleep(0.1)

def land():
    print("Landing...")
    vehicle.mode = VehicleMode("LAND")
    while vehicle.armed:
        print("Waiting for landing...")
        time.sleep(1)
    print("Landed!")

# Main mission
arm_and_takeoff(1)

# Move forward (North) by 3 meters
send_ned_velocity(1, 0, 0, 3)  # 1 m/s for 3 seconds

# Move left (West, negative East) by 2 meters
send_ned_velocity(0, -1, 0, 2)  # 1 m/s for 2 seconds

# Land
land()

# Close vehicle connection
vehicle.close()