"""
Tests for feature processing module.

Tests compute_angle_features, compute_hamer_features, load_features_from_h5,
and edge cases with numerical data.
"""

from pathlib import Path
from unittest.mock import MagicMock, patch

import numpy as np
import pytest


class TestComputeBodyAngles:
    """Tests for compute_body_angles function."""

    def test_compute_body_angles_shape(self, sample_body_poses):
        """Test that body angles have correct output shape."""
        from sltk.processing.features import compute_body_angles

        result = compute_body_angles(sample_body_poses)

        assert result.shape == (100, 22)
        assert result.dtype == np.float32

    def test_compute_body_angles_zero_input(self):
        """Test with zero input data."""
        from sltk.processing.features import compute_body_angles

        zeros = np.zeros((50, 33, 3), dtype=np.float32)
        result = compute_body_angles(zeros)

        assert result.shape == (50, 22)
        # Should not crash on zero vectors

    def test_compute_body_angles_single_frame(self):
        """Test with single frame."""
        from sltk.processing.features import compute_body_angles

        single = np.random.randn(1, 33, 3).astype(np.float32)
        result = compute_body_angles(single)

        assert result.shape == (1, 22)

    def test_compute_body_angles_values_range(self, sample_body_poses):
        """Test that angle values are in valid range."""
        from sltk.processing.features import compute_body_angles

        result = compute_body_angles(sample_body_poses)

        # Angles should be between 0 and pi
        angle_values = result[:, :12]  # First 12 are actual angles
        assert np.all(angle_values >= 0) or np.allclose(
            angle_values[angle_values < 0], 0, atol=1e-6
        )
        assert np.all(angle_values <= np.pi + 1e-6)


class TestComputeHandAngles:
    """Tests for compute_hand_angles function."""

    def test_compute_hand_angles_shape(self, sample_hand_poses):
        """Test that hand angles have correct output shape."""
        from sltk.processing.features import compute_hand_angles

        result = compute_hand_angles(sample_hand_poses)

        assert result.shape == (100, 41)
        assert result.dtype == np.float32

    def test_compute_hand_angles_zero_input(self):
        """Test with zero input data."""
        from sltk.processing.features import compute_hand_angles

        zeros = np.zeros((50, 21, 3), dtype=np.float32)
        result = compute_hand_angles(zeros)

        assert result.shape == (50, 41)

    def test_compute_hand_angles_single_frame(self):
        """Test with single frame."""
        from sltk.processing.features import compute_hand_angles

        single = np.random.randn(1, 21, 3).astype(np.float32)
        result = compute_hand_angles(single)

        assert result.shape == (1, 41)


class TestManoToEulerAngles:
    """Tests for mano_to_euler_angles function."""

    def test_mano_to_euler_4d_input(self):
        """Test with 4D MANO input (T, 15, 3, 3)."""
        from sltk.processing.features import mano_to_euler_angles

        mano_pose = np.random.randn(100, 15, 3, 3).astype(np.float32)
        result = mano_to_euler_angles(mano_pose)

        assert result.shape == (100, 41)
        assert result.dtype == np.float32

    def test_mano_to_euler_3d_input(self):
        """Test with 3D MANO input (15, 3, 3) - single frame."""
        from sltk.processing.features import mano_to_euler_angles

        mano_pose = np.random.randn(15, 3, 3).astype(np.float32)
        result = mano_to_euler_angles(mano_pose)

        assert result.shape == (1, 41)

    def test_mano_to_euler_identity_matrices(self):
        """Test with identity rotation matrices."""
        from sltk.processing.features import mano_to_euler_angles

        # Identity matrices should give zero angles
        identity = np.eye(3).reshape(1, 1, 3, 3).astype(np.float32)
        identity = np.tile(identity, (10, 15, 1, 1))

        result = mano_to_euler_angles(identity)

        assert result.shape == (10, 41)
        # Angles from identity should be close to zero
        assert np.allclose(result, 0, atol=1e-5)

    def test_mano_to_euler_values_range(self):
        """Test that Euler angles are in valid range (-pi, pi)."""
        from sltk.processing.features import mano_to_euler_angles

        mano_pose = np.random.randn(50, 15, 3, 3).astype(np.float32)
        result = mano_to_euler_angles(mano_pose)

        assert np.all(result >= -np.pi - 1e-6)
        assert np.all(result <= np.pi + 1e-6)


