orchestra.py

#!/usr/bin/env python3
"""
Fast-LIO Drone Orchestra Controller
-----------------------------------
1. Launches MAVROS, Ouster, Fast-LIO, and Robot State Publisher.
2. Streams Fast-LIO odometry to the FCU via MAVLink (external navigation feed).
3. Records odometry samples for reverse trajectories.
4. Detects radio cuts and replays the latest path in reverse via reverse_trajectory_follower.
"""

import csv
import os
import signal
import subprocess
import threading
import time
from typing import Callable, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseStamped, TwistStamped
from mavros_msgs.msg import RCIn, State, StatusText
from pymavlink import mavutil

# =========================== CONFIGURATION ===========================
ROS_DISTRO = "humble"
WORKSPACE = os.path.expanduser("~/LIONav/ros2_ws")
URDF_PATH = os.path.expanduser("~/LIONav/ros2_ws/src/my_drone_description/urdf/drone_body.urdf")

FAST_LIO_TOPIC = os.environ.get("FAST_LIO_TOPIC", "/Odometry")
FAST_LIO_WAIT_SEC = float(os.environ.get("FAST_LIO_WAIT_SEC", 20.0))
FAST_LIO_LAUNCH = os.environ.get(
    "FAST_LIO_LAUNCH",
    "ros2 launch fast_lio mapping.launch.py config_file:=ouster32.yaml",
)
OUSTER_HOSTNAME = os.environ.get("OUSTER_HOSTNAME", "os-122511001246.local")
OUSTER_LAUNCH = (
    "ros2 launch ouster_ros sensor.launch.xml "
    f"sensor_hostname:={OUSTER_HOSTNAME} "
    "timestamp_mode:=TIME_FROM_ROS_TIME viz:=false"
)
MAVROS_LAUNCH = (
    "ros2 run mavros mavros_node "
    "--ros-args "
    "-p fcu_url:=serial:///dev/ttyTHS1:921600 "
    "-p fcu_protocol:=v2.0 "
    "-p use_companion_time:=true "
    "-p timesync_mode:=MAVLINK "
    "-p timesync_rate:=10.0 "
    "-p system_time_rate:=1.0 "
    "-p plugin_allowlist:=\"['sys_time','sys_status','imu','odometry','setpoint_position','setpoint_velocity','setpoint_attitude','actuator_control','rc_io']\" "
    "-p plugin_denylist:=\"['distance_sensor']\""
)
ROBOT_STATE_PUBLISHER_CMD = (
    f"ros2 run robot_state_publisher robot_state_publisher --ros-args -p robot_description:=\"$(xacro {URDF_PATH})\""
)

MAVLINK_TARGET = os.environ.get("MAVLINK_TARGET", "192.168.1.49:5762")
EXTERNAL_NAV_RATE_HZ = float(os.environ.get("EXTERNAL_NAV_RATE_HZ", 30.0))
REVERSE_SCRIPT = os.path.expanduser("~/LIONav/utils/reverse_trajectory_follower.py")
TRAJECTORY_FILE = os.path.expanduser("~/rtl_fastlio_latest.csv")
TRAJ_SAMPLE_HZ = float(os.environ.get("TRAJ_SAMPLE_HZ", 20.0))
TRAJ_MIN_BYTES = int(os.environ.get("TRAJ_MIN_BYTES", 2048))
RC_TIMEOUT_SEC = float(os.environ.get("RC_TIMEOUT_SEC", 3.0))
RC_VALID_MIN = 900
RC_VALID_MAX = 2100

POSE_COVARIANCE = [float("nan")] * 21
VEL_COVARIANCE = [float("nan")] * 21

# ====================================================================

processes = []
stop_event = threading.Event()
trajectory_logger: Optional["TrajectoryLogger"] = None
nav_bridge: Optional["ExternalNavBridge"] = None
monitor_node_ref: Optional["FastLioOrchestraNode"] = None
reverse_lock = threading.Lock()


