"""Unified interface for converting motion features to joint positions."""

import numpy as np

from visualization.HumanML3D263.recovery import recover_joint_positions as recover_joint_positions_263
from visualization.HumanML3D263.recovery import StreamJointRecovery as StreamJointRecovery263
from visualization.MotionStreamer272.recovery import recover_joint_positions as recover_joint_positions_272
from visualization.MotionStreamer272.recovery import StreamJointRecovery as StreamJointRecovery272
from visualization.MEI138.recovery import recover_joint_positions as recover_joint_positions_138
from visualization.MEI138.recovery import StreamJointRecovery as StreamJointRecovery138

def get_stream_joint_recovery(dim: int, **kwargs):
    """Create a StreamJointRecovery instance for the given feature dimension.

    Args:
        dim: feature dimension (263 or 138)
        **kwargs: passed to the constructor (joints_num, smoothing_alpha, ...)

    Returns:
        StreamJointRecovery instance
    """
    if dim == 263:
        return StreamJointRecovery263(**kwargs)
    elif dim == 138:
        return StreamJointRecovery138(**kwargs)
    elif dim == 272:
        return StreamJointRecovery272(**kwargs)
    else:
        raise ValueError(f"No StreamJointRecovery for dim={dim}. Supported: [263, 138, 272]")


def convert_motion_to_joints(
    motion_data: np.ndarray,
    dim: int,
    mean: np.ndarray = None,
    std: np.ndarray = None,
):
    """Convert motion features to 22-joint positions.

    Args:
        motion_data: (K, dim) motion features.
        dim: feature dimension (263, 272, or 138).
        mean, std: normalization statistics (optional).

    Returns:
        joints: (K, 22, 3) joint positions.
    """
    if mean is not None and std is not None:
        motion_data = motion_data * std + mean
    if dim == 263:
        return recover_joint_positions_263(motion_data)
    elif dim == 272:
        return recover_joint_positions_272(motion_data)
    elif dim == 138:
        return recover_joint_positions_138(motion_data)
    else:
        raise ValueError(f"Unsupported motion data dimension: {dim}. Supported: [263, 272, 138]")