# Copyright (c) 2025, Yutang Li, SIAT (yt.li2@siat.ac.cn)
# SPDX-License-Identifier: BSD-3-Clause

"""Head tracking via XR headset (HMD) for robot head joint control."""

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

from xr_utils import transform_xr_pose, is_valid_quaternion


class HeadTracker:
    """Extracts yaw/pitch from HMD pose relative to a calibrated neutral orientation."""

    def __init__(self, yaw_limit: float = 1.57, pitch_limit: float = 0.52):
        self.yaw_limit = yaw_limit      # ±90°
        self.pitch_limit = pitch_limit   # ±30°
        self.neutral_rot: R | None = None

    def reset(self):
        """Clear neutral calibration; next call to get_targets() will re-calibrate."""
        self.neutral_rot = None

    def get_targets(self, xr_client) -> np.ndarray:
        """Return [yaw, pitch] in radians relative to neutral HMD pose.

        First call after reset() records neutral orientation and returns zeros.
        """
        try:
            raw_pose = xr_client.get_pose("headset")
            if not is_valid_quaternion(raw_pose[3:]):
                return np.zeros(2, dtype=np.float32)
            _, head_world_quat_xyzw = transform_xr_pose(raw_pose[:3], raw_pose[3:])
            head_rot = R.from_quat(head_world_quat_xyzw)
            if self.neutral_rot is None:
                self.neutral_rot = head_rot
                return np.zeros(2, dtype=np.float32)
            rel_rot = head_rot * self.neutral_rot.inv()
            euler_zyx = rel_rot.as_euler("ZYX")
            yaw = float(np.clip(euler_zyx[0], -self.yaw_limit, self.yaw_limit))
            pitch = float(np.clip(euler_zyx[1], -self.pitch_limit, self.pitch_limit))
            return np.array([yaw, pitch], dtype=np.float32)
        except Exception:
            return np.zeros(2, dtype=np.float32)