"""
Feature extraction for sign language segmentation.

Computes:
- Angle features (104 dims): body (22) + right hand (41) + left hand (41)
- HaMeR features (288 dims): MANO hand pose rotation matrices flattened

These features were used by the legacy segmenter. The v2 segmenter uses WiLoR H5 features directly.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np


def compute_angle_features(
    body_poses: np.ndarray,
    right_hand_poses: np.ndarray,
    left_hand_poses: np.ndarray,
    mano_right: dict[str, np.ndarray] | None = None,
    mano_left: dict[str, np.ndarray] | None = None,
) -> np.ndarray:
    """
    Compute angle features for segmentation (104 dimensions).

    Args:
        body_poses: MediaPipe body keypoints (T, 33, 3)
        right_hand_poses: Right hand keypoints (T, 21, 3)
        left_hand_poses: Left hand keypoints (T, 21, 3)
        mano_right: Optional MANO params for right hand with 'hand_pose' key
        mano_left: Optional MANO params for left hand with 'hand_pose' key

    Returns:
        Angle features array (T, 104)
    """
    # Body angles (22 dims) - computed from body pose joint angles
    body_angles = compute_body_angles(body_poses)  # (T, 22)

    # Hand angles (41 dims each) - from MANO or computed from keypoints
    if mano_right is not None and "hand_pose" in mano_right:
        right_angles = mano_to_euler_angles(mano_right["hand_pose"])
    else:
        right_angles = compute_hand_angles(right_hand_poses)

    if mano_left is not None and "hand_pose" in mano_left:
        left_angles = mano_to_euler_angles(mano_left["hand_pose"])
    else:
        left_angles = compute_hand_angles(left_hand_poses)

    # Ensure all have same length
    min_len = min(len(body_angles), len(right_angles), len(left_angles))
    body_angles = body_angles[:min_len]
    right_angles = right_angles[:min_len]
    left_angles = left_angles[:min_len]

    # Concatenate: body (22) + right (41) + left (41) = 104
    angle_features = np.concatenate([body_angles, right_angles, left_angles], axis=-1)

    return angle_features.astype(np.float32)


def compute_body_angles(body_poses: np.ndarray) -> np.ndarray:
    """
    Compute body joint angles from MediaPipe body keypoints.

    Computes relative angles between connected body joints.

    Args:
        body_poses: MediaPipe body keypoints (T, 33, 3)

    Returns:
        Body angle features (T, 22)
    """
    T = len(body_poses)

    # MediaPipe body joint connections for angle computation
    # Format: (parent, joint, child) - compute angle at joint
    body_joints = [
        # Torso/spine
        (11, 23, 25),  # left hip angle
        (12, 24, 26),  # right hip angle
        (23, 11, 13),  # left shoulder angle
        (24, 12, 14),  # right shoulder angle
        # Arms
        (11, 13, 15),  # left elbow
        (12, 14, 16),  # right elbow
        (13, 15, 17),  # left wrist
        (14, 16, 18),  # right wrist
        # Legs (optional but helps with body orientation)
        (23, 25, 27),  # left knee
        (24, 26, 28),  # right knee
        (25, 27, 29),  # left ankle
        (26, 28, 30),  # right ankle
    ]

    angles = np.zeros((T, 22), dtype=np.float32)

    for i, (parent, joint, child) in enumerate(body_joints):
        if (
            parent < body_poses.shape[1]
            and joint < body_poses.shape[1]
            and child < body_poses.shape[1]
        ):
            v1 = body_poses[:, parent, :] - body_poses[:, joint, :]
            v2 = body_poses[:, child, :] - body_poses[:, joint, :]

            # Compute angle between vectors
            cos_angle = np.sum(v1 * v2, axis=-1) / (
                np.linalg.norm(v1, axis=-1) * np.linalg.norm(v2, axis=-1) + 1e-8
            )
            cos_angle = np.clip(cos_angle, -1, 1)
            angles[:, i] = np.arccos(cos_angle)

    # Fill remaining dims with velocity/orientation features
    # Shoulder orientation (roll, pitch, yaw approximations)
    left_shoulder = body_poses[:, 11, :]
    right_shoulder = body_poses[:, 12, :]
    shoulder_vec = right_shoulder - left_shoulder

    # Normalize and use components as features
    shoulder_norm = np.linalg.norm(shoulder_vec, axis=-1, keepdims=True) + 1e-8
    shoulder_unit = shoulder_vec / shoulder_norm
    angles[:, 12:15] = shoulder_unit  # 3 dims

    # Hip orientation
    left_hip = body_poses[:, 23, :]
    right_hip = body_poses[:, 24, :]
    hip_vec = right_hip - left_hip
    hip_norm = np.linalg.norm(hip_vec, axis=-1, keepdims=True) + 1e-8
    hip_unit = hip_vec / hip_norm
    angles[:, 15:18] = hip_unit  # 3 dims

    # Wrist positions relative to shoulders (normalized)
    left_wrist = body_poses[:, 15, :]
    right_wrist = body_poses[:, 16, :]
    angles[:, 18:20] = (left_wrist[:, :2] - left_shoulder[:, :2]) / shoulder_norm  # 2 dims
    angles[:, 20:22] = (right_wrist[:, :2] - right_shoulder[:, :2]) / shoulder_norm  # 2 dims

    return angles


def compute_hand_angles(hand_poses: np.ndarray) -> np.ndarray:
    """
    Compute hand joint angles from 21 keypoints.

    Args:
        hand_poses: Hand keypoints (T, 21, 3)

    Returns:
        Hand angle features (T, 41)
    """
    T = len(hand_poses)
    angles = np.zeros((T, 41), dtype=np.float32)

    # Hand joint connections (21 keypoints, ~20 bones)
    # Each finger: MCP, PIP, DIP, TIP
    finger_bases = [1, 5, 9, 13, 17]  # MCP joints

    idx = 0
    for finger_idx, base in enumerate(finger_bases):
        # Finger joints: base, base+1, base+2, base+3
        joints = (
            [0, base, base + 1, base + 2, base + 3]
            if finger_idx == 0
            else [base, base + 1, base + 2, base + 3]
        )

        for i in range(len(joints) - 2):
            if idx >= 41:
                break
            parent, joint, child = joints[i], joints[i + 1], joints[i + 2]

            v1 = hand_poses[:, parent, :] - hand_poses[:, joint, :]
            v2 = hand_poses[:, child, :] - hand_poses[:, joint, :]

            # Angle between vectors
            cos_angle = np.sum(v1 * v2, axis=-1) / (
                np.linalg.norm(v1, axis=-1) * np.linalg.norm(v2, axis=-1) + 1e-8
            )
            cos_angle = np.clip(cos_angle, -1, 1)
            angles[:, idx] = np.arccos(cos_angle)
            idx += 1

            # Also compute rotation components (spread across 3 dims per joint)
            if idx + 2 < 41:
                cross = np.cross(v1, v2)
                cross_norm = np.linalg.norm(cross, axis=-1, keepdims=True) + 1e-8
                angles[:, idx : idx + 3] = cross / cross_norm
                idx += 3

    # Wrist rotation relative to palm (global orientation)
    wrist = hand_poses[:, 0, :]
    middle_mcp = hand_poses[:, 9, :]
    palm_vec = middle_mcp - wrist
    palm_norm = np.linalg.norm(palm_vec, axis=-1, keepdims=True) + 1e-8

    remaining = 41 - idx
    if remaining > 0:
        angles[:, idx : idx + min(3, remaining)] = (palm_vec / palm_norm)[:, : min(3, remaining)]

    return angles


def compute_body_angles_from_smplx(joints2d: np.ndarray) -> np.ndarray:
    """
    Compute body joint angles from SMPL-X 55-joint positions (2D or 3D).

    Maps SMPL-X body joint indices to the same angle features that
    ``compute_body_angles`` produces from MediaPipe landmarks.

    SMPL-X body joint order (first 22):
        0=pelvis, 1=L_hip, 2=R_hip, 3=spine1, 4=L_knee, 5=R_knee,
        6=spine2, 7=L_ankle, 8=R_ankle, 9=spine3, 10=L_foot, 11=R_foot,
        12=neck, 13=L_collar, 14=R_collar, 15=head,
        16=L_shoulder, 17=R_shoulder, 18=L_elbow, 19=R_elbow,
        20=L_wrist, 21=R_wrist

    Args:
        joints2d: SMPL-X joint positions (T, 55, 2) or (T, 55, 3).
            Only the first 22 body joints are used, plus hand joint 25
            (left index MCP) and 40 (right index MCP) as wrist targets.

    Returns:
        Body angle features (T, 22), matching ``compute_body_angles`` output.
    """
    T = len(joints2d)
    ndim = joints2d.shape[-1]  # 2 or 3

    # Pad to 3D if needed (set z=0)
    if ndim == 2:
        joints3d = np.zeros((T, joints2d.shape[1], 3), dtype=np.float32)
        joints3d[:, :, :2] = joints2d
    else:
        joints3d = joints2d.astype(np.float32)

    # SMPL-X equivalents of the MediaPipe angle triplets (parent, joint, child)
    smplx_angle_triplets = [
        (3, 1, 4),  # left hip angle: spine1 -> L_hip -> L_knee
        (3, 2, 5),  # right hip angle: spine1 -> R_hip -> R_knee
        (1, 16, 18),  # left shoulder angle: L_hip -> L_shoulder -> L_elbow
        (2, 17, 19),  # right shoulder angle: R_hip -> R_shoulder -> R_elbow
        (16, 18, 20),  # left elbow: L_shoulder -> L_elbow -> L_wrist
        (17, 19, 21),  # right elbow: R_shoulder -> R_elbow -> R_wrist
        (18, 20, 25),  # left wrist: L_elbow -> L_wrist -> L_index_MCP
        (19, 21, 40),  # right wrist: R_elbow -> R_wrist -> R_index_MCP
        (1, 4, 7),  # left knee: L_hip -> L_knee -> L_ankle
        (2, 5, 8),  # right knee: R_hip -> R_knee -> R_ankle
        (4, 7, 10),  # left ankle: L_knee -> L_ankle -> L_foot
        (5, 8, 11),  # right ankle: R_knee -> R_ankle -> R_foot
    ]

    angles = np.zeros((T, 22), dtype=np.float32)

    for i, (parent, joint, child) in enumerate(smplx_angle_triplets):
        if parent >= joints3d.shape[1] or joint >= joints3d.shape[1] or child >= joints3d.shape[1]:
            continue
        v1 = joints3d[:, parent, :] - joints3d[:, joint, :]
        v2 = joints3d[:, child, :] - joints3d[:, joint, :]

        cos_angle = np.sum(v1 * v2, axis=-1) / (
            np.linalg.norm(v1, axis=-1) * np.linalg.norm(v2, axis=-1) + 1e-8
        )
        cos_angle = np.clip(cos_angle, -1, 1)
        angles[:, i] = np.arccos(cos_angle)

    # Orientation features (same dims 12-21 as MediaPipe version)
    left_shoulder = joints3d[:, 16, :]
    right_shoulder = joints3d[:, 17, :]
    shoulder_vec = right_shoulder - left_shoulder
    shoulder_norm = np.linalg.norm(shoulder_vec, axis=-1, keepdims=True) + 1e-8
    shoulder_unit = shoulder_vec / shoulder_norm
    angles[:, 12:15] = shoulder_unit

    left_hip = joints3d[:, 1, :]
    right_hip = joints3d[:, 2, :]
    hip_vec = right_hip - left_hip
    hip_norm = np.linalg.norm(hip_vec, axis=-1, keepdims=True) + 1e-8
    hip_unit = hip_vec / hip_norm
    angles[:, 15:18] = hip_unit

    left_wrist = joints3d[:, 20, :]
    right_wrist = joints3d[:, 21, :]
    angles[:, 18:20] = (left_wrist[:, :2] - left_shoulder[:, :2]) / shoulder_norm
    angles[:, 20:22] = (right_wrist[:, :2] - right_shoulder[:, :2]) / shoulder_norm

    return angles


def load_features_from_nlf_wilor(
    nlf_path: str | Path,
    wilor_path: str | Path,
    num_frames: int | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Load and compute segmentation features from NLF body + WiLoR hand H5 files.

    This is the equivalent of ``load_features_from_h5`` but for BT2's
    NLF (SMPL-X body) + WiLoR (MANO hands) extraction pipeline instead of
    the MediaPipe + WiLoR pipeline.

    Args:
        nlf_path: Path to NLF body H5 file (contains ``joints2d``).
        wilor_path: Path to WiLoR hand H5 file (contains ``mano/`` group).
        num_frames: Expected number of video frames (for alignment).

    Returns:
        Tuple of (angle_features, hamer_features) ready for ``SignSegmenter``.
    """
    import h5py

    nlf_path = Path(nlf_path)
    wilor_path = Path(wilor_path)

    # --- Load NLF body joints ---
    with h5py.File(nlf_path, "r") as f:
        joints2d = f["joints2d"][:].astype(np.float32)  # (N_det, 55, 2)
        frame_idx = f["frame_idx"][:] if "frame_idx" in f else None
        img_idx = f["img_idx"][:] if "img_idx" in f else None

        if num_frames is None:
            if frame_idx is not None:
                num_frames = len(frame_idx)
            elif img_idx is not None:
                num_frames = int(img_idx.max()) + 1
            else:
                num_frames = len(joints2d)

        # Build dense (T, 55, 2) array — take first detection per frame
        dense_joints = np.zeros((num_frames, 55, 2), dtype=np.float32)
        if frame_idx is not None:
            from sltk.io.h5 import get_frame_detections as _gfd

            for t in range(num_frames):
                start, end = _gfd(f, t, frame_idx)
                if start >= 0 and start < len(joints2d):
                    dense_joints[t] = joints2d[start]
        elif img_idx is not None:
            for det_idx in range(len(img_idx)):
                t = int(img_idx[det_idx])
                if t < num_frames and np.all(dense_joints[t] == 0):
                    dense_joints[t] = joints2d[det_idx]

    # Compute body angles from SMPL-X joints
    body_angles = compute_body_angles_from_smplx(dense_joints)  # (T, 22)

    # --- Load WiLoR MANO parameters ---
    with h5py.File(wilor_path, "r") as f:
        w_frame_idx = f["frame_idx"][:] if "frame_idx" in f else None
        w_img_idx = f["img_idx"][:] if "img_idx" in f else None
        right_mask = f["right"][:] if "right" in f else None

        # Initialize per-frame MANO arrays
        mano_hand_pose_right = np.zeros((num_frames, 15, 3, 3), dtype=np.float32)
        mano_hand_pose_left = np.zeros((num_frames, 15, 3, 3), dtype=np.float32)
        mano_go_right = np.tile(
            np.eye(3, dtype=np.float32)[np.newaxis, np.newaxis], (num_frames, 1, 1, 1)
        )
        mano_go_left = np.tile(
            np.eye(3, dtype=np.float32)[np.newaxis, np.newaxis], (num_frames, 1, 1, 1)
        )

        if "mano" in f:
            hp_data = f["mano/hand_pose"][:].astype(np.float32)
            go_data = f["mano/global_orient"][:].astype(np.float32)

            def _assign(det_idx: int, frame: int) -> None:
                is_right = right_mask[det_idx] if right_mask is not None else True
                if is_right:
                    mano_hand_pose_right[frame] = hp_data[det_idx]
                    mano_go_right[frame] = go_data[det_idx]
                else:
                    mano_hand_pose_left[frame] = hp_data[det_idx]
                    mano_go_left[frame] = go_data[det_idx]

            if w_frame_idx is not None:
                from sltk.io.h5 import get_frame_detections as _gfd

                for t in range(num_frames):
                    start, end = _gfd(f, t, w_frame_idx)
                    if start >= 0:
                        for i in range(start, min(end, len(hp_data))):
                            _assign(i, t)
            elif w_img_idx is not None:
                for det_idx in range(len(w_img_idx)):
                    t = int(w_img_idx[det_idx])
                    if t < num_frames and det_idx < len(hp_data):
                        _assign(det_idx, t)

    # Hand angles from MANO rotation matrices
    right_angles = mano_to_euler_angles(mano_hand_pose_right)  # (T, 41)
    left_angles = mano_to_euler_angles(mano_hand_pose_left)  # (T, 41)

    angle_features = np.concatenate([body_angles, right_angles, left_angles], axis=-1)
    angle_features = angle_features.astype(np.float32)

    # HaMeR features from flattened rotation matrices
    hamer_features = compute_hamer_features(
        mano_go_right,
        mano_hand_pose_right,
        mano_go_left,
        mano_hand_pose_left,
    )

    return angle_features, hamer_features


