"""Tests for pose normalization module."""

import numpy as np

from sltk.data.core import PoseSequence
from sltk.processing.normalization import (
    KEYPOINT_INDICES,
    compute_acceleration,
    compute_velocity,
    normalize_poses,
    scale_to_unit_shoulders,
    scale_to_unit_torso,
    shoulder_center,
)


class TestShoulderCenter:
    """Tests for shoulder centering."""

    def test_centers_at_shoulder_midpoint(self):
        """Test that poses are centered at shoulder midpoint."""
        # Create pose with known shoulder positions
        data = np.zeros((10, 30, 3), dtype=np.float32)

        # Set shoulders (using SMPLX indices)
        left_shoulder_idx = KEYPOINT_INDICES["smplx"]["left_shoulder"]
        right_shoulder_idx = KEYPOINT_INDICES["smplx"]["right_shoulder"]

        data[:, left_shoulder_idx, :] = [1.0, 2.0, 0.0]
        data[:, right_shoulder_idx, :] = [3.0, 2.0, 0.0]

        centered = shoulder_center(data, KEYPOINT_INDICES["smplx"])

        # Midpoint should be at [2.0, 2.0, 0.0], so centered values should be shifted
        expected_left = [-1.0, 0.0, 0.0]  # [1, 2, 0] - [2, 2, 0]
        expected_right = [1.0, 0.0, 0.0]  # [3, 2, 0] - [2, 2, 0]

        np.testing.assert_array_almost_equal(centered[0, left_shoulder_idx, :], expected_left)
        np.testing.assert_array_almost_equal(centered[0, right_shoulder_idx, :], expected_right)


class TestScaling:
    """Tests for scaling normalization."""

    def test_scale_to_unit_shoulders(self):
        """Test scaling to unit shoulder width."""
        data = np.zeros((10, 30, 3), dtype=np.float32)

        left_idx = KEYPOINT_INDICES["smplx"]["left_shoulder"]
        right_idx = KEYPOINT_INDICES["smplx"]["right_shoulder"]

        # Shoulder width of 2.0
        data[:, left_idx, :] = [0.0, 0.0, 0.0]
        data[:, right_idx, :] = [2.0, 0.0, 0.0]

        scaled = scale_to_unit_shoulders(data, KEYPOINT_INDICES["smplx"])

        # After scaling, shoulder width should be 1.0
        scaled_width = np.linalg.norm(scaled[0, right_idx, :] - scaled[0, left_idx, :])
        np.testing.assert_almost_equal(scaled_width, 1.0)

    def test_scale_to_unit_torso(self):
        """Test scaling to unit torso height."""
        data = np.zeros((10, 30, 3), dtype=np.float32)

        left_shoulder = KEYPOINT_INDICES["smplx"]["left_shoulder"]
        left_hip = KEYPOINT_INDICES["smplx"]["left_hip"]

        # Torso height of 2.0
        data[:, left_shoulder, :] = [0.0, 0.0, 0.0]
        data[:, left_hip, :] = [0.0, 2.0, 0.0]

        scaled = scale_to_unit_torso(data, KEYPOINT_INDICES["smplx"])

        # After scaling, torso height should be 1.0
        scaled_height = np.linalg.norm(scaled[0, left_shoulder, :] - scaled[0, left_hip, :])
        np.testing.assert_almost_equal(scaled_height, 1.0)


class TestNormalizePoses:
    """Tests for the main normalize_poses function."""

    def test_normalize_with_posesequence(self):
        """Test normalization with PoseSequence input."""
        data = np.random.randn(100, 55, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="smplx")

        normalized = normalize_poses(
            poses,
            method="shoulder_centered",
            scale="unit_torso",
            align_facing=False,
        )

        assert normalized.format == "smplx"
        assert normalized.fps == 25.0
        assert normalized.metadata.get("normalized") is True

    def test_normalize_updates_metadata(self):
        """Test that normalization updates metadata."""
        data = np.random.randn(50, 55, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="smplx")

        normalized = normalize_poses(
            poses,
            method="hip_centered",
            scale="unit_shoulders",
        )

        assert normalized.metadata["norm_method"] == "hip_centered"
        assert normalized.metadata["norm_scale"] == "unit_shoulders"


class TestVelocityAcceleration:
    """Tests for velocity and acceleration computation."""

    def test_compute_velocity(self):
        """Test velocity computation."""
        # Linear motion: x increases by 1 per frame
        data = np.zeros((100, 10, 3), dtype=np.float32)
        data[:, :, 0] = np.arange(100)[:, None]

        poses = PoseSequence(data=data, fps=25.0, format="test")
        velocity = compute_velocity(poses, smooth_window=1)

        # Velocity should be 25 (1 unit per frame * 25 fps)
        expected_vel = 25.0
        np.testing.assert_array_almost_equal(
            velocity[:, 0, 0], np.full(99, expected_vel), decimal=3
        )

    def test_compute_acceleration(self):
        """Test acceleration computation."""
        # Quadratic motion: position = t^2
        t = np.arange(100).astype(np.float32)
        data = np.zeros((100, 10, 3), dtype=np.float32)
        data[:, :, 0] = (t**2)[:, None]

        poses = PoseSequence(data=data, fps=25.0, format="test")
        acceleration = compute_acceleration(poses, smooth_window=1)

        # For x = t^2, acceleration = 2 * fps^2 = 2 * 625 = 1250
        # (discrete approximation may vary slightly)
        assert acceleration.shape[0] == 98  # T - 2

    def test_velocity_shape(self):
        """Test that velocity has correct shape."""
        data = np.random.randn(50, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="wilor")

        velocity = compute_velocity(poses)

        assert velocity.shape == (49, 21, 3)  # T-1 frames