class TrajectoryLogger:
    """CSV logger for odometry samples."""

    HEADER = [
        "stamp_sec",
        "stamp_nanosec",
        "x",
        "y",
        "z",
        "qx",
        "qy",
        "qz",
        "qw",
        "vx",
        "vy",
        "vz",
    ]

    def __init__(self, filepath: str, sample_hz: float):
        self.filepath = filepath
        self.sample_interval = 1.0 / sample_hz if sample_hz > 0 else 0.0
        self.file = None
        self.writer = None
        self.last_sample_time = None
        self.samples = 0
        self.lock = threading.Lock()

    def start(self):
        os.makedirs(os.path.dirname(self.filepath) or ".", exist_ok=True)
        self.file = open(self.filepath, "w", newline="")
        self.writer = csv.writer(self.file)
        self.writer.writerow(self.HEADER)
        self.file.flush()
        self.samples = 0
        self.last_sample_time = None
        print(f"✅ Trajectory logger writing to {self.filepath}")

    def stop(self):
        with self.lock:
            if self.file:
                self.file.flush()
                self.file.close()
            self.file = None
            self.writer = None
            print(f"✅ Trajectory logger stopped after {self.samples} samples")

    def try_write(self, msg: Odometry):
        if not self.writer:
            return

        stamp_sec = msg.header.stamp.sec
        stamp_nsec = msg.header.stamp.nanosec
        if stamp_sec == 0 and stamp_nsec == 0:
            now = time.time()
            stamp_sec = int(now)
            stamp_nsec = int((now - stamp_sec) * 1e9)
        stamp_float = stamp_sec + stamp_nsec * 1e-9

        if self.sample_interval and self.last_sample_time is not None:
            if (stamp_float - self.last_sample_time) < self.sample_interval:
                return

        pose = msg.pose.pose
        twist = msg.twist.twist
        row = [
            stamp_sec,
            stamp_nsec,
            pose.position.x,
            pose.position.y,
            pose.position.z,
            pose.orientation.x,
            pose.orientation.y,
            pose.orientation.z,
            pose.orientation.w,
            twist.linear.x,
            twist.linear.y,
            twist.linear.z,
        ]

        with self.lock:
            self.writer.writerow(row)
            self.file.flush()
            self.samples += 1
            self.last_sample_time = stamp_float


class ExternalNavBridge:
    """Streams Fast-LIO odometry to the FCU via MAVLink ODOMETRY."""

    def __init__(self, target_tcp: str, rate_hz: float):
        self.target_tcp = target_tcp
        self.min_interval = 1.0 / rate_hz if rate_hz > 0 else 0.0
        self.master = None
        self.lock = threading.Lock()
        self.latest_payload = None
        self.last_send = 0.0
        self.running = True
        self.thread = threading.Thread(target=self._pump_loop, daemon=True)
        self.thread.start()

    def _connect(self) -> bool:
        with self.lock:
            if self.master:
                return True
            try:
                self.master = mavutil.mavlink_connection(f"tcp:{self.target_tcp}", autoreconnect=True)
                self.master.wait_heartbeat(timeout=5)
                print(f"✅ External nav stream connected to {self.target_tcp}")
                return True
            except Exception as exc:
                print(f"⚠️  Unable to reach {self.target_tcp} for external nav: {exc}")
                self.master = None
                return False

    def send_odometry(self, msg: Odometry):
        payload = self._extract_payload(msg)
        if payload is None:
            return
        with self.lock:
            self.latest_payload = payload

    def _pump_loop(self):
        while self.running and not stop_event.is_set():
            payload = None
            with self.lock:
                payload = self.latest_payload
            if payload:
                self._emit(payload)
            time.sleep(self.min_interval or 0.02)

    def _extract_payload(self, msg: Odometry):
        pose = msg.pose.pose
        twist = msg.twist.twist
        quat = pose.orientation
        return {
            "position": (pose.position.x, pose.position.y, pose.position.z),
            "orientation": (quat.w, quat.x, quat.y, quat.z),
            "linear": (twist.linear.x, twist.linear.y, twist.linear.z),
            "angular": (twist.angular.x, twist.angular.y, twist.angular.z),
        }

    def _emit(self, payload):
        now = time.time()
        if not self.master and not self._connect():
            return
        time_usec = int(time.time() * 1e6)
        try:
            self.master.mav.odometry_send(
                time_usec,
                mavutil.mavlink.MAV_FRAME_LOCAL_ENU,
                mavutil.mavlink.MAV_FRAME_BODY_FRD,
                *payload["position"],
                payload["orientation"],
                *payload["linear"],
                *payload["angular"],
                POSE_COVARIANCE,
                VEL_COVARIANCE,
                estimator_type=mavutil.mavlink.MAV_ESTIMATOR_TYPE_VISION,
                quality=100,
            )
            self.last_send = now
        except Exception as exc:
            print(f"⚠️  Failed to stream odometry: {exc}")
            with self.lock:
                self.master = None

    def close(self):
        self.running = False
        if self.thread:
            self.thread.join(timeout=1.0)
        with self.lock:
            if self.master:
                try:
                    self.master.close()
                except Exception:
                    pass
            self.master = None


