"""
Joint mapping between SMPL (24 joints) and MediaPipe Pose (33 landmarks).
Handles coordinate system conversion and normalization.
"""

import numpy as np
from typing import List, Dict


# SMPL to MediaPipe joint mapping
# Maps SMPL joint indices to MediaPipe landmark indices
# Format: {mediapipe_index: smpl_index}
SMPL_TO_MEDIAPIPE_MAP = {
    # Torso and core
    11: 16,  # Left shoulder (MP) -> lshoulder (SMPL)
    12: 17,  # Right shoulder (MP) -> rshoulder (SMPL)
    23: 1,  # Left hip (MP) -> lhip (SMPL)
    24: 2,  # Right hip (MP) -> rhip (SMPL)
    # Arms
    13: 18,  # Left elbow (MP) -> lelbow (SMPL)
    14: 19,  # Right elbow (MP) -> relbow (SMPL)
    15: 20,  # Left wrist (MP) -> lwrist (SMPL)
    16: 21,  # Right wrist (MP) -> rwrist (SMPL)
    # Legs
    25: 4,  # Left knee (MP) -> lknee (SMPL)
    26: 5,  # Right knee (MP) -> rknee (SMPL)
    27: 7,  # Left ankle (MP) -> lankle (SMPL)
    28: 8,  # Right ankle (MP) -> rankle (SMPL)
    # Head/Face
    0: 15,  # Nose (MP) -> head (SMPL) - approximate
    # Feet
    31: 10,  # Left foot index (MP) -> ltoes (SMPL)
    32: 11,  # Right foot index (MP) -> rtoes (SMPL)
}


class JointMapper:
    """Maps SMPL joints to MediaPipe landmarks with coordinate conversion."""

    def __init__(self):
        """Initialize joint mapper with coordinate conversion parameters."""
        # These will be set based on the SMPL data range
        self.smpl_bounds = None
        self.center = None
        self.scale = None

    def calibrate_from_motion(self, full_pose: np.ndarray):
        """
        Calibrate normalization parameters from the full motion sequence.

        Args:
            full_pose: Motion data of shape (num_frames, 24, 3)
        """
        # Find min/max across all frames and joints
        all_joints = full_pose.reshape(-1, 3)  # (num_frames * 24, 3)

        min_coords = np.min(all_joints, axis=0)
        max_coords = np.max(all_joints, axis=0)

        # Calculate center and scale
        self.center = (min_coords + max_coords) / 2
        self.scale = np.max(max_coords - min_coords)

        # Store bounds
        self.smpl_bounds = {
            "min": min_coords,
            "max": max_coords,
            "center": self.center,
            "scale": self.scale,
        }

    def smpl_to_mediapipe_landmarks(
        self, smpl_joints: np.ndarray
    ) -> List[Dict[str, float]]:
        """
        Convert SMPL joint positions to MediaPipe landmark format.

        Args:
            smpl_joints: SMPL joint positions of shape (24, 3) in meters

        Returns:
            List of 33 MediaPipe landmarks with x, y, z, visibility
        """
        # Initialize all 33 MediaPipe landmarks with low visibility (missing)
        landmarks = []
        for _ in range(33):
            landmarks.append(
                {
                    "x": 0.5,
                    "y": 0.5,
                    "z": 0.0,
                    "visibility": 0.1,  # Low visibility for unmapped landmarks
                }
            )

        # Normalize SMPL coordinates if calibrated
        if self.center is not None and self.scale is not None:
            normalized_joints = (smpl_joints - self.center) / self.scale

            # SMPL is Y-up (vertical), Z-forward, X-right
            # Currently seeing top view, so rotate 90 degrees:
            # Swap Y and Z to get front view
            # - X (left-right) -> screen X
            # - Z (depth) -> screen Y (vertical, flipped)
            # - Y (up) -> ignore (becomes depth)

            # Create output with 90-degree rotation (front view)
            output_joints = np.zeros_like(normalized_joints)
            output_joints[:, 0] = normalized_joints[:, 0]  # X -> screen X (horizontal)
            output_joints[:, 1] = normalized_joints[:, 2]  # Z -> screen Y (vertical)
            output_joints[:, 2] = normalized_joints[:, 1]  # Y -> depth (ignored)

            # Shift to 0-1 range and flip Y (MediaPipe Y is top-to-bottom)
            output_joints = output_joints * 0.6 + 0.5  # Scale to better fill canvas
            output_joints[:, 1] = (
                1.0 - output_joints[:, 1]
            )  # Flip Y axis (top is 0, bottom is 1)

            normalized_joints = output_joints
        else:
            # Fallback: simple normalization
            normalized_joints = smpl_joints.copy()

        # Map SMPL joints to MediaPipe landmarks
        for mp_idx, smpl_idx in SMPL_TO_MEDIAPIPE_MAP.items():
            joint_pos = normalized_joints[smpl_idx]
            landmarks[mp_idx] = {
                "x": float(joint_pos[0]),
                "y": float(joint_pos[1]),
                "z": float(joint_pos[2]),
                "visibility": 0.9,  # High visibility for mapped joints
            }

        return landmarks

    def get_mapped_joints_for_comparison(self) -> List[int]:
        """
        Get list of MediaPipe landmark indices that are mapped from SMPL.
        These should be used for pose comparison.

        Returns:
            List of MediaPipe landmark indices
        """
        return sorted(list(SMPL_TO_MEDIAPIPE_MAP.keys()))


def create_mapper_from_motion(full_pose: np.ndarray) -> JointMapper:
    """
    Create and calibrate a JointMapper from motion data.

    Args:
        full_pose: Motion data of shape (num_frames, 24, 3)

    Returns:
        Calibrated JointMapper instance
    """
    mapper = JointMapper()
    mapper.calibrate_from_motion(full_pose)
    return mapper