def mano_to_euler_angles(hand_pose: np.ndarray) -> np.ndarray:
    """
    Convert MANO hand pose rotation matrices to Euler angles.

    Args:
        hand_pose: MANO hand pose rotations (T, 15, 3, 3) or (N, 15, 3, 3)

    Returns:
        Euler angles (T, 41) - 15 joints × ~3 angles each ≈ 41
    """
    if hand_pose.ndim == 3:
        hand_pose = hand_pose[np.newaxis, ...]

    T = hand_pose.shape[0]
    angles = np.zeros((T, 41), dtype=np.float32)

    # Convert each rotation matrix to Euler angles (ZYX convention)
    for t in range(T):
        idx = 0
        for j in range(min(15, hand_pose.shape[1])):
            if idx + 3 > 41:
                break
            R = hand_pose[t, j]

            # Extract Euler angles from rotation matrix
            sy = np.sqrt(R[0, 0] ** 2 + R[1, 0] ** 2)
            singular = sy < 1e-6

            if not singular:
                x = np.arctan2(R[2, 1], R[2, 2])
                y = np.arctan2(-R[2, 0], sy)
                z = np.arctan2(R[1, 0], R[0, 0])
            else:
                x = np.arctan2(-R[1, 2], R[1, 1])
                y = np.arctan2(-R[2, 0], sy)
                z = 0

            # Store angles (clamped for numerical stability)
            if idx < 41:
                angles[t, idx] = np.clip(x, -np.pi, np.pi)
                idx += 1
            if idx < 41:
                angles[t, idx] = np.clip(y, -np.pi, np.pi)
                idx += 1
            if idx < 41:
                angles[t, idx] = np.clip(z, -np.pi, np.pi)
                idx += 1

    return angles