class FastLioOrchestraNode(Node):
    """ROS 2 node that bridges Fast-LIO odometry and monitors radio health."""

    def __init__(self, nav_bridge: ExternalNavBridge):
        super().__init__("fast_lio_orchestra")
        self.nav_bridge = nav_bridge
        self.trajectory_logger: Optional[TrajectoryLogger] = None
        self.odometry_ready = False
        self.monitoring_enabled = False
        self.rc_connected = False
        self.rc_ever_received = False
        self.last_rc_time = time.time()
        self.rtl_triggered = False
        self.mavros_connected = False
        self.radio_failsafe = False
        self.connection_loss_cb: Optional[Callable[[str], None]] = None

        self.create_subscription(Odometry, FAST_LIO_TOPIC, self._handle_odometry, 50)
        self.create_subscription(RCIn, "/mavros/rc/in", self._handle_rc, 10)
        self.create_subscription(State, "/mavros/state", self._handle_state, 10)
        statustext_qos = QoSProfile(depth=10)
        statustext_qos.reliability = ReliabilityPolicy.BEST_EFFORT
        self.create_subscription(StatusText, "/mavros/statustext/recv", self._handle_statustext, statustext_qos)
        self.create_timer(1.0, self._rc_watchdog)

        qos_out = QoSProfile(history=HistoryPolicy.KEEP_LAST, depth=10)
        qos_out.reliability = ReliabilityPolicy.BEST_EFFORT
        self.vision_pose_pub = self.create_publisher(PoseStamped, "/mavros/vision_pose/pose", qos_out)
        self.vision_speed_pub = self.create_publisher(TwistStamped, "/mavros/vision_speed/speed_twist", qos_out)

    def attach_logger(self, logger: Optional[TrajectoryLogger]):
        self.trajectory_logger = logger

    def enable_monitoring(self):
        self.monitoring_enabled = True
        self.rc_connected = self.rc_ever_received
        self.last_rc_time = time.time()
        self.get_logger().info("🔍 RC monitoring ENABLED")

    def set_connection_loss_callback(self, callback: Callable[[str], None]):
        self.connection_loss_cb = callback

    def _handle_odometry(self, msg: Odometry):
        if not self.odometry_ready:
            self.odometry_ready = True
            self.get_logger().info("✅ Fast-LIO odometry detected")
        if self.nav_bridge:
            self.nav_bridge.send_odometry(msg)
        if self.trajectory_logger:
            self.trajectory_logger.try_write(msg)
        self._publish_mavros_vision(msg)

    def _publish_mavros_vision(self, msg: Odometry):
        stamp = msg.header.stamp
        if stamp.sec == 0 and stamp.nanosec == 0:
            stamp = self.get_clock().now().to_msg()
        pose_msg = PoseStamped()
        pose_msg.header.stamp = stamp
        pose_msg.header.frame_id = msg.header.frame_id or "map"
        pose_msg.pose = msg.pose.pose
        speed_msg = TwistStamped()
        speed_msg.header.stamp = stamp
        speed_msg.header.frame_id = msg.child_frame_id or "base_link"
        speed_msg.twist = msg.twist.twist
        self.vision_pose_pub.publish(pose_msg)
        self.vision_speed_pub.publish(speed_msg)

    def _handle_rc(self, msg: RCIn):
        valid = any(
            RC_VALID_MIN < ch < RC_VALID_MAX and ch != 65535
            for ch in msg.channels[:4]
        ) if msg.channels else False
        if valid:
            self.last_rc_time = time.time()
            if not self.rc_ever_received:
                self.rc_ever_received = True
                self.get_logger().info(f"✅ RC detected: {list(msg.channels[:4])}")
            if not self.rc_connected and self.monitoring_enabled:
                self.rc_connected = True
                self.get_logger().info("✅ RC link restored")
        elif self.rc_connected and self.monitoring_enabled:
            self.get_logger().warn(f"⚠️  Invalid RC values: {list(msg.channels[:4])}")

    def _handle_state(self, msg: State):
        if self.monitoring_enabled and self.mavros_connected and not msg.connected:
            self.get_logger().error("⚠️  MAVROS connection lost")
        self.mavros_connected = msg.connected
        if self.monitoring_enabled and msg.mode.upper() in {"RTL", "LAND", "AUTO.RTL", "AUTO.LAND"}:
            self._trigger_connection_loss(f"FCU entered {msg.mode}")

    def _handle_statustext(self, msg: StatusText):
        text = msg.text.lower()
        if "radio failsafe" in text and "cleared" not in text:
            self._trigger_connection_loss("Radio failsafe from FCU")
        elif "radio failsafe" in text and self.radio_failsafe:
            self.radio_failsafe = False
            self.get_logger().info(f"✅ Radio failsafe cleared: {msg.text}")

    def _rc_watchdog(self):
        if not self.monitoring_enabled or not self.rc_ever_received or self.rtl_triggered:
            return
        if (time.time() - self.last_rc_time) > RC_TIMEOUT_SEC:
            self._trigger_connection_loss("RC timeout")

    def _trigger_connection_loss(self, reason: str):
        if self.rtl_triggered:
            return
        self.rtl_triggered = True
        self.get_logger().error(f"🚨 Radio cut detected: {reason}")
        if self.connection_loss_cb:
            threading.Thread(target=self.connection_loss_cb, args=(reason,), daemon=False).start()


