"""
Transform 数据类：四元数/旋转矩阵、逆变换、序列化。
不依赖 manager，供 yaml_io、transform_math、manager 使用。
"""
import numpy as np
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass


@dataclass
class Transform:
    """表示一个坐标变换"""
    source_frame: str
    target_frame: str
    translation: np.ndarray  # [x, y, z]
    rotation: np.ndarray     # [qx, qy, qz, qw] 四元数

    def to_matrix(self) -> np.ndarray:
        """转换为4x4齐次变换矩阵"""
        qx, qy, qz, qw = self.rotation

        # 四元数转旋转矩阵
        R = np.array([
            [1 - 2*(qy**2 + qz**2), 2*(qx*qy - qz*qw), 2*(qx*qz + qy*qw)],
            [2*(qx*qy + qz*qw), 1 - 2*(qx**2 + qz**2), 2*(qy*qz - qx*qw)],
            [2*(qx*qz - qy*qw), 2*(qy*qz + qx*qw), 1 - 2*(qx**2 + qy**2)]
        ])

        # 构建齐次变换矩阵
        T = np.eye(4)
        T[:3, :3] = R
        T[:3, 3] = self.translation
        return T

    @staticmethod
    def from_matrix(matrix: np.ndarray, source: str, target: str) -> "Transform":
        """从4x4齐次变换矩阵创建Transform"""
        translation = matrix[:3, 3]
        R = matrix[:3, :3]

        # 旋转矩阵转四元数
        trace = np.trace(R)
        if trace > 0:
            s = 0.5 / np.sqrt(trace + 1.0)
            qw = 0.25 / s
            qx = (R[2, 1] - R[1, 2]) * s
            qy = (R[0, 2] - R[2, 0]) * s
            qz = (R[1, 0] - R[0, 1]) * s
        else:
            if R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
                s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
                qw = (R[2, 1] - R[1, 2]) / s
                qx = 0.25 * s
                qy = (R[0, 1] + R[1, 0]) / s
                qz = (R[0, 2] + R[2, 0]) / s
            elif R[1, 1] > R[2, 2]:
                s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
                qw = (R[0, 2] - R[2, 0]) / s
                qx = (R[0, 1] + R[1, 0]) / s
                qy = 0.25 * s
                qz = (R[1, 2] + R[2, 1]) / s
            else:
                s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
                qw = (R[1, 0] - R[0, 1]) / s
                qx = (R[0, 2] + R[2, 0]) / s
                qy = (R[1, 2] + R[2, 1]) / s
                qz = 0.25 * s

        rotation = np.array([qx, qy, qz, qw])
        return Transform(source, target, translation, rotation)

    def inverse(self) -> "Transform":
        """返回逆变换"""
        T_inv = np.linalg.inv(self.to_matrix())
        return Transform.from_matrix(T_inv, self.target_frame, self.source_frame)

    def to_dict(self) -> dict:
        """转换为字典格式（用于保存YAML）"""
        return {
            'source_frame': self.source_frame,
            'target_frame': self.target_frame,
            'translation': {
                'x': float(self.translation[0]),
                'y': float(self.translation[1]),
                'z': float(self.translation[2])
            },
            'rotation': {
                'qx': float(self.rotation[0]),
                'qy': float(self.rotation[1]),
                'qz': float(self.rotation[2]),
                'qw': float(self.rotation[3])
            }
        }
