"""The pose that crosses the driver boundary, and the tool-tip offset applied to it."""

from dataclasses import dataclass

import numpy as np
from scipy.spatial.transform import Rotation


@dataclass
class RobotPose:
    """A robot pose: position in mm, orientation in degrees, both in the base frame.

    The controller reports and accepts the flange pose. ``tcp_pose`` on the
    dataflow is this pose pushed out to the mounted tool's tip by
    :func:`with_tool_offset`.
    """

    x: float
    y: float
    z: float
    roll: float = 180.0
    pitch: float = 0.0
    yaw: float = 0.0

    @classmethod
    def from_sequence(cls, values) -> "RobotPose":
        """Build a pose from ``[x, y, z, roll, pitch, yaw]``; a position alone is enough."""
        return cls(**dict(zip(("x", "y", "z", "roll", "pitch", "yaw"), values, strict=False)))

    def as_list(self) -> list[float]:
        """Return the pose as ``[x, y, z, roll, pitch, yaw]``, the order the SDK uses."""
        return [self.x, self.y, self.z, self.roll, self.pitch, self.yaw]


def rotation_matrix(pose: RobotPose) -> np.ndarray:
    """Return the 3x3 base-frame rotation of ``pose``."""
    return Rotation.from_euler("xyz", [pose.roll, pose.pitch, pose.yaw], degrees=True).as_matrix()


def with_tool_offset(pose: RobotPose, offset_mm: float) -> RobotPose:
    """Return ``pose`` moved ``offset_mm`` along its own Z axis.

    The offset is along the *tool* Z, not the base Z, so a tilted flange puts the
    tip where the tool actually reaches rather than straight down from it.
    Orientation is unchanged: sliding along the axis does not turn anything.
    """
    if offset_mm == 0.0:
        return pose
    tip = np.array([pose.x, pose.y, pose.z]) + rotation_matrix(pose) @ np.array([0.0, 0.0, offset_mm])
    return RobotPose(tip[0], tip[1], tip[2], pose.roll, pose.pitch, pose.yaw)