# =========================== UTILITIES ===============================

def run_in_tab(title: str, command: str):
    cmd = [
        "gnome-terminal",
        "--tab",
        "--title",
        title,
        "--",
        "zsh",
        "-c",
        f"source /opt/ros/{ROS_DISTRO}/setup.zsh; "
        f"source {WORKSPACE}/install/setup.zsh; "
        f"{command}; exec zsh",
    ]
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    processes.append(proc)
    return proc


def launch_ros_stack():
    print("🚀 Launching ROS 2 stack...")
    run_in_tab("MAVROS", MAVROS_LAUNCH)
    time.sleep(3)
    run_in_tab("Ouster ROS", OUSTER_LAUNCH)
    time.sleep(2)
    run_in_tab("FAST-LIO", FAST_LIO_LAUNCH)
    time.sleep(2)
    run_in_tab("Robot State Publisher", ROBOT_STATE_PUBLISHER_CMD)
    print("✅ All core nodes launched")


def wait_for_fast_lio(node: FastLioOrchestraNode, timeout: float = FAST_LIO_WAIT_SEC):
    print(f"⏳ Waiting for Fast-LIO odometry (timeout {timeout}s)...")
    start = time.time()
    while not node.odometry_ready and (time.time() - start) < timeout and not stop_event.is_set():
        rclpy.spin_once(node, timeout_sec=0.5)
    if node.odometry_ready:
        print("✅ Fast-LIO odometry ready")
        return True
    print("⚠️  Fast-LIO odometry timeout; continuing anyway")
    return False


def start_trajectory_recording(node: FastLioOrchestraNode):
    global trajectory_logger
    if trajectory_logger:
        return
    trajectory_logger = TrajectoryLogger(TRAJECTORY_FILE, TRAJ_SAMPLE_HZ)
    trajectory_logger.start()
    node.attach_logger(trajectory_logger)


def stop_trajectory_recording():
    global trajectory_logger
    if trajectory_logger:
        trajectory_logger.stop()
        trajectory_logger = None
    if monitor_node_ref:
        monitor_node_ref.attach_logger(None)


def handle_connection_loss(reason: str):
    if stop_event.is_set():
        return
    with reverse_lock:
        print("\n" + "=" * 60)
        print("🚨 EMERGENCY SEQUENCE INITIATED")
        print(f"Reason: {reason}")
        print("=" * 60)
        stop_trajectory_recording()
        time.sleep(0.5)
        if not os.path.exists(TRAJECTORY_FILE):
            print(f"❌ Trajectory file missing: {TRAJECTORY_FILE}")
            send_direct_rtl()
            return
        file_size = os.path.getsize(TRAJECTORY_FILE)
        if file_size < TRAJ_MIN_BYTES:
            print(f"❌ Trajectory file too small ({file_size} bytes)")
            send_direct_rtl()
            return
        run_reverse_trajectory()
        print("=" * 60)


