"""
Pose normalization methods for consistent coordinate systems.

Provides methods to normalize poses across different signers and recording setups,
enabling fair comparisons and better model generalization.
"""

from __future__ import annotations

from typing import Literal

import numpy as np

from sltk.data.core import PoseSequence

# Standard keypoint indices for different pose formats
KEYPOINT_INDICES = {
    "smplx": {
        "left_shoulder": 16,
        "right_shoulder": 17,
        "left_hip": 1,
        "right_hip": 2,
        "pelvis": 0,
        "nose": 24,
        "neck": 12,
        "left_ear": 25,
        "right_ear": 26,
    },
    "mediapipe": {
        "left_shoulder": 11,
        "right_shoulder": 12,
        "left_hip": 23,
        "right_hip": 24,
        "nose": 0,
        "left_ear": 7,
        "right_ear": 8,
    },
    "wilor": {
        # WiLoR is hand-only, wrist is index 0
        "wrist": 0,
        "middle_mcp": 9,  # Middle finger base for palm center
    },
}


CenterMethod = Literal["shoulder_centered", "hip_centered", "face_centered", "wrist_centered"]
ScaleMethod = Literal["unit_torso", "unit_shoulders", "unit_palm", "none"]


def normalize_poses(
    poses: PoseSequence,
    method: CenterMethod = "shoulder_centered",
    scale: ScaleMethod = "unit_torso",
    align_facing: bool = True,
    keypoint_indices: dict | None = None,
) -> PoseSequence:
    """
    Normalize poses to a consistent coordinate system.

    Args:
        poses: Input PoseSequence
        method: Centering method
            - "shoulder_centered": Center between shoulders
            - "hip_centered": Center at pelvis/hip midpoint
            - "face_centered": Center at nose/face
            - "wrist_centered": Center at wrist (for hand-only data)
        scale: Scaling method
            - "unit_torso": Scale so torso height = 1
            - "unit_shoulders": Scale so shoulder width = 1
            - "unit_palm": Scale so palm size = 1 (for hand data)
            - "none": No scaling
        align_facing: Whether to rotate poses to face forward (Z-axis)
        keypoint_indices: Optional custom keypoint index mapping

    Returns:
        Normalized PoseSequence

    Example:
        >>> normalized = normalize_poses(
        ...     poses,
        ...     method="shoulder_centered",
        ...     scale="unit_torso",
        ...     align_facing=True
        ... )
    """
    data = poses.data.copy()

    # Get keypoint indices for this format
    if keypoint_indices is None:
        keypoint_indices = KEYPOINT_INDICES.get(poses.format, {})

    # Apply centering
    if method == "shoulder_centered":
        data = shoulder_center(data, keypoint_indices)
    elif method == "hip_centered":
        data = hip_center(data, keypoint_indices)
    elif method == "face_centered":
        data = face_center(data, keypoint_indices)
    elif method == "wrist_centered":
        data = wrist_center(data, keypoint_indices)

    # Apply scaling
    if scale == "unit_torso":
        data = scale_to_unit_torso(data, keypoint_indices)
    elif scale == "unit_shoulders":
        data = scale_to_unit_shoulders(data, keypoint_indices)
    elif scale == "unit_palm":
        data = scale_to_unit_palm(data, keypoint_indices)

    # Apply rotation alignment
    if align_facing and data.shape[-1] >= 3:
        data = align_facing_direction(data, keypoint_indices)

    return PoseSequence(
        data=data,
        fps=poses.fps,
        format=poses.format,
        confidence=poses.confidence,
        metadata={**poses.metadata, "normalized": True, "norm_method": method, "norm_scale": scale},
    )


