"""
RobotState/emotes.py - Expressive Emote States (Wave, Laola Wave)
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import numpy as np
import math

from config import cfg
import DataTypes as dt
from states.State import State
if TYPE_CHECKING:
    from Robot import Robot


class WaveEmoteState(State):
    """Front leg wave greeting emote state."""

    def __init__(
        self,
        cycles: int = 3,
        duration: float = 1.5,
        tickpersec: float = 20.0,
        height: float = 40.0,
        amplitude: float = 25.0,
    ):
        self.cycles = cycles
        self.duration = duration
        self.tickpersec = tickpersec
        self.ticks = int(duration * tickpersec)
        self.height = height
        self.amplitude = amplitude

        self.current_cycle = 0
        self.current_tick = 0
        self.base_pos: Optional[dt.PosArray] = None

    def enter(self, robot: Robot) -> None:
        self.current_cycle = 0
        self.current_tick = 0
        self.base_pos = dt.PosArray(np.copy(robot.current_pos.data))

    def execute(self, robot: Robot) -> Optional[str]:
        t = self.current_tick / float(self.ticks)
        tick_pos = np.copy(self.base_pos.data)
        leg_id = 0  # Front leg

        cx, cy, cz = self.base_pos[leg_id]

        # Wave trajectory calculation
        y_wave = math.sin(4.0 * math.pi * t)
        y_offset = self.amplitude * (0.5 * (y_wave + 1.0))
        z_offset = self.height * math.sin(math.pi * t)

        max_y_dev = 30.0
        new_y = cx
        new_y = max(cy - max_y_dev, min(cy + max_y_dev, cy + y_offset))

        tick_pos[leg_id] = [cx + 10.0, new_y, cz + z_offset]

        pos_array = dt.PosArray(tick_pos)
        target_rad = robot.compute_ik(pos_array)

        robot.current_pos = pos_array
        robot.set_joint_angles(target_rad)

        self.current_tick += 1

        if self.current_tick >= self.ticks:
            self.current_tick = 0
            self.current_cycle += 1

        if self.current_cycle >= self.cycles:
            return "idle"

        return None

    def exit(self, robot: Robot) -> None:
        robot.robot_state = "idle"


class LaolaWaveEmoteState(State):
    """Laola Wave side-to-side leg wave emote state."""

    def __init__(
        self,
        cycles: int = 3,
        duration: float = 1.5,
        tickpersec: float = cfg.standard_tickpersec,
        height: float = 10.0,
        amplitude: float = 30.0,
    ):
        self.cycles = cycles
        self.duration = duration
        self.ticks = int(duration * tickpersec)
        self.height = height
        self.amplitude = amplitude

        self.current_cycle = 0
        self.current_tick = 0
        self.base_pos: Optional[dt.PosArray] = None

    def enter(self, robot: Robot) -> None:
        self.current_cycle = 0
        self.current_tick = 0
        self.base_pos = dt.PosArray(np.copy(robot.current_pos.data))

    def execute(self, robot: Robot) -> Optional[str]:
        t = self.current_tick / float(self.ticks)
        tick_pos = np.copy(self.base_pos.data)
        wave_legs = [1, 4]

        for leg_id in wave_legs:
            cx, cy, cz = self.base_pos[leg_id]
            phase = 0.0 if leg_id == 1 else math.pi

            y_offset = self.amplitude * math.sin(2.0 * math.pi * t + phase)
            z_offset = self.height * math.sin(2.0 * math.pi * t + phase)

            tick_pos[leg_id] = [cx, cy + y_offset, cz + z_offset]

        pos_array = dt.PosArray(tick_pos)
        target_rad = robot.compute_ik(pos_array)

        robot.current_pos = pos_array
        robot.set_joint_angles(target_rad)

        self.current_tick += 1

        if self.current_tick >= self.ticks:
            self.current_tick = 0
            self.current_cycle += 1

        if self.current_cycle >= self.cycles:
            return "idle"

        return None

    def exit(self, robot: Robot) -> None:
        robot.robot_state = "idle"