app.py

"""Reachy Mini Controller - Fun, Queue-Based Control Interface"""

import sys
import subprocess
import threading
import time
from dataclasses import dataclass
from typing import List, Optional

import gradio as gr
import numpy as np
import cv2

from reachy_mini import ReachyMini
from reachy_mini.daemon.backend.mujoco.video_udp import UDPJPEGFrameReceiver
from reachy_mini.utils import create_head_pose


@dataclass
class Movement:
    name: str
    x: float = 0
    y: float = 0
    z: float = 0
    roll: float = 0
    pitch: float = 0
    yaw: float = 0
    left_antenna: Optional[float] = None
    right_antenna: Optional[float] = None
    duration: float = 1.0


# Preset movements library
PRESET_MOVEMENTS = {
    "Home": Movement("Home", 0, 0, 0, 0, 0, 0, 0, 0),
    "Look Left": Movement("Look Left", 0, 0, 0, 0, 0, -30),
    "Look Right": Movement("Look Right", 0, 0, 0, 0, 0, 30),
    "Look Up": Movement("Look Up", 0, 0, 0, 0, -20, 0),
    "Look Down": Movement("Look Down", 0, 0, 0, 0, 15, 0),
    "Tilt Left": Movement("Tilt Left", 0, 0, 0, -20, 0, 0),
    "Tilt Right": Movement("Tilt Right", 0, 0, 0, 20, 0, 0),
    "Curious": Movement("Curious", 10, 0, 10, 15, -10, -15, 45, -45),
    "Excited": Movement("Excited", 0, 0, 20, 0, -15, 0, 90, 90),
    "Shy": Movement("Shy", -10, 0, -10, 10, 10, 20, -30, 30),
}

# Preset sequences
PRESET_SEQUENCES = {
    "Wave": ["Home", "Look Left", "Look Right", "Look Left", "Look Right", "Home"],
    "Nod": ["Home", "Look Down", "Look Up", "Look Down", "Home"],
    "Excited Dance": ["Home", "Excited", "Tilt Left", "Tilt Right", "Tilt Left", "Home"],
    "Look Around": ["Home", "Look Left", "Look Up", "Look Right", "Look Down", "Home"],
    "Curious Peek": ["Home", "Curious", "Look Right", "Look Left", "Home"],
}