def shoulder_center(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """
    Center poses at the midpoint between shoulders.

    Args:
        data: Pose array (T, N, C)
        indices: Keypoint index mapping

    Returns:
        Centered pose array
    """
    left_idx = indices.get("left_shoulder")
    right_idx = indices.get("right_shoulder")

    if left_idx is None or right_idx is None:
        # Fallback: center at mean of all keypoints
        center = np.mean(data, axis=1, keepdims=True)
    else:
        center = (data[:, left_idx : left_idx + 1, :] + data[:, right_idx : right_idx + 1, :]) / 2

    return data - center


def hip_center(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Center poses at the hip/pelvis midpoint."""
    pelvis_idx = indices.get("pelvis")

    if pelvis_idx is not None:
        center = data[:, pelvis_idx : pelvis_idx + 1, :]
    else:
        left_idx = indices.get("left_hip")
        right_idx = indices.get("right_hip")

        if left_idx is not None and right_idx is not None:
            center = (
                data[:, left_idx : left_idx + 1, :] + data[:, right_idx : right_idx + 1, :]
            ) / 2
        else:
            center = np.mean(data, axis=1, keepdims=True)

    return data - center


def face_center(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Center poses at the nose/face position."""
    nose_idx = indices.get("nose")

    if nose_idx is not None:
        center = data[:, nose_idx : nose_idx + 1, :]
    else:
        # Try ear midpoint as fallback
        left_ear = indices.get("left_ear")
        right_ear = indices.get("right_ear")

        if left_ear is not None and right_ear is not None:
            center = (
                data[:, left_ear : left_ear + 1, :] + data[:, right_ear : right_ear + 1, :]
            ) / 2
        else:
            center = np.mean(data, axis=1, keepdims=True)

    return data - center


def wrist_center(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Center poses at the wrist (for hand-only data)."""
    wrist_idx = indices.get("wrist", 0)
    center = data[:, wrist_idx : wrist_idx + 1, :]
    return data - center


def scale_to_unit_torso(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Scale poses so that torso height equals 1."""
    # Estimate torso height from shoulders to hips
    left_shoulder = indices.get("left_shoulder")
    left_hip = indices.get("left_hip")

    if left_shoulder is not None and left_hip is not None:
        shoulder_pos = data[:, left_shoulder, :]
        hip_pos = data[:, left_hip, :]
        torso_length = np.linalg.norm(shoulder_pos - hip_pos, axis=-1, keepdims=True)
        torso_length = np.maximum(torso_length, 1e-6)  # Avoid division by zero

        # Average torso length across frames for stability
        avg_torso = np.mean(torso_length)
        if avg_torso > 1e-6:
            data = data / avg_torso

    return data


def scale_to_unit_shoulders(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Scale poses so that shoulder width equals 1."""
    left_idx = indices.get("left_shoulder")
    right_idx = indices.get("right_shoulder")

    if left_idx is not None and right_idx is not None:
        left_pos = data[:, left_idx, :]
        right_pos = data[:, right_idx, :]
        shoulder_width = np.linalg.norm(left_pos - right_pos, axis=-1, keepdims=True)

        avg_width = np.mean(shoulder_width)
        if avg_width > 1e-6:
            data = data / avg_width

    return data


def scale_to_unit_palm(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """Scale hand poses so that palm size equals 1."""
    wrist_idx = indices.get("wrist", 0)
    middle_mcp = indices.get("middle_mcp", 9)

    wrist_pos = data[:, wrist_idx, :]
    middle_pos = data[:, middle_mcp, :]
    palm_size = np.linalg.norm(middle_pos - wrist_pos, axis=-1, keepdims=True)

    avg_palm = np.mean(palm_size)
    if avg_palm > 1e-6:
        data = data / avg_palm

    return data


def align_facing_direction(
    data: np.ndarray,
    indices: dict,
) -> np.ndarray:
    """
    Rotate poses so the person faces the camera (positive Z direction).

    Uses shoulder orientation to determine facing direction.
    """
    left_idx = indices.get("left_shoulder")
    right_idx = indices.get("right_shoulder")

    if left_idx is None or right_idx is None or data.shape[-1] < 3:
        return data

    # Calculate average shoulder vector
    left_shoulder = np.mean(data[:, left_idx, :], axis=0)
    right_shoulder = np.mean(data[:, right_idx, :], axis=0)

    shoulder_vec = right_shoulder - left_shoulder
    shoulder_vec[2] = 0  # Project to XY plane

    # Calculate rotation angle to align with X-axis
    if np.linalg.norm(shoulder_vec[:2]) > 1e-6:
        angle = np.arctan2(shoulder_vec[1], shoulder_vec[0])

        # Create rotation matrix (around Z-axis)
        cos_a, sin_a = np.cos(-angle), np.sin(-angle)
        rot_matrix = np.array([[cos_a, -sin_a, 0], [sin_a, cos_a, 0], [0, 0, 1]])

        # Apply rotation to all frames and keypoints
        original_shape = data.shape
        data_flat = data.reshape(-1, 3)
        data_rotated = data_flat @ rot_matrix.T
        data = data_rotated.reshape(original_shape)

    return data


def compute_velocity(poses: PoseSequence, smooth_window: int = 5) -> np.ndarray:
    """
    Compute velocity (first derivative) of pose sequence.

    Args:
        poses: Input PoseSequence
        smooth_window: Optional smoothing window size

    Returns:
        Velocity array of shape (T-1, N, C)
    """
    data = poses.data
    velocity = np.diff(data, axis=0) * poses.fps

    if smooth_window > 1:
        # Simple moving average smoothing
        kernel = np.ones(smooth_window) / smooth_window
        smoothed = np.zeros_like(velocity)
        for i in range(velocity.shape[1]):
            for j in range(velocity.shape[2]):
                smoothed[:, i, j] = np.convolve(velocity[:, i, j], kernel, mode="same")
        velocity = smoothed

    return velocity


def compute_acceleration(poses: PoseSequence, smooth_window: int = 5) -> np.ndarray:
    """
    Compute acceleration (second derivative) of pose sequence.

    Args:
        poses: Input PoseSequence
        smooth_window: Optional smoothing window size

    Returns:
        Acceleration array of shape (T-2, N, C)
    """
    data = poses.data
    velocity = np.diff(data, axis=0) * poses.fps
    acceleration = np.diff(velocity, axis=0) * poses.fps

    if smooth_window > 1:
        kernel = np.ones(smooth_window) / smooth_window
        smoothed = np.zeros_like(acceleration)
        for i in range(acceleration.shape[1]):
            for j in range(acceleration.shape[2]):
                smoothed[:, i, j] = np.convolve(acceleration[:, i, j], kernel, mode="same")
        acceleration = smoothed

    return acceleration
