---

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Автономный полёт по трассе на полигоне ArUco-меток (7×7, маркеры 0–48).
Платформа: Clover + Gazebo + ROS (simple_offboard).

Маршрут:
  42 (взлёт, белый LED) → 44 → 38 (розовый) →
  правый борт (оранжевый, голубой) → 12 (фиолетовый) →
  11 → 17 (красный) → 9 → 8 (жёлтый) →
  левый борт (зелёный, тёмно-синий) → 38 (розовый) →
  42 (посадка, LED выкл.)
"""

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

# --- Параметры полёта ---
FLIGHT_Z = 1.0
SPEED = 0.6
TOLERANCE = 0.3
MAP_WAIT_TIMEOUT = 30.0   # макс. ожидание фрейма aruco_map, с
ARRIVAL_TIMEOUT = 90.0    # макс. ожидание одной точки, с

# Фрейм карты меток (заполняется после взлёта)
MAP_FRAME = 'aruco_map'

rospy.init_node('autonomous_flight')

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

tf_buffer = tf2_ros.Buffer()
tf_listener = tf2_ros.TransformListener(tf_buffer)


def set_led(r, g, b, effect='fill'):
    set_effect(effect=effect, r=r, g=g, b=b)


def led_off():
    set_led(0, 0, 0)


WHITE = (255, 255, 255)
PINK = (255, 0, 255)
ORANGE = (255, 128, 0)
CYAN = (0, 255, 255)
PURPLE = (128, 0, 255)
RED = (255, 0, 0)
YELLOW = (255, 255, 0)
GREEN = (0, 255, 0)
DARK_BLUE = (0, 0, 139)


def marker_xy(marker_id):
    """
    Координаты центра маркера в aruco_map.
    Начало — маркер №42 (левый нижний угол), шаг 1 м.
    """
    row = marker_id // 7
    col = marker_id % 7
    return col * 1.0, (6 - row) * 1.0


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


GREEN_BLOCK = midpoint(14, 21)
DARK_BLUE_BLOCK = midpoint(22, 23)
CYAN_BLOCK = midpoint(19, 20)
ORANGE_BLOCK = midpoint(26, 33)


def wait_for_map_frame():
    """Ждём появления фрейма карты меток в TF-дереве."""
    global MAP_FRAME
    deadline = rospy.Time.now() + rospy.Duration(MAP_WAIT_TIMEOUT)
    candidates = ('aruco_map', 'aruco_map_detected')

    rospy.loginfo('Waiting for aruco_map frame...')
    rate = rospy.Rate(5)
    while not rospy.is_shutdown() and rospy.Time.now() < deadline:
        for frame in candidates:
            try:
                tf_buffer.lookup_transform('map', frame, rospy.Time(0), rospy.Duration(0.5))
                MAP_FRAME = frame
                rospy.loginfo('Map frame found: %s', MAP_FRAME)
                return
            except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
                pass
        rate.sleep()

    rospy.logwarn('aruco_map frame not found, continuing anyway')


def wait_arrival(tolerance=TOLERANCE, timeout=ARRIVAL_TIMEOUT):
    start = rospy.Time.now()
    while not rospy.is_shutdown():
        if (rospy.Time.now() - start).to_sec() > timeout:
            rospy.logwarn('Waypoint timeout')
            return
        telem = get_telemetry(frame_id='navigate_target')
        if math.sqrt(telem.x ** 2 + telem.y ** 2 + telem.z ** 2) < tolerance:
            return
        rospy.sleep(0.2)


def fly_to(x, y, z=FLIGHT_Z, speed=SPEED, label=''):
    if label:
        rospy.loginfo('Fly to %s (%.1f, %.1f, %.1f) [%s]', label, x, y, z, MAP_FRAME)
    navigate(x=x, y=y, z=z, speed=speed, frame_id=MAP_FRAME)
    wait_arrival()


def fly_to_marker(marker_id, z=FLIGHT_Z, speed=SPEED):
    x, y = marker_xy(marker_id)
    fly_to(x, y, z=z, speed=speed, label='marker %d' % marker_id)


def fly_to_block(xy, color, label='', z=FLIGHT_Z, speed=SPEED):
    set_led(*color)
    fly_to(xy[0], xy[1], z=z, speed=speed, label=label)


def takeoff():
    set_led(*WHITE)
    rospy.loginfo('Takeoff to %.1f m', FLIGHT_Z)
    navigate(x=0, y=0, z=FLIGHT_Z, speed=0.4, frame_id='body', auto_arm=True)
    wait_arrival()
    wait_for_map_frame()
    rospy.sleep(2.0)


def land_at_home():
    rospy.loginfo('Return to marker 42 and land')
    fly_to_marker(42)
    land()
    while get_telemetry().armed:
        rospy.sleep(0.2)
    led_off()


def run_mission():
    takeoff()

    fly_to_marker(44)
    set_led(*PINK)
    fly_to_marker(38)

    fly_to_block(ORANGE_BLOCK, ORANGE, label='orange block')
    fly_to_block(CYAN_BLOCK, CYAN, label='cyan block')
    set_led(*PURPLE)
    fly_to_marker(12)

    fly_to_marker(11)
    set_led(*RED)
    fly_to_marker(17)

    fly_to_marker(9)
    set_led(*YELLOW)
    fly_to_marker(8)

    fly_to_block(GREEN_BLOCK, GREEN, label='green block')
    fly_to_block(DARK_BLUE_BLOCK, DARK_BLUE, label='dark blue block')
    set_led(*PINK)
    fly_to_marker(38)

    land_at_home()
    rospy.loginfo('Mission complete')


if __name__ == '__main__':
    try:
        run_mission()
    except rospy.ROSInterruptException:
        pass
    except KeyboardInterrupt:
        rospy.loginfo('Interrupted, landing')
        try:
            land()
            led_off()
        except rospy.ServiceException:
            pass