def compute_hamer_features(
    mano_global_orient_right: np.ndarray,
    mano_hand_pose_right: np.ndarray,
    mano_global_orient_left: np.ndarray,
    mano_hand_pose_left: np.ndarray,
) -> np.ndarray:
    """
    Compute HaMeR features from MANO parameters (288 dimensions).

    HaMeR features are flattened rotation matrices:
    - 2 hands × 15 joints × 3×3 rotation = 270
    - 2 hands × 1 global × 3×3 rotation = 18
    - Total: 288

    Args:
        mano_global_orient_right: Right hand global orientation (T, 1, 3, 3)
        mano_hand_pose_right: Right hand joint rotations (T, 15, 3, 3)
        mano_global_orient_left: Left hand global orientation (T, 1, 3, 3)
        mano_hand_pose_left: Left hand joint rotations (T, 15, 3, 3)

    Returns:
        HaMeR features (T, 288)
    """
    T = mano_hand_pose_right.shape[0]

    # Flatten rotation matrices
    # hand_pose: (T, 15, 3, 3) -> (T, 135)
    # global_orient: (T, 1, 3, 3) -> (T, 9)

    hand_pose_right_flat = mano_hand_pose_right.reshape(T, -1)  # (T, 135)
    hand_pose_left_flat = mano_hand_pose_left.reshape(T, -1)  # (T, 135)
    global_orient_right_flat = mano_global_orient_right.reshape(T, -1)  # (T, 9)
    global_orient_left_flat = mano_global_orient_left.reshape(T, -1)  # (T, 9)

    # Concatenate: [right_hand_pose, left_hand_pose, right_global, left_global]
    # = 135 + 135 + 9 + 9 = 288
    hamer_features = np.concatenate(
        [
            hand_pose_right_flat,
            hand_pose_left_flat,
            global_orient_right_flat,
            global_orient_left_flat,
        ],
        axis=-1,
    )

    return hamer_features.astype(np.float32)