class TestComputeAngleFeatures:
    """Tests for compute_angle_features function."""

    def test_compute_angle_features_shape(self, sample_body_poses, sample_hand_poses):
        """Test that combined angle features have correct shape."""
        from sltk.processing.features import compute_angle_features

        result = compute_angle_features(
            body_poses=sample_body_poses[:100],
            right_hand_poses=sample_hand_poses,
            left_hand_poses=sample_hand_poses,
        )

        assert result.shape == (100, 104)  # 22 + 41 + 41
        assert result.dtype == np.float32

    def test_compute_angle_features_with_mano(
        self, sample_body_poses, sample_hand_poses, sample_mano_params
    ):
        """Test with MANO parameters for hand angles."""
        from sltk.processing.features import compute_angle_features

        result = compute_angle_features(
            body_poses=sample_body_poses[:100],
            right_hand_poses=sample_hand_poses,
            left_hand_poses=sample_hand_poses,
            mano_right=sample_mano_params,
            mano_left=sample_mano_params,
        )

        assert result.shape == (100, 104)

    def test_compute_angle_features_length_mismatch(self):
        """Test with mismatched input lengths."""
        from sltk.processing.features import compute_angle_features

        body = np.random.randn(100, 33, 3).astype(np.float32)
        right = np.random.randn(80, 21, 3).astype(np.float32)
        left = np.random.randn(90, 21, 3).astype(np.float32)

        result = compute_angle_features(body, right, left)

        # Should use minimum length
        assert result.shape[0] == 80

    def test_compute_angle_features_empty_input(self):
        """Test with empty input arrays."""
        from sltk.processing.features import compute_angle_features

        body = np.random.randn(0, 33, 3).astype(np.float32)
        right = np.random.randn(0, 21, 3).astype(np.float32)
        left = np.random.randn(0, 21, 3).astype(np.float32)

        result = compute_angle_features(body, right, left)

        assert result.shape[0] == 0


class TestComputeHamerFeatures:
    """Tests for compute_hamer_features function."""

    def test_compute_hamer_features_shape(self):
        """Test that HaMeR features have correct shape."""
        from sltk.processing.features import compute_hamer_features

        T = 100
        result = compute_hamer_features(
            mano_global_orient_right=np.random.randn(T, 1, 3, 3).astype(np.float32),
            mano_hand_pose_right=np.random.randn(T, 15, 3, 3).astype(np.float32),
            mano_global_orient_left=np.random.randn(T, 1, 3, 3).astype(np.float32),
            mano_hand_pose_left=np.random.randn(T, 15, 3, 3).astype(np.float32),
        )

        assert result.shape == (T, 288)  # 2*(15*9 + 9) = 2*(135+9) = 288
        assert result.dtype == np.float32

    def test_compute_hamer_features_identity(self):
        """Test with identity rotation matrices."""
        from sltk.processing.features import compute_hamer_features

        T = 10
        eye = np.eye(3).reshape(1, 1, 3, 3).astype(np.float32)

        result = compute_hamer_features(
            mano_global_orient_right=np.tile(eye, (T, 1, 1, 1)),
            mano_hand_pose_right=np.tile(eye, (T, 15, 1, 1)),
            mano_global_orient_left=np.tile(eye, (T, 1, 1, 1)),
            mano_hand_pose_left=np.tile(eye, (T, 15, 1, 1)),
        )

        assert result.shape == (T, 288)

    def test_compute_hamer_features_single_frame(self):
        """Test with single frame."""
        from sltk.processing.features import compute_hamer_features

        result = compute_hamer_features(
            mano_global_orient_right=np.random.randn(1, 1, 3, 3).astype(np.float32),
            mano_hand_pose_right=np.random.randn(1, 15, 3, 3).astype(np.float32),
            mano_global_orient_left=np.random.randn(1, 1, 3, 3).astype(np.float32),
            mano_hand_pose_left=np.random.randn(1, 15, 3, 3).astype(np.float32),
        )

        assert result.shape == (1, 288)


