"""H5 feature extraction pipeline for WiLoR hand pose data.

Converts WiLoR HDF5 files to 192-dimensional feature tensors:
  - Left hand: global_orient (6D) + hand_pose (15 joints x 6D) = 96 dims
  - Right hand: same = 96 dims
  - Total: 192 dims

The 6D rotation representation drops the last column of each 3x3 rotation
matrix (Zhou et al. 2019).

Adapted from segmenterv2/segmenter/data/h5_loader.py.
"""

from __future__ import annotations

import logging
from pathlib import Path

import h5py
import numpy as np
import torch

try:
    import hdf5plugin  # noqa: F401 - registers LZ4/other compression filters
except ImportError:
    pass

logger = logging.getLogger(__name__)


def matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray:
    """Convert 3x3 rotation matrices to 6D representation.

    Args:
        matrix: Array of shape (..., 3, 3).

    Returns:
        Array of shape (..., 6).
    """
    batch_dim = matrix.shape[:-2]
    return matrix[..., :2, :].copy().reshape(*batch_dim, 6)


def _read_mano_h5(file_path: Path) -> dict:
    """Read MANO data from a WiLoR H5 file."""
    params: dict = {"data": {}}
    with h5py.File(file_path, "r") as hf:
        params["num_frames"] = hf.attrs["num_frames"]
        params["frame_idx"] = hf["frame_idx"][:]
        params["right"] = hf["right"][:]
        for key in hf["mano"].keys():
            params["data"][key] = hf["mano"][key][:]
    return params


def _convert_to_per_frame(mano: dict) -> np.ndarray:
    """Convert per-detection MANO data to per-frame 2-hand features.

    Returns:
        Array of shape (num_frames, 2, 96) with 6D rotation features.
    """
    num_frames = mano["num_frames"]
    frame_idx = mano["frame_idx"]

    features = np.zeros((num_frames, 2, 96), dtype=np.float32)

    for i, (start, count) in enumerate(frame_idx):
        if count == 0 or start == -1:
            continue

        right_flags = mano["right"][start : start + count]
        unique_sides, first_idx = np.unique(right_flags, return_index=True)
        data_slots = unique_sides.astype(int)

        det_indices = start + first_idx
        go = mano["data"]["global_orient"][det_indices].squeeze(1)
        hp = mano["data"]["hand_pose"][det_indices]

        go_6d = matrix_to_rotation_6d(go)
        n_hands, n_joints = hp.shape[:2]
        hp_flat = hp.reshape(n_hands * n_joints, 3, 3)
        hp_6d = matrix_to_rotation_6d(hp_flat).reshape(n_hands, n_joints, 6)

        hp_flat_96 = hp_6d.reshape(n_hands, -1)
        per_hand = np.concatenate([hp_flat_96, go_6d], axis=-1)

        features[i, data_slots] = per_hand

    return features


def h5_to_features(file_path: Path) -> torch.Tensor:
    """Convert a WiLoR H5 file to a feature tensor.

    Args:
        file_path: Path to the .h5 file.

    Returns:
        Tensor of shape (num_frames, 192).
    """
    file_path = Path(file_path)
    mano = _read_mano_h5(file_path)
    features = _convert_to_per_frame(mano)
    tensor = torch.tensor(features, dtype=torch.float32)
    return tensor.reshape(tensor.shape[0], -1)


def load_h5_directory(dir_path: Path, recursive: bool = True) -> dict[str, torch.Tensor]:
    """Load all H5 files from a directory.

    Args:
        dir_path: Directory containing .h5 files.
        recursive: Whether to search subdirectories.

    Returns:
        Dict mapping video name (stem) to feature tensor (N, 192).
    """
    dir_path = Path(dir_path)
    pattern = "**/*.h5" if recursive else "*.h5"
    h5_files = sorted(dir_path.glob(pattern))

    result = {}
    for h5_path in h5_files:
        try:
            features = h5_to_features(h5_path)
            result[h5_path.stem] = features
        except Exception as e:
            logger.warning("Error processing %s: %s", h5_path.name, e)

    return result