def run_reverse_trajectory():
    print("🔄 Launching reverse trajectory follower...")
    env = os.environ.copy()
    env["MP_TCP"] = MAVLINK_TARGET
    env["TRAJ_FILE"] = TRAJECTORY_FILE
    env.setdefault("CURRENT_ALT", "5")
    env.setdefault("TIME_SCALE", "1.0")
    try:
        result = subprocess.run(
            ["python3", REVERSE_SCRIPT],
            env=env,
            capture_output=True,
            text=True,
            timeout=300,
        )
        if result.returncode == 0:
            print("✅ Reverse trajectory completed successfully")
        else:
            print(f"⚠️  Reverse trajectory failed (code {result.returncode})")
            print(result.stdout)
            print(result.stderr)
            send_direct_rtl()
    except subprocess.TimeoutExpired:
        print("⚠️  Reverse follower timed out")
        send_direct_rtl()
    except FileNotFoundError:
        print(f"❌ Reverse script not found: {REVERSE_SCRIPT}")
        send_direct_rtl()
    except Exception as exc:
        print(f"❌ Reverse follower error: {exc}")
        send_direct_rtl()


def send_direct_rtl():
    print("📡 Sending direct RTL to autopilot...")
    try:
        master = mavutil.mavlink_connection(f"tcp:{MAVLINK_TARGET}", timeout=5)
        master.wait_heartbeat(timeout=5)
        master.mav.command_long_send(
            master.target_system,
            master.target_component,
            mavutil.mavlink.MAV_CMD_NAV_RETURN_TO_LAUNCH,
            0,
            0, 0, 0, 0, 0, 0, 0,
        )
        print("✅ RTL command sent")
    except Exception as exc:
        print(f"❌ Failed to send RTL: {exc}")


def monitor_ros_node(node: FastLioOrchestraNode):
    print("👀 Monitoring ROS 2 callbacks (Ctrl+C to exit)...")
    try:
        while not stop_event.is_set():
            rclpy.spin_once(node, timeout_sec=0.5)
    except KeyboardInterrupt:
        pass


def cleanup():
    print("\n🧹 Cleaning up processes and resources...")
    stop_trajectory_recording()
    if nav_bridge:
        nav_bridge.close()
    for proc in processes:
        if proc.poll() is None:
            proc.terminate()
    for proc in processes:
        try:
            proc.wait(timeout=2)
        except subprocess.TimeoutExpired:
            proc.kill()
    print("✅ Cleanup complete")


# =============================== MAIN ===============================

def main():
    global monitor_node_ref, nav_bridge
    print("=" * 60)
    print("🎼 FAST-LIO DRONE ORCHESTRA")
    print("=" * 60)
    print(f"Workspace: {WORKSPACE}")
    print(f"Fast-LIO topic: {FAST_LIO_TOPIC}")
    print(f"External navigation target: {MAVLINK_TARGET}")
    print(f"Trajectory file: {TRAJECTORY_FILE}")
    print("=" * 60 + "\n")

    nav_bridge = ExternalNavBridge(MAVLINK_TARGET, EXTERNAL_NAV_RATE_HZ)
    rclpy.init()
    node = FastLioOrchestraNode(nav_bridge)
    node.set_connection_loss_callback(handle_connection_loss)
    monitor_node_ref = node

    try:
        launch_ros_stack()
        time.sleep(5)
        wait_for_fast_lio(node)
        start_trajectory_recording(node)
        print("⏳ Waiting for RC signal...")
        rc_wait_start = time.time()
        while not node.rc_ever_received and (time.time() - rc_wait_start) < 30 and not stop_event.is_set():
            rclpy.spin_once(node, timeout_sec=0.5)
        if node.rc_ever_received:
            print("✅ RC signal detected")
            node.enable_monitoring()
        else:
            print("⚠️  No RC detected; monitoring disabled")
        monitor_ros_node(node)
    except KeyboardInterrupt:
        print("\n⏹️  Interrupted by user")
    except Exception as exc:
        print(f"\n❌ Fatal error: {exc}")
        import traceback
        traceback.print_exc()
    finally:
        stop_event.set()
        cleanup()
        node.destroy_node()
        rclpy.shutdown()
        print("\n✅ Orchestra shut down cleanly")


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    main()