class TestLoadFeaturesFromH5:
    """Tests for load_features_from_h5 function."""

    def test_load_features_from_h5(self, tmp_path):
        """Test loading features from H5 files."""
        import h5py

        from sltk.processing.features import load_features_from_h5

        # Create mock MediaPipe H5
        mp_path = tmp_path / "mediapipe.h5"
        with h5py.File(mp_path, "w") as f:
            f.create_dataset("poses", data=np.random.randn(100, 33, 3))
            f.create_dataset("left_hand", data=np.random.randn(100, 21, 3))
            f.create_dataset("right_hand", data=np.random.randn(100, 21, 3))

        # Create mock WiLoR H5
        wilor_path = tmp_path / "wilor.h5"
        with h5py.File(wilor_path, "w") as f:
            f.create_dataset(
                "frame_idx",
                data=np.column_stack(
                    [
                        np.arange(100),
                        np.ones(100, dtype=np.int32),
                    ]
                ),
            )
            mano = f.create_group("mano")
            mano.create_dataset("hand_pose", data=np.random.randn(100, 15, 3, 3))
            mano.create_dataset("global_orient", data=np.random.randn(100, 1, 3, 3))
            f.create_dataset("right", data=np.ones(100, dtype=bool))

        angle_features, hamer_features = load_features_from_h5(mp_path, wilor_path)

        assert angle_features.shape == (100, 104)
        assert hamer_features.shape == (100, 288)

    def test_load_features_body_key_fallback(self, tmp_path):
        """Test fallback to 'body' key in H5 file."""
        import h5py

        from sltk.processing.features import load_features_from_h5

        # Create H5 with 'body' key instead of 'poses'
        mp_path = tmp_path / "mediapipe.h5"
        with h5py.File(mp_path, "w") as f:
            f.create_dataset("body", data=np.random.randn(50, 33, 3))

        # Create mock WiLoR H5
        wilor_path = tmp_path / "wilor.h5"
        with h5py.File(wilor_path, "w") as f:
            f.create_dataset(
                "frame_idx",
                data=np.column_stack(
                    [
                        np.arange(50),
                        np.ones(50, dtype=np.int32),
                    ]
                ),
            )

        angle_features, hamer_features = load_features_from_h5(mp_path, wilor_path, num_frames=50)

        assert angle_features.shape[0] == 50

    def test_load_features_missing_body_raises(self, tmp_path):
        """Test that missing body poses raises error."""
        import h5py

        from sltk.processing.features import load_features_from_h5

        # Create H5 without body poses
        mp_path = tmp_path / "mediapipe.h5"
        with h5py.File(mp_path, "w") as f:
            f.create_dataset("other_data", data=np.random.randn(50, 10))

        wilor_path = tmp_path / "wilor.h5"
        with h5py.File(wilor_path, "w") as f:
            f.create_dataset(
                "frame_idx",
                data=np.column_stack(
                    [
                        np.arange(50),
                        np.ones(50, dtype=np.int32),
                    ]
                ),
            )

        with pytest.raises(ValueError, match="Could not find body poses"):
            load_features_from_h5(mp_path, wilor_path)


class TestSaveFeatures:
    """Tests for save_features function."""

    def test_save_features_pt_format(self, tmp_path):
        """Test saving angle features as .pt file."""
        pytest.importorskip("torch")
        from sltk.processing.features import save_features

        angle_features = np.random.randn(100, 104).astype(np.float32)
        hamer_features = np.random.randn(100, 288).astype(np.float32)

        angles_path = tmp_path / "angles.pt"
        hamer_path = tmp_path / "hamer.lzma"

        result = save_features(angle_features, hamer_features, angles_path, hamer_path)

        assert result[0] == angles_path
        assert result[1] == hamer_path
        assert angles_path.exists()
        assert hamer_path.exists()

    def test_save_features_npy_format(self, tmp_path):
        """Test saving features as .npy files."""
        pytest.importorskip("torch")
        from sltk.processing.features import save_features

        angle_features = np.random.randn(100, 104).astype(np.float32)
        hamer_features = np.random.randn(100, 288).astype(np.float32)

        angles_path = tmp_path / "angles.npy"
        hamer_path = tmp_path / "hamer.npy"

        save_features(angle_features, hamer_features, angles_path, hamer_path)

        assert angles_path.exists()
        assert hamer_path.exists()

        # Verify contents
        loaded = np.load(angles_path)
        assert loaded.shape == (100, 104)

    def test_save_features_creates_directories(self, tmp_path):
        """Test that save_features creates parent directories."""
        pytest.importorskip("torch")
        from sltk.processing.features import save_features

        angle_features = np.random.randn(50, 104).astype(np.float32)
        hamer_features = np.random.randn(50, 288).astype(np.float32)

        angles_path = tmp_path / "subdir1" / "angles.npy"
        hamer_path = tmp_path / "subdir2" / "hamer.npy"

        save_features(angle_features, hamer_features, angles_path, hamer_path)

        assert angles_path.parent.exists()
        assert hamer_path.parent.exists()


