"""Tests for pose format loading and conversion."""

import tempfile
from pathlib import Path

import numpy as np
import pytest

from sltk.data.core import PoseSequence
from sltk.data.formats import (
    list_formats,
    load_poses,
    save_poses,
)


class TestGenericFormats:
    """Tests for generic numpy format loading/saving."""

    def test_npy_roundtrip(self):
        """Test saving and loading from .npy format."""
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
            path = Path(f.name)

        try:
            save_poses(poses, path)
            loaded = load_poses(path, format="test", fps=25.0)

            np.testing.assert_array_almost_equal(poses.data, loaded.data)
            assert loaded.fps == 25.0

        finally:
            path.unlink(missing_ok=True)

    def test_npz_roundtrip(self):
        """Test saving and loading from .npz format."""
        data = np.random.randn(50, 55, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=30.0, format="test")

        with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as f:
            path = Path(f.name)

        try:
            save_poses(poses, path)
            loaded = load_poses(path, format="test", fps=30.0)

            np.testing.assert_array_almost_equal(poses.data, loaded.data)

        finally:
            path.unlink(missing_ok=True)

    def test_2d_to_3d_reshape(self):
        """Test automatic reshaping of 2D data to 3D."""
        # Simulated flattened data: (T, N*3)
        data_2d = np.random.randn(100, 63).astype(np.float32)  # 21 keypoints * 3

        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
            path = Path(f.name)
            np.save(path, data_2d)

        try:
            loaded = load_poses(path, format="test", fps=25.0)

            # Should be reshaped to (100, 21, 3)
            assert loaded.data.shape == (100, 21, 3)

        finally:
            Path(path).unlink(missing_ok=True)


class TestListFormats:
    """Tests for format listing."""

    def test_list_formats_returns_dict(self):
        """Test that list_formats returns a dictionary."""
        formats = list_formats()

        assert isinstance(formats, dict)
        assert "wilor" in formats
        assert "smplx" in formats
        assert "mediapipe" in formats

    def test_format_capabilities(self):
        """Test that capabilities are reported."""
        formats = list_formats()

        for name, info in formats.items():
            assert "can_load" in info
            assert "can_save" in info
            assert isinstance(info["can_load"], bool)
            assert isinstance(info["can_save"], bool)


class TestPoseSequenceLoadSave:
    """Tests for PoseSequence load/save methods."""

    def test_load_via_posesequence(self):
        """Test loading via PoseSequence.load class method."""
        data = np.random.randn(50, 21, 3).astype(np.float32)

        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
            path = Path(f.name)
            np.save(path, data)

        try:
            poses = PoseSequence.load(path, format="test", fps=25.0)

            np.testing.assert_array_almost_equal(poses.data, data)
            assert poses.format == "test"

        finally:
            path.unlink(missing_ok=True)

    def test_save_via_posesequence(self):
        """Test saving via PoseSequence.save method."""
        data = np.random.randn(50, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
            path = Path(f.name)

        try:
            poses.save(path)
            loaded = np.load(path)

            np.testing.assert_array_almost_equal(loaded, data)

        finally:
            path.unlink(missing_ok=True)


class TestH5FormatLoading:
    """Tests for H5 format loading (requires h5py)."""

    @pytest.fixture
    def sample_wilor_h5(self, tmp_path):
        """Create a sample WiLoR-format H5 file."""
        try:
            import h5py
        except ImportError:
            pytest.skip("h5py not installed")

        h5_path = tmp_path / "test_wilor.h5"
        num_frames = 100
        num_detections = 150  # More detections than frames (some frames have multiple)

        with h5py.File(h5_path, "w") as f:
            # Create frame_idx (sparse indexing)
            frame_idx = np.zeros((num_frames, 2), dtype=np.int32)
            detection_idx = 0
            for frame in range(num_frames):
                num_dets = 1 if frame % 10 != 0 else 2  # Some frames have 2 hands
                frame_idx[frame] = [detection_idx, num_dets]
                detection_idx += num_dets

            f.create_dataset("frame_idx", data=frame_idx)
            f.create_dataset(
                "kpts_2d", data=np.random.randn(num_detections, 21, 2).astype(np.float16)
            )
            f.create_dataset(
                "kpts_3d", data=np.random.randn(num_detections, 21, 3).astype(np.float16)
            )
            f.create_dataset("right", data=np.random.randint(0, 2, num_detections).astype(bool))

            # MANO group
            mano = f.create_group("mano")
            mano.create_dataset(
                "global_orient", data=np.random.randn(num_detections, 1, 3, 3).astype(np.float16)
            )
            mano.create_dataset(
                "hand_pose", data=np.random.randn(num_detections, 15, 3, 3).astype(np.float16)
            )
            mano.create_dataset(
                "betas", data=np.random.randn(num_detections, 10).astype(np.float16)
            )

            # Metadata
            f.attrs["video_name"] = "test_video"
            f.attrs["fps"] = 25.0

        return h5_path

    def test_wilor_format_load(self, sample_wilor_h5):
        """Test loading WiLoR format H5 file."""
        poses = load_poses(sample_wilor_h5, format="wilor", fps=25.0, keypoints="3d")

        assert poses.num_frames == 100
        assert poses.format == "wilor"
        assert poses.fps == 25.0
        assert "source_path" in poses.metadata
