---

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Autonomous flight on 7x7 ArUco field (markers 0-48, marker 42 = origin).
Official Clover API: simple_offboard + aruco_map + snippets.

Docs:
  https://clovercoex.tech/ru/simple_offboard.html
  https://clovercoex.tech/ru/aruco_map.html
  https://clovercoex.tech/ru/snippets.html

NAV_MODE:
  'aruco_map'       - absolute coords on marker field (recommended by docs)
  'navigate_target' - relative steps, z=0 (no altitude change per docs)
"""

import math
import rospy
from clover import srv
from clover.srv import SetLEDEffect
from std_srvs.srv import Trigger

# --- settings ---
NAV_MODE = 'aruco_map'   # or 'navigate_target'
FLIGHT_Z = 1.0
SPEED = 0.5
YAW = float('nan')         # keep current yaw (docs)
TOLERANCE = 0.25

rospy.init_node('flight')

get_telemetry = rospy.ServiceProxy('get_telemetry', srv.GetTelemetry)
navigate = rospy.ServiceProxy('navigate', srv.Navigate)
set_altitude = rospy.ServiceProxy('set_altitude', srv.SetAltitude)
land = rospy.ServiceProxy('land', Trigger)
set_effect = rospy.ServiceProxy('led/set_effect', SetLEDEffect)


# --- official snippets ---
def navigate_wait(x=0, y=0, z=0, yaw=YAW, speed=SPEED, frame_id='map',
                  auto_arm=False, tolerance=TOLERANCE):
    navigate(x=x, y=y, z=z, yaw=yaw, speed=speed,
             frame_id=frame_id, auto_arm=auto_arm)
    while not rospy.is_shutdown():
        telem = get_telemetry(frame_id='navigate_target')
        dist = math.sqrt(telem.x ** 2 + telem.y ** 2 + telem.z ** 2)
        if dist < tolerance:
            return
        rospy.sleep(0.2)


def land_wait():
    land()
    while get_telemetry().armed:
        rospy.sleep(0.2)


def set_led(r, g, b):
    try:
        set_effect(r=r, g=g, b=b)
    except rospy.ServiceException:
        pass


def marker_xy(marker_id):
    """Marker center in aruco_map (origin = marker 42, step 1 m)."""
    row = marker_id // 7
    col = marker_id % 7
    return col * 1.0, (6 - row) * 1.0


def midpoint(a, b):
    xa, ya = marker_xy(a)
    xb, yb = marker_xy(b)
    return (xa + xb) / 2.0, (ya + yb) / 2.0


def wait_aruco_map(timeout=30.0):
    deadline = rospy.Time.now() + rospy.Duration(timeout)
    while rospy.Time.now() < deadline and not rospy.is_shutdown():
        telem = get_telemetry(frame_id='aruco_map')
        if not math.isnan(telem.z):
            return True
        rospy.sleep(0.5)
    return False


# (x, y) or (dx, dy), LED rgb or None
WAYPOINTS = [
    (marker_xy(44), None),
    (marker_xy(38), (255, 0, 255)),       # pink
    (midpoint(26, 33), (255, 165, 0)),    # orange
    (midpoint(19, 20), (0, 255, 255)),    # cyan
    (marker_xy(12), (128, 0, 255)),       # purple
    (marker_xy(11), None),
    (marker_xy(17), (255, 0, 0)),         # red
    (marker_xy(9), None),
    (marker_xy(8), (255, 255, 0)),        # yellow
    (midpoint(14, 21), (0, 255, 0)),      # green
    (midpoint(22, 23), (0, 0, 139)),      # dark blue
    (marker_xy(38), (255, 0, 255)),       # pink
    (marker_xy(42), None),
]

# relative legs for navigate_target mode (from marker 42)
RELATIVE_LEGS = [
    (2.0, 0.0, None),
    (1.0, 1.0, (255, 0, 255)),
    (2.0, 1.5, (255, 165, 0)),
    (0.5, 1.5, (0, 255, 255)),
    (-0.5, 1.0, (128, 0, 255)),
    (-1.0, 0.0, None),
    (-1.0, -1.0, (255, 0, 0)),
    (-1.0, 1.0, None),
    (-1.0, 0.0, (255, 255, 0)),
    (-1.0, -1.5, (0, 255, 0)),
    (1.5, -0.5, (0, 0, 139)),
    (1.5, -2.0, (255, 0, 255)),
    (-3.0, -1.0, None),
]


def fly_aruco_map():
    """Absolute flight in aruco_map frame (official aruco_map docs)."""
    set_led(255, 255, 255)
    navigate_wait(z=FLIGHT_Z, speed=0.4, frame_id='body', auto_arm=True)
    rospy.sleep(5.0)

    if not wait_aruco_map():
        rospy.logwarn('aruco_map not detected yet')

    set_altitude(z=FLIGHT_Z, frame_id='aruco_map')

    for (x, y), color in WAYPOINTS:
        if color:
            set_led(*color)
            rospy.sleep(1.0)
        navigate_wait(x=x, y=y, z=FLIGHT_Z, frame_id='aruco_map', speed=SPEED)


def fly_navigate_target():
    """
    Relative flight in navigate_target frame.
    Docs: z=0 means NO altitude change (do NOT use z=1 each step!).
    """
    set_led(255, 255, 255)
    navigate_wait(z=FLIGHT_Z, speed=0.4, frame_id='body', auto_arm=True)
    rospy.sleep(5.0)
    set_altitude(z=FLIGHT_Z, frame_id='body')

    for dx, dy, color in RELATIVE_LEGS:
        if color:
            set_led(*color)
            rospy.sleep(1.0)
        # z=0: horizontal only (simple_offboard docs)
        navigate_wait(x=dx, y=dy, z=0, frame_id='navigate_target', speed=SPEED)


def run_mission():
    if NAV_MODE == 'navigate_target':
        fly_navigate_target()
    else:
        fly_aruco_map()
    land_wait()
    set_led(0, 0, 0)


if __name__ == '__main__':
    try:
        run_mission()
    except rospy.ROSInterruptException:
        pass
    except KeyboardInterrupt:
        try:
            land_wait()
            set_led(0, 0, 0)
        except rospy.ServiceException:
            pass