class TestNumericalEdgeCases:
    """Tests for numerical edge cases."""

    def test_nan_in_poses(self):
        """Test handling of NaN values in pose data."""
        from sltk.processing.features import compute_body_angles

        data = np.random.randn(50, 33, 3).astype(np.float32)
        data[10:15] = np.nan

        # Should not raise, but may produce NaN output
        result = compute_body_angles(data)

        assert result.shape == (50, 22)

    def test_inf_in_poses(self):
        """Test handling of Inf values in pose data."""
        from sltk.processing.features import compute_body_angles

        data = np.random.randn(50, 33, 3).astype(np.float32)
        data[10:15] = np.inf

        # Should not crash
        result = compute_body_angles(data)

        assert result.shape == (50, 22)

    def test_very_small_vectors(self):
        """Test handling of very small (near-zero) vectors."""
        from sltk.processing.features import compute_body_angles

        data = np.random.randn(50, 33, 3).astype(np.float32) * 1e-10

        result = compute_body_angles(data)

        assert result.shape == (50, 22)
        # Should handle division by near-zero norm

    def test_very_large_values(self):
        """Test handling of very large values."""
        from sltk.processing.features import compute_body_angles

        data = np.random.randn(50, 33, 3).astype(np.float32) * 1e10

        result = compute_body_angles(data)

        assert result.shape == (50, 22)

    def test_mixed_nan_inf(self):
        """Test handling of mixed NaN and Inf."""
        from sltk.processing.features import compute_hand_angles

        data = np.random.randn(50, 21, 3).astype(np.float32)
        data[5:10] = np.nan
        data[15:20] = np.inf

        result = compute_hand_angles(data)

        assert result.shape == (50, 41)

    def test_negative_zero_vectors(self):
        """Test handling of negative zero (-0.0)."""
        from sltk.processing.features import compute_body_angles

        data = np.zeros((50, 33, 3), dtype=np.float32)
        data = -data  # Create negative zeros

        result = compute_body_angles(data)

        assert result.shape == (50, 22)

    def test_orthogonal_vectors(self):
        """Test angle computation for orthogonal vectors (90 degrees)."""
        from sltk.processing.features import compute_body_angles

        # Create poses where connected joints form 90-degree angles
        data = np.zeros((10, 33, 3), dtype=np.float32)

        # Set up specific keypoints for first angle (hip)
        # Joints 11 (left shoulder), 23 (left hip), 25 (left knee)
        for t in range(10):
            data[t, 11] = [0, 1, 0]  # Above hip
            data[t, 23] = [0, 0, 0]  # Hip at origin
            data[t, 25] = [1, 0, 0]  # To the right of hip

        result = compute_body_angles(data)

        # First angle should be close to pi/2 (90 degrees)
        # Note: depending on implementation, might be 0 or pi/2
        assert result.shape == (10, 22)

    def test_collinear_vectors(self):
        """Test angle computation for collinear vectors (0 or 180 degrees)."""
        from sltk.processing.features import compute_body_angles

        data = np.zeros((10, 33, 3), dtype=np.float32)

        # Set up collinear keypoints
        for t in range(10):
            data[t, 11] = [0, 2, 0]
            data[t, 23] = [0, 1, 0]
            data[t, 25] = [0, 0, 0]

        result = compute_body_angles(data)

        # Angle should be 0 (vectors point in same direction) or pi (opposite)
        assert result.shape == (10, 22)


class TestFeatureConsistency:
    """Tests for feature computation consistency."""

    def test_deterministic_output(self, sample_body_poses, sample_hand_poses):
        """Test that output is deterministic for same input."""
        from sltk.processing.features import compute_angle_features

        result1 = compute_angle_features(
            sample_body_poses,
            sample_hand_poses,
            sample_hand_poses,
        )

        result2 = compute_angle_features(
            sample_body_poses,
            sample_hand_poses,
            sample_hand_poses,
        )

        np.testing.assert_array_equal(result1, result2)

    def test_feature_dimension_breakdown(self, sample_body_poses, sample_hand_poses):
        """Test that feature dimensions add up correctly."""
        from sltk.processing.features import (
            compute_angle_features,
            compute_body_angles,
            compute_hand_angles,
        )

        body_angles = compute_body_angles(sample_body_poses)
        right_angles = compute_hand_angles(sample_hand_poses)
        left_angles = compute_hand_angles(sample_hand_poses)

        assert body_angles.shape[1] == 22
        assert right_angles.shape[1] == 41
        assert left_angles.shape[1] == 41

        combined = compute_angle_features(
            sample_body_poses,
            sample_hand_poses,
            sample_hand_poses,
        )

        assert combined.shape[1] == 22 + 41 + 41

    def test_different_dtype_handling(self):
        """Test handling of different input dtypes."""
        from sltk.processing.features import compute_body_angles

        # Test float64 input
        data_f64 = np.random.randn(50, 33, 3).astype(np.float64)
        result = compute_body_angles(data_f64)

        assert result.dtype == np.float32  # Should always output float32
