from dronekit import connect, VehicleMode
import time
import threading
import keyboard  # pip install keyboard

# -----------------------------
# Connect to Pixhawk (USB)
# -----------------------------
vehicle = connect('COM7', baud=57600, wait_ready=True)
vehicle.mode = VehicleMode("MANUAL")
time.sleep(1)
print("Mode:", vehicle.mode.name)

# -----------------------------
# RC override values
# Array order matches physical outputs:
# [Thruster1, Thruster2, Steering1, Steering2] → OUT1,2,3,4
# -----------------------------
rc_values = [1500, 1500, 1500, 1500]
INCREMENT = 10
MAX_RC = 2200
MIN_RC = 800

# -----------------------------
# Helper functions
# -----------------------------
def clamp(val):
    return max(MIN_RC, min(MAX_RC, val))

def send_rc():
    vehicle.channels.overrides = {
        '1': rc_values[0],  # Thruster1 → OUT1
        '2': rc_values[1],  # Thruster2 → OUT2
        '3': rc_values[2],  # Steering1 → OUT3
        '4': rc_values[3],  # Steering2 → OUT4
    }

def set_throttle(value=1500):
    rc_values[0] = clamp(value)
    rc_values[1] = clamp(value)
    send_rc()

def set_steering(position=1500):
    rc_values[2] = clamp(position)
    rc_values[3] = clamp(position)
    send_rc()

# -----------------------------
# Keyboard control loop
# -----------------------------
def update_rc():
    global rc_values
    last_values = rc_values.copy()

    while True:
        key_pressed = None

        # Arm vehicle
        if keyboard.is_pressed('ctrl+space'):
            if not vehicle.armed:
                vehicle.armed = True
                print("Forcing arming... Vehicle armed:", vehicle.armed)
                time.sleep(0.5)

        # ---------------- Throttle control ----------------
        if keyboard.is_pressed('w'):
            rc_values[0] = clamp(rc_values[0] + INCREMENT)
            rc_values[1] = clamp(rc_values[1] + INCREMENT)
            key_pressed = "Throttle Up"
            time.sleep(0.2)
        elif keyboard.is_pressed('s'):
            rc_values[0] = clamp(rc_values[0] - INCREMENT)
            rc_values[1] = clamp(rc_values[1] - INCREMENT)
            key_pressed = "Throttle Down"
            time.sleep(0.2)

        # ---------------- Incremental positional steering ----------------
        if keyboard.is_pressed('a'):
            rc_values[2] = clamp(rc_values[2] - INCREMENT)
            rc_values[3] = clamp(rc_values[3] - INCREMENT)
            key_pressed = "Steer Left"
            time.sleep(0.2)
        elif keyboard.is_pressed('d'):
            rc_values[2] = clamp(rc_values[2] + INCREMENT)
            rc_values[3] = clamp(rc_values[3] + INCREMENT)
            key_pressed = "Steer Right"
            time.sleep(0.2)
        elif keyboard.is_pressed('c'):
            rc_values[2] = rc_values[3] = 1500
            key_pressed = "Center Steering"
            time.sleep(0.2)

        # ---------------- Stop thrusters ----------------
        if keyboard.is_pressed('v'):
            rc_values[0] = rc_values[1] = 1500
            key_pressed = "Stop Thrusters"
            time.sleep(0.2)

        send_rc()

        if rc_values != last_values and key_pressed:
            print(f"[{key_pressed}] Thruster1: {rc_values[0]}, Thruster2: {rc_values[1]}, "
                  f"Steering1: {rc_values[2]}, Steering2: {rc_values[3]}")
            last_values = rc_values.copy()

# -----------------------------
# Start keyboard listener
# -----------------------------
thread = threading.Thread(target=update_rc)
thread.daemon = True
thread.start()

print("USV remote active.")
print("W/S=Throttle | A/D=Steer | C=Center Steering | V=Stop Thrusters | Ctrl+Space=Arm | Ctrl+C=Exit")

# -----------------------------
# Main loop
# -----------------------------
try:
    while True:
        time.sleep(1)
except KeyboardInterrupt:
    print("\nExiting...")
    vehicle.channels.overrides = {}
    vehicle.close()