import numpy as np

class SkeletonGenerator:
    def __init__(self, custom_mapping=None):
        # Define default Maya-compatible bone names and their corresponding MediaPipe indices
        self.bone_mapping = {
            # Root & Spine
            'Hips': 23,          # HIPS/Pelvis
            'Spine': 24,         # Lower spine
            'Spine1': 12,        # Mid spine
            'Spine2': 11,        # Upper spine
            
            # Head & Neck
            'Neck': 11,          # BASE_NECK
            'Head': 0,           # HEAD (using nose as reference)
            
            # Left Arm Chain
            'LeftShoulder': 11,  # LEFT_SHOULDER
            'LeftArm': 13,       # LEFT_UPPER_ARM
            'LeftForeArm': 15,   # LEFT_LOWER_ARM
            'LeftHand': 15,      # LEFT_HAND (using wrist as reference)
            
            # Right Arm Chain
            'RightShoulder': 12, # RIGHT_SHOULDER
            'RightArm': 14,      # RIGHT_UPPER_ARM
            'RightForeArm': 16,  # RIGHT_LOWER_ARM
            'RightHand': 16,     # RIGHT_HAND (using wrist as reference)
            
            # Left Leg Chain
            'LeftUpLeg': 23,     # LEFT_UPPER_LEG
            'LeftLeg': 25,       # LEFT_LOWER_LEG
            'LeftFoot': 27,      # LEFT_FOOT
            'LeftToeBase': 31,   # LEFT_TOE
            
            # Right Leg Chain
            'RightUpLeg': 24,    # RIGHT_UPPER_LEG
            'RightLeg': 26,      # RIGHT_LOWER_LEG
            'RightFoot': 28,     # RIGHT_FOOT
            'RightToeBase': 32,  # RIGHT_TOE
        }

    def generate_skeleton(self, landmarks):
        """
        Convert MediaPipe landmarks to Maya-compatible skeleton data
        """
        if landmarks is None:
            return None

        skeleton_data = {}

        # Convert landmarks to skeleton joint positions
        for bone_name, landmark_idx in self.bone_mapping.items():
            position = landmarks[landmark_idx]
            rotation = self._calculate_bone_rotation(landmarks, landmark_idx)

            # Convert NumPy arrays to regular Python lists
            skeleton_data[bone_name] = {
                'position': position.tolist(),
                'rotation': rotation.tolist()
            }

        return skeleton_data

    def _calculate_bone_rotation(self, landmarks, landmark_idx):
        """
        Calculate bone rotation based on connected landmarks
        """
        # Simple rotation calculation - can be enhanced for better accuracy
        rotation = np.zeros(3)

        if landmark_idx > 0:
            # Calculate rotation based on parent-child relationship
            parent_idx = landmark_idx - 1
            direction = landmarks[landmark_idx] - landmarks[parent_idx]

            # Convert direction to euler angles (simplified)
            rotation[0] = np.arctan2(direction[1], direction[2])  # pitch
            rotation[1] = np.arctan2(direction[0], direction[2])  # yaw
            rotation[2] = np.arctan2(direction[0], direction[1])  # roll

        return rotation