class ReachyController:
    def __init__(self):
        self.daemon_process = None
        self.reachy_mini = None
        self.receiver = None
        self.frame_thread = None
        self.running = False
        self.frame_received_event = threading.Event()
        self.current_frame = np.zeros((1080, 1080, 3), dtype=np.uint8)
        
        self.movement_queue: List[Movement] = []
        self.is_playing = False
        self.playback_speed = 1.0
        self.play_thread = None
        self.auto_play = True  # Auto-play mode enabled by default
        
    def start_daemon(self):
        """Start the Reachy Mini daemon"""
        try:
            if self.daemon_process is not None:
                return "⚠️ Daemon already running"
            
            python_cmd = "mjpython" if sys.platform == "darwin" else "python"
            self.daemon_process = subprocess.Popen(
                [python_cmd, "-m", "reachy_mini.daemon.app.main", "--sim",
                 "--scene", "minimal", "--headless", "--stream-robot-view"],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            self.start_frame_listener()
            frame_received = self.frame_received_event.wait(timeout=10)
            
            if self.daemon_process.poll() is not None:
                return "❌ Daemon failed to start"
            
            return "✅ Daemon started" if frame_received else "⚠️ Daemon started (no video)"
        except Exception as e:
            return f"❌ Daemon error: {str(e)}"
    
    def start_frame_listener(self):
        if self.receiver is None:
            self.receiver = UDPJPEGFrameReceiver(listen_port=5010)
        
        if self.frame_thread and self.frame_thread.is_alive():
            return
        
        self.running = True
        self.frame_thread = threading.Thread(target=self._recv_loop, daemon=True)
        self.frame_thread.start()
    
    def _recv_loop(self):
        while self.running:
            frame = self.receiver.recv_frame()
            if frame is not None:
                # Resize frame from 640x640 to 1080x1080
                resized_frame = cv2.resize(frame, (1080, 1080), interpolation=cv2.INTER_CUBIC)
                self.current_frame = resized_frame
                self.frame_received_event.set()
    
    def stream_frames(self):
        while True:
            yield self.current_frame
            time.sleep(0.04)
    
    def initialize_robot(self):
        """Initialize robot connection"""
        try:
            if self.reachy_mini is not None:
                return "⚠️ Robot already connected"
            
            self.reachy_mini = ReachyMini(media_backend="default_no_video")
            return "✅ Robot connected - Ready to control!"
        except Exception as e:
            return f"❌ Connection failed: {str(e)}"
    
    def auto_start(self):
        """Auto-start daemon and robot connection"""
        status_msgs = []
        
        # Start daemon
        msg = self.start_daemon()
        status_msgs.append(msg)
        yield "\n".join(status_msgs)
        
        # Initialize robot
        msg = self.initialize_robot()
        status_msgs.append(msg)
        yield "\n".join(status_msgs)
    
    def restart_system(self):
        """Restart daemon and robot connection"""
        self.is_playing = False
        
        # Stop everything
        if self.daemon_process:
            self.daemon_process.terminate()
            self.daemon_process.wait(timeout=5)
            self.daemon_process = None
        
        self.running = False
        
        if self.frame_thread:
            self.frame_thread.join(timeout=2)
        
        if self.receiver:
            self.receiver.close()
            self.receiver = None
        
        if self.reachy_mini:
            try:
                self.reachy_mini.__exit__(None, None, None)
            except:
                pass
            self.reachy_mini = None
        
        yield "🔄 Restarting..."
        
        # Restart
        for status in self.auto_start():
            yield status
    
    def stop_all(self):
        """Stop everything"""
        self.is_playing = False
        
        if self.daemon_process:
            self.daemon_process.terminate()
            self.daemon_process.wait(timeout=5)
            self.daemon_process = None
        
        self.running = False
        
        if self.frame_thread:
            self.frame_thread.join()
        
        if self.receiver:
            self.receiver.close()
            self.receiver = None
        
        if self.reachy_mini:
            try:
                self.reachy_mini.__exit__(None, None, None)
            except:
                pass
            self.reachy_mini = None
        
        return "✅ Stopped"
    
    def add_to_queue(self, movement_name, x, y, z, roll, pitch, yaw, 
                     left_ant, right_ant, duration):
        """Add a movement to the queue"""
        movement = Movement(
            name=movement_name or f"Custom {len(self.movement_queue) + 1}",
            x=x, y=y, z=z,
            roll=roll, pitch=pitch, yaw=yaw,
            left_antenna=left_ant,
            right_antenna=right_ant,
            duration=duration
        )
        self.movement_queue.append(movement)
        
        # Auto-play if enabled and not already playing
        if self.auto_play and not self.is_playing:
            self._start_auto_play()
        
        return self.format_queue(), f"✅ Added: {movement.name}"
    
    def add_preset(self, preset_name):
        """Add a preset movement to queue"""
        if preset_name not in PRESET_MOVEMENTS:
            return self.format_queue(), f"❌ Unknown preset: {preset_name}"
        
        self.movement_queue.append(PRESET_MOVEMENTS[preset_name])
        
        # Auto-play if enabled and not already playing
        if self.auto_play and not self.is_playing:
            self._start_auto_play()
        
        return self.format_queue(), f"✅ Added: {preset_name}"
    
    def add_sequence(self, sequence_name):
        """Add a preset sequence to queue"""
        if sequence_name not in PRESET_SEQUENCES:
            return self.format_queue(), f"❌ Unknown sequence"
        
        for preset_name in PRESET_SEQUENCES[sequence_name]:
            self.movement_queue.append(PRESET_MOVEMENTS[preset_name])
        
        # Auto-play if enabled and not already playing
        if self.auto_play and not self.is_playing:
            self._start_auto_play()
        
        return self.format_queue(), f"✅ Added sequence: {sequence_name}"
    
    def clear_queue(self):
        """Clear the movement queue"""
        self.movement_queue.clear()
        self.is_playing = False  # Stop playback when clearing
        return self.format_queue(), "🗑️ Queue cleared"
    
    def remove_last(self):
        """Remove last movement from queue"""
        if self.movement_queue:
            removed = self.movement_queue.pop()
            return self.format_queue(), f"🗑️ Removed: {removed.name}"
        return self.format_queue(), "⚠️ Queue is empty"
    
    def format_queue(self):
        """Format queue for display"""
        if not self.movement_queue:
            return "📋 Queue is empty\n\nAdd movements using presets or custom controls"
        
        lines = ["📋 Movement Queue:\n"]
        total_duration = 0
        
        for i, mov in enumerate(self.movement_queue, 1):
            total_duration += mov.duration
            emoji = "▶️" if i == 1 else "⏸️"
            
            # Format head position
            head_str = f"Head: x={mov.x:.0f} y={mov.y:.0f} z={mov.z:.0f} r={mov.roll:.0f}° p={mov.pitch:.0f}° y={mov.yaw:.0f}°"
            
            # Format antennas if present
            ant_str = ""
            if mov.left_antenna is not None and mov.right_antenna is not None:
                ant_str = f"\n   Antennas: L={mov.left_antenna:.0f}° R={mov.right_antenna:.0f}°"
            
            lines.append(
                f"{emoji} {i}. {mov.name} ({mov.duration}s)\n"
                f"   {head_str}{ant_str}"
            )
        
        lines.append(f"\n⏱️ Total duration: {total_duration:.1f}s")
        lines.append(f"{'🔄 Auto-play: ON' if self.auto_play else '⏸️ Auto-play: OFF'}")
        return "\n".join(lines)
    
    def play_queue(self, speed):
        """Execute the movement queue"""
        if not self.movement_queue:
            return self.format_queue(), "⚠️ Queue is empty"
        
        if self.reachy_mini is None:
            return self.format_queue(), "❌ Robot not initialized"
        
        if self.is_playing:
            return self.format_queue(), "⚠️ Already playing"
        
        self.playback_speed = speed
        self.is_playing = True
        self.play_thread = threading.Thread(target=self._play_loop, daemon=True)
        self.play_thread.start()
        
        return self.format_queue(), f"▶️ Playing at {speed}x speed..."
    
    def _play_loop(self):
        """Background thread to execute movements"""
        try:
            current_index = 0
            while self.is_playing:
                # Check if there are movements to play
                if current_index < len(self.movement_queue):
                    movement = self.movement_queue[current_index]
                    
                    # Create pose
                    pose = create_head_pose(
                        x=movement.x, y=movement.y, z=movement.z,
                        roll=movement.roll, pitch=movement.pitch, yaw=movement.yaw,
                        degrees=True, mm=True
                    )
                    
                    # Adjust duration by playback speed
                    actual_duration = movement.duration / self.playback_speed
                    
                    
                    if movement.left_antenna is not None and movement.right_antenna is not None:
                        self.reachy_mini.goto_target(
                            head=pose,
                            antennas=[
                                np.deg2rad(movement.right_antenna),
                                np.deg2rad(movement.left_antenna)
                            ],
                            duration=actual_duration
                        )
                    else:
                        self.reachy_mini.goto_target(head=pose, duration=actual_duration)
                    
                    current_index += 1
                else:
                    # No more movements, stop if not in auto-play mode
                    if not self.auto_play:
                        break
                    # In auto-play mode, wait for new movements
                    time.sleep(0.1)
            
        except Exception as e:
            print(f"Error during playback: {e}")
        finally:
            self.is_playing = False
    
    def _start_auto_play(self):
        """Start auto-play mode"""
        if self.reachy_mini is None:
            return
        
        if not self.is_playing:
            self.is_playing = True
            self.play_thread = threading.Thread(target=self._play_loop, daemon=True)
            self.play_thread.start()
    
    def toggle_auto_play(self, enabled):
        """Toggle auto-play mode"""
        self.auto_play = enabled
        
        if self.auto_play and self.movement_queue and not self.is_playing:
            self._start_auto_play()
        
        return self.format_queue(), f"{'🔄 Auto-play enabled' if enabled else '⏸️ Auto-play disabled'}"
    
    def update_speed(self, speed):
        """Update playback speed in real-time"""
        self.playback_speed = speed
        return f"⚡ Speed: {speed}x"
    
    def stop_playback(self):
        """Stop current playback"""
        self.is_playing = False
        if self.play_thread:
            self.play_thread.join(timeout=2)
        
        # If auto-play is still enabled, inform user
        msg = "⏹️ Stopped"
        if self.auto_play:
            msg += " (auto-play still enabled)"
        
        return self.format_queue(), msg


# Create manager
manager = ReachyController()

# Build Gradio interface with improved layout
with gr.Blocks(title="Reachy Controller", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Reachy Mini Controller")
    gr.Markdown("Create fun movement sequences for your robot!")
    
    with gr.Row():
        # Left panel - Controls (narrower)
        with gr.Column(scale=2):
            # Status section
            gr.Markdown("### 📡 System Status")
            status = gr.Textbox(
                label="Status",
                lines=3,
                interactive=False,
                value="🔄 Initializing system..."
            )
            restart_btn = gr.Button("🔄 Restart System", variant="secondary", size="sm")
            
            gr.Markdown("### 🎮 Playback Controls")
            
            auto_play_toggle = gr.Checkbox(
                label="🔄 Auto-play",
                value=True,
                info="Execute movements automatically when added"
            )
            
            speed_slider = gr.Slider(
                0.25, 3.0, 1.0,
                label="⚡ Speed Multiplier",
                info="Adjust playback speed"
            )
            
            with gr.Row():
                play_btn = gr.Button("▶️ Play All", variant="primary", scale=2)
                stop_play_btn = gr.Button("⏹️ Stop", scale=1)
            
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear All")
                remove_btn = gr.Button("↶ Remove Last")
            
            # Queue display
            queue_display = gr.Textbox(
                label="📋 Movement Queue",
                lines=20,
                interactive=False,
                value=manager.format_queue()
            )
        
        # Right panel - Simulation view (larger and more square)
        with gr.Column(scale=3):
            sim_view = gr.Image(
                label="🎬 Robot Simulation",
                type="numpy",
                height=1080,
                width=1080,
                show_label=True
            )
    
    # Movement builder section below
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🎨 Quick Presets")
            
            with gr.Row():
                preset_btns = []
                for preset in list(PRESET_MOVEMENTS.keys())[:5]:
                    btn = gr.Button(preset, size="sm")
                    preset_btns.append((btn, preset))
            
            with gr.Row():
                for preset in list(PRESET_MOVEMENTS.keys())[5:]:
                    btn = gr.Button(preset, size="sm")
                    preset_btns.append((btn, preset))
        
        with gr.Column():
            gr.Markdown("### 🎬 Sequences")
            with gr.Row():
                sequence_dropdown = gr.Dropdown(
                    choices=list(PRESET_SEQUENCES.keys()),
                    label="Select Sequence",
                    value=None,
                    scale=3
                )
                add_seq_btn = gr.Button("➕ Add", scale=1)
    
    # Custom movement controls in accordion
    with gr.Accordion("🎯 Custom Movement Builder", open=False):
        custom_name = gr.Textbox(label="Movement Name", placeholder="My Move")
        
        with gr.Row():
            x = gr.Slider(-50, 50, 0, label="X (mm)", step=5)
            y = gr.Slider(-50, 50, 0, label="Y (mm)", step=5)
            z = gr.Slider(-20, 50, 0, label="Z (mm)", step=5)
        
        with gr.Row():
            roll = gr.Slider(-30, 30, 0, label="Roll (°)", step=5)
            pitch = gr.Slider(-30, 30, 0, label="Pitch (°)", step=5)
            yaw = gr.Slider(-45, 45, 0, label="Yaw (°)", step=5)
        
        with gr.Row():
            left_ant = gr.Slider(-180, 180, 0, label="Left Antenna (°)", step=15)
            right_ant = gr.Slider(-180, 180, 0, label="Right Antenna (°)", step=15)
        
        duration = gr.Slider(0.3, 3.0, 1.0, label="Duration (s)", step=0.1)
        
        add_custom_btn = gr.Button("➕ Add to Queue", variant="primary")
    
    # Auto-start on load
    demo.load(
        fn=manager.auto_start,
        outputs=[status]
    )
    
    # Stream video
    demo.load(fn=manager.stream_frames, outputs=sim_view)
    
    # Connect events - System control
    restart_btn.click(fn=manager.restart_system, outputs=[status])
    
    # Connect events - Playback control
    auto_play_toggle.change(
        fn=manager.toggle_auto_play,
        inputs=[auto_play_toggle],
        outputs=[queue_display, status]
    )
    
    speed_slider.change(
        fn=manager.update_speed,
        inputs=[speed_slider],
        outputs=[status]
    )
    
    play_btn.click(
        fn=manager.play_queue,
        inputs=[speed_slider],
        outputs=[queue_display, status]
    )
    stop_play_btn.click(fn=manager.stop_playback, outputs=[queue_display, status])
    clear_btn.click(fn=manager.clear_queue, outputs=[queue_display, status])
    remove_btn.click(fn=manager.remove_last, outputs=[queue_display, status])
    
    # Connect preset buttons
    for btn, preset_name in preset_btns:
        btn.click(
            fn=lambda p=preset_name: manager.add_preset(p),
            outputs=[queue_display, status]
        )
    
    # Connect sequence dropdown
    add_seq_btn.click(
        fn=manager.add_sequence,
        inputs=[sequence_dropdown],
        outputs=[queue_display, status]
    )
    
    # Connect custom movement
    add_custom_btn.click(
        fn=manager.add_to_queue,
        inputs=[custom_name, x, y, z, roll, pitch, yaw, left_ant, right_ant, duration],
        outputs=[queue_display, status]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)