"""
Motion data handler for loading and managing SMPL motion pickle files.
Handles .pkl files containing dance motion data with SMPL format.
"""

import pickle
import numpy as np
from typing import Dict, Optional, Tuple


class MotionHandler:
    """Handles loading and accessing motion data from pickle files."""

    def __init__(self, fps: int = 30):
        """
        Initialize motion handler.

        Args:
            fps: Frames per second for the motion data (default: 30)
        """
        self.fps = fps
        self.motion_data: Optional[Dict] = None
        self.full_pose: Optional[np.ndarray] = None
        self.num_frames: int = 0
        self.duration: float = 0.0

    def load_pickle(self, file_bytes: bytes) -> Dict:
        """
        Load motion data from pickle file bytes.

        Args:
            file_bytes: Raw bytes from uploaded pickle file

        Returns:
            Dictionary containing motion metadata

        Raises:
            ValueError: If pickle file format is invalid
        """
        try:
            # Load pickle data
            self.motion_data = pickle.loads(file_bytes)

            # Validate required keys
            if "full_pose" not in self.motion_data:
                raise ValueError("Pickle file must contain 'full_pose' field")

            # Extract full_pose: (num_frames, 24, 3)
            self.full_pose = self.motion_data["full_pose"]

            # Validate shape
            if (
                len(self.full_pose.shape) != 3
                or self.full_pose.shape[1] != 24
                or self.full_pose.shape[2] != 3
            ):
                raise ValueError(
                    f"Invalid full_pose shape: {self.full_pose.shape}. "
                    f"Expected (num_frames, 24, 3)"
                )

            # Calculate metadata
            self.num_frames = self.full_pose.shape[0]
            self.duration = self.num_frames / self.fps

            return {
                "num_frames": int(self.num_frames),
                "duration": float(self.duration),
                "fps": self.fps,
                "num_joints": 24,
                "shape": list(self.full_pose.shape),
            }

        except pickle.UnpicklingError as e:
            raise ValueError(f"Failed to unpickle file: {str(e)}")
        except Exception as e:
            raise ValueError(f"Error loading motion data: {str(e)}")

    def get_frame_at_timestamp(self, timestamp: float) -> Optional[np.ndarray]:
        """
        Get joint positions for a specific timestamp.

        Args:
            timestamp: Time in seconds

        Returns:
            Joint positions array of shape (24, 3) or None if timestamp is out of bounds
        """
        if self.full_pose is None:
            return None

        # Calculate frame index
        frame_idx = int(timestamp * self.fps)

        # Clamp to valid range
        if frame_idx < 0:
            frame_idx = 0
        elif frame_idx >= self.num_frames:
            frame_idx = self.num_frames - 1

        # Return frame data: (24, 3)
        return self.full_pose[frame_idx]

    def get_frame_by_index(self, frame_idx: int) -> Optional[np.ndarray]:
        """
        Get joint positions for a specific frame index.

        Args:
            frame_idx: Frame index (0-based)

        Returns:
            Joint positions array of shape (24, 3) or None if invalid index
        """
        if self.full_pose is None or frame_idx < 0 or frame_idx >= self.num_frames:
            return None

        return self.full_pose[frame_idx]

    def is_loaded(self) -> bool:
        """Check if motion data is loaded."""
        return self.full_pose is not None


# SMPL Joint Names (24 joints)
SMPL_JOINT_NAMES = [
    "root",  # 0
    "lhip",  # 1
    "rhip",  # 2
    "belly",  # 3
    "lknee",  # 4
    "rknee",  # 5
    "spine",  # 6
    "lankle",  # 7
    "rankle",  # 8
    "chest",  # 9
    "ltoes",  # 10
    "rtoes",  # 11
    "neck",  # 12
    "linshoulder",  # 13
    "rinshoulder",  # 14
    "head",  # 15
    "lshoulder",  # 16
    "rshoulder",  # 17
    "lelbow",  # 18
    "relbow",  # 19
    "lwrist",  # 20
    "rwrist",  # 21
    "lhand",  # 22
    "rhand",  # 23
]