def load_features_from_h5(
    mediapipe_path: str | Path,
    wilor_path: str | Path,
    num_frames: int | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Load and compute features from MediaPipe and WiLoR H5 files.

    Args:
        mediapipe_path: Path to MediaPipe H5 file
        wilor_path: Path to WiLoR H5 file
        num_frames: Expected number of frames (for alignment)

    Returns:
        Tuple of (angle_features, hamer_features)
    """
    import h5py

    try:
        import hdf5plugin  # noqa: F401 — registers LZ4/BLOSC HDF5 filters
    except ImportError:
        pass  # Optional: only needed for compressed H5 files (e.g. from WiLoR Docker)

    mediapipe_path = Path(mediapipe_path)
    wilor_path = Path(wilor_path)

    # Load MediaPipe body poses
    with h5py.File(mediapipe_path, "r") as f:
        if "poses" in f:
            body_poses = f["poses"][:]
        elif "body" in f:
            body_poses = f["body"][:]
        else:
            # Try to find body poses in common locations
            body_poses = None
            for key in ["body_poses", "pose_landmarks", "landmarks"]:
                if key in f:
                    body_poses = f[key][:]
                    break
            if body_poses is None:
                raise ValueError(f"Could not find body poses in {mediapipe_path}")

        # Get hand poses if available
        left_hand = f["left_hand"][:] if "left_hand" in f else np.zeros((len(body_poses), 21, 3))
        right_hand = f["right_hand"][:] if "right_hand" in f else np.zeros((len(body_poses), 21, 3))

    # Load WiLoR MANO parameters
    with h5py.File(wilor_path, "r") as f:
        # Support both formats:
        #   frame_idx: shape (num_frames, 2), each entry is (start_idx, count)
        #   img_idx: shape (num_detections,), each entry is the frame number
        frame_idx = f["frame_idx"][:] if "frame_idx" in f else None
        img_idx = f["img_idx"][:] if "img_idx" in f else None

        # Get number of video frames
        if num_frames is None:
            if frame_idx is not None:
                num_frames = len(frame_idx)
            elif img_idx is not None:
                num_frames = int(img_idx.max()) + 1
            else:
                num_frames = len(body_poses)

        # Initialize arrays for per-frame MANO params
        mano_hand_pose_right = np.zeros((num_frames, 15, 3, 3), dtype=np.float32)
        mano_hand_pose_left = np.zeros((num_frames, 15, 3, 3), dtype=np.float32)
        mano_global_orient_right = np.eye(3, dtype=np.float32)[np.newaxis, np.newaxis, ...].repeat(
            num_frames, axis=0
        )
        mano_global_orient_left = np.eye(3, dtype=np.float32)[np.newaxis, np.newaxis, ...].repeat(
            num_frames, axis=0
        )

        # Load MANO parameters
        if "mano" in f:
            mano_group = f["mano"]
            hand_pose_data = (
                mano_group["hand_pose"][:].astype(np.float32) if "hand_pose" in mano_group else None
            )
            global_orient_data = (
                mano_group["global_orient"][:].astype(np.float32)
                if "global_orient" in mano_group
                else None
            )

            # Get handedness
            n_dets = len(hand_pose_data) if hand_pose_data is not None else num_frames
            right_mask = f["right"][:] if "right" in f else np.ones(n_dets, dtype=np.bool_)

            def _assign_detection(det_idx):
                """Assign a single detection to the appropriate hand array."""
                is_right = right_mask[det_idx] if det_idx < len(right_mask) else True
                if is_right:
                    mano_hand_pose_right[frame] = hand_pose_data[det_idx]
                    if global_orient_data is not None:
                        mano_global_orient_right[frame] = global_orient_data[det_idx]
                else:
                    mano_hand_pose_left[frame] = hand_pose_data[det_idx]
                    if global_orient_data is not None:
                        mano_global_orient_left[frame] = global_orient_data[det_idx]

            if hand_pose_data is not None:
                if frame_idx is not None:
                    # Legacy format: frame_idx[frame] = (start_idx, count)
                    for frame in range(num_frames):
                        if frame < len(frame_idx):
                            start_idx, count = frame_idx[frame]
                            if start_idx >= 0 and count > 0:
                                for i in range(min(count, 2)):
                                    det_idx = int(start_idx + i)
                                    if det_idx < len(hand_pose_data):
                                        _assign_detection(det_idx)
                elif img_idx is not None:
                    # WiLoR format: img_idx[det] = frame_number
                    for det_idx in range(len(img_idx)):
                        frame = int(img_idx[det_idx])
                        if frame < num_frames and det_idx < len(hand_pose_data):
                            _assign_detection(det_idx)

    # Compute angle features
    angle_features = compute_angle_features(
        body_poses=body_poses[:num_frames],
        right_hand_poses=right_hand[:num_frames],
        left_hand_poses=left_hand[:num_frames],
    )

    # Compute HaMeR features
    hamer_features = compute_hamer_features(
        mano_global_orient_right=mano_global_orient_right,
        mano_hand_pose_right=mano_hand_pose_right,
        mano_global_orient_left=mano_global_orient_left,
        mano_hand_pose_left=mano_hand_pose_left,
    )

    return angle_features, hamer_features


def extract_wilor_features_for_segmentor_v2(
    wilor_path: str | Path,
    num_frames: int | None = None,
) -> np.ndarray:
    """
    Extract WiLoR features for SignSegmenterV2 (192 dimensions).

    SignSegmenterV2 expects (T, 192) where:
    - 192 = 2 hands × 96 dims per hand
    - 96 dims = hand_pose (45) + global_orient (3) = 48 dims axis-angle,
      then duplicated/padded to 96

    Args:
        wilor_path: Path to WiLoR hand H5 file (contains mano/ group)
        num_frames: Expected number of frames (auto-detected if None)

    Returns:
        WiLoR features array (T, 192)
    """
    import h5py
    from scipy.spatial.transform import Rotation

    wilor_path = Path(wilor_path)

    with h5py.File(wilor_path, "r") as f:
        w_frame_idx = f["frame_idx"][:] if "frame_idx" in f else None
        w_img_idx = f["img_idx"][:] if "img_idx" in f else None
        right_mask = f["right"][:] if "right" in f else None

        # Get number of frames
        if num_frames is None:
            if w_frame_idx is not None:
                num_frames = len(w_frame_idx)
            elif w_img_idx is not None:
                num_frames = int(w_img_idx.max()) + 1
            else:
                num_frames = 100  # fallback

        # Initialize axis-angle arrays (not rotation matrices)
        hand_pose_right = np.zeros((num_frames, 15, 3), dtype=np.float32)  # 45 dims
        hand_pose_left = np.zeros((num_frames, 15, 3), dtype=np.float32)  # 45 dims
        global_orient_right = np.zeros((num_frames, 3), dtype=np.float32)  # 3 dims
        global_orient_left = np.zeros((num_frames, 3), dtype=np.float32)  # 3 dims

        if "mano" in f:
            hp_data = f["mano/hand_pose"][:].astype(np.float32)  # (N, 15, 3, 3)
            go_data = f["mano/global_orient"][:].astype(np.float32)  # (N, 1, 3, 3)

            # Convert rotation matrices to axis-angle
            def rotmat_to_aa(R):
                """Convert 3x3 rotation matrix to axis-angle (3,)."""
                rot = Rotation.from_matrix(R)
                return rot.as_rotvec()

            def _assign(det_idx: int, frame: int) -> None:
                is_right = right_mask[det_idx] if right_mask is not None else True

                # Convert hand_pose rotation matrices to axis-angle
                for j in range(15):
                    aa = rotmat_to_aa(hp_data[det_idx, j])
                    if is_right:
                        hand_pose_right[frame, j] = aa
                    else:
                        hand_pose_left[frame, j] = aa

                # Convert global_orient rotation matrix to axis-angle
                go_aa = rotmat_to_aa(go_data[det_idx, 0])
                if is_right:
                    global_orient_right[frame] = go_aa
                else:
                    global_orient_left[frame] = go_aa

            if w_frame_idx is not None:
                from sltk.io.h5 import get_frame_detections as _gfd

                for t in range(num_frames):
                    start, end = _gfd(f, t, w_frame_idx)
                    if start >= 0:
                        for i in range(start, min(end, len(hp_data))):
                            _assign(i, t)
            elif w_img_idx is not None:
                for det_idx in range(len(w_img_idx)):
                    t = int(w_img_idx[det_idx])
                    if t < num_frames and det_idx < len(hp_data):
                        _assign(det_idx, t)

    # Flatten to per-frame vectors
    # Each hand: 15 joints × 3 (axis-angle) + 1 global × 3 = 48 dims
    # Need to pad to 96 dims per hand
    hand_right_flat = np.concatenate(
        [hand_pose_right.reshape(num_frames, 45), global_orient_right], axis=-1
    )  # (T, 48)
    hand_left_flat = np.concatenate(
        [hand_pose_left.reshape(num_frames, 45), global_orient_left], axis=-1
    )  # (T, 48)

    # Pad to 96 dims per hand (model expects 96, we have 48)
    # Duplicate the features to fill to 96 dims
    hand_right_96 = np.concatenate([hand_right_flat, hand_right_flat], axis=-1)  # (T, 96)
    hand_left_96 = np.concatenate([hand_left_flat, hand_left_flat], axis=-1)  # (T, 96)

    # Concatenate left + right → 192 total
    # Note: model reshapes as (batch, seq, 2, 96) with [:, :, 0, :] = left, [:, :, 1, :] = right
    features = np.concatenate([hand_left_96, hand_right_96], axis=-1)  # (T, 192)

    return features


def save_features(
    angle_features: np.ndarray,
    hamer_features: np.ndarray,
    angles_path: str | Path,
    hamer_path: str | Path,
) -> tuple[Path, Path]:
    """
    Save computed features to files.

    Args:
        angle_features: Angle features (T, 104)
        hamer_features: HaMeR features (T, 288)
        angles_path: Output path for angle features (.pt or .npy)
        hamer_path: Output path for HaMeR features (.lzma or .npy)

    Returns:
        Tuple of (angles_path, hamer_path)
    """
    import lzma
    import pickle

    try:
        import torch
    except ImportError:
        torch = None

    angles_path = Path(angles_path)
    hamer_path = Path(hamer_path)

    # Ensure directories exist
    angles_path.parent.mkdir(parents=True, exist_ok=True)
    hamer_path.parent.mkdir(parents=True, exist_ok=True)

    # Save angle features
    if angles_path.suffix == ".pt":
        if torch is None:
            raise ImportError("torch is required to save .pt files: pip install torch")
        torch.save(
            {
                "body_angles": angle_features[:, :22],
                "right_angles": angle_features[:, 22:63],
                "left_angles": angle_features[:, 63:104],
            },
            angles_path,
        )
    else:
        np.save(angles_path, angle_features)

    # Save HaMeR features
    if hamer_path.suffix == ".lzma":
        with lzma.open(hamer_path, "wb") as f:
            pickle.dump({"features": hamer_features}, f)
    else:
        np.save(hamer_path, hamer_features)

    return angles_path, hamer_path
