import threading
from reachy_mini import ReachyMini, ReachyMiniApp
from reachy_mini.utils import create_head_pose
import numpy as np
import time
from pydantic import BaseModel


frequency = 100  # Hz
base_sleep_time = 1 / frequency


class Takina(ReachyMiniApp):
    # Optional: URL to a custom configuration page for the app
    # eg. "http://localhost:8042"
    custom_app_url: str | None = "http://0.0.0.0:8042"
    # Optional: specify a media backend ("gstreamer", "gstreamer_no_video", "default", etc.)
    # On the wireless, use gstreamer_no_video to optimise CPU usage if the app does not use video streaming
    request_media_backend: str | None = None

    def run(self, reachy_mini: ReachyMini, stop_event: threading.Event):
        t0 = time.time()

        antennas_enabled = True

        # You can ignore this part if you don't want to add settings to your app. If you set custom_app_url to None, you have to remove this part as well.
        # === vvv ===
        class AntennaState(BaseModel):
            enabled: bool

        @self.settings_app.post("/antennas")
        def update_antennas_state(state: AntennaState):
            nonlocal antennas_enabled
            antennas_enabled = state.enabled
            return {"antennas_enabled": antennas_enabled}

        # === ^^^ ===
        
        def get_antennas_deg():
            if antennas_enabled:
                amp_deg = 25.0
                a = amp_deg * np.sin(2.0 * np.pi * 0.5 * t)
                return np.array([a, -a])
            return np.array([0.0, 0.0])
            
        
        def roll_head_pose(t):
            roll_deg = 20.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(roll=roll_deg, degrees=True)
        
        def pitch_head_pose(t):
            pitch_deg = 20.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(pitch=pitch_deg, degrees=True)
        
        def yaw_head_pose(t):
            yaw_deg = 20.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(yaw=yaw_deg, degrees=True)
        
        def x_head_pose(t):
            x = 30.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(x=x, mm=True)
        
        def y_head_pose(t):
            y = 30.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(y=y, mm=True)
        
        def z_head_pose(t):
            z = 30.0 * np.sin(2.0 * np.pi * 0.5 * t)
            return create_head_pose(z=z, mm=True)
        
        def init_pos():
            reachy_mini.set_target(
                head = create_head_pose(roll=0, pitch=0, yaw=0, x=0, y=0, z=0, degrees=True, mm=True),
                antennas = np.array([0.0, 0.0]),
            )
        
        init_pos()

        # Main control loop
        while not stop_event.is_set():
            t = time.time() - t0

            antennas_deg_r = 0
            antennas_deg_l = 0

            # Antennas angle in rad
            antennas_rad_r = np.deg2rad(antennas_deg_r)
            antennas_rad_l = np.deg2rad(antennas_deg_l)
            antennas_pose = np.array([antennas_rad_l, -antennas_rad_r])
            
            # Head's position
            head_pose = create_head_pose(roll=0, pitch=0, yaw=0, x=0, y=0, z=0, degrees=True, mm=True)
            
            reachy_mini.set_target(
                head=head_pose,
                antennas=antennas_pose,
            )
            
            duration_loop = (time.time() - t0) - t
            sleep_time = base_sleep_time - duration_loop
            if sleep_time < 0:
                sleep_time = 0
            time.sleep(sleep_time)
        
        init_pos()


if __name__ == "__main__":
    app = Takina()
    try:
        app.wrapped_run()
    except KeyboardInterrupt:
        app.stop()