"""Tests for the data soup module."""

import numpy as np
import pytest

from neon.data.action_space import ControlMode, G1ActionSpace
from neon.data.data_soup import ActionMapper, DataSourceConfig, DataSoupConfig, NeonEpisode


class TestActionMapper:
    """Test the cross-embodiment action mapper."""

    def test_g1_identity_mapping(self):
        """G1 → G1 should be identity."""
        space = G1ActionSpace(mode=ControlMode.WHOLE_BODY, include_locomotion=False)
        mapper = ActionMapper(space)

        actions = np.random.randn(10, 29).astype(np.float32)
        mapped = mapper.map_actions(actions, "g1", normalize=False)
        # Should have correct shape
        assert mapped.shape == (10, space.action_dim)

    def test_franka_to_g1(self):
        """Franka (8 DoF) should map to G1 right arm."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=False)
        mapper = ActionMapper(space)

        actions = np.ones((5, 8), dtype=np.float32)  # 7 joints + gripper
        mapped = mapper.map_actions(actions, "franka", normalize=False)
        assert mapped.shape == (5, space.action_dim)

    def test_unknown_embodiment(self):
        """Unknown embodiment should zero-pad."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=False)
        mapper = ActionMapper(space)

        actions = np.ones((3, 6), dtype=np.float32)
        mapped = mapper.map_actions(actions, "unknown_robot", normalize=False)
        assert mapped.shape == (3, space.action_dim)

    def test_agibot_mapping(self):
        """Agibot (14 DoF bimanual) should map to G1 arms."""
        space = G1ActionSpace(mode=ControlMode.ARMS_ONLY, include_locomotion=False)
        mapper = ActionMapper(space)

        actions = np.random.randn(8, 14).astype(np.float32)
        mapped = mapper.map_actions(actions, "agibot", normalize=False)
        assert mapped.shape == (8, space.action_dim)


class TestNeonEpisode:
    """Test the episode container."""

    def test_episode_creation(self):
        """Episode should store data correctly."""
        ep = NeonEpisode(
            images=[None] * 10,
            actions=np.zeros((10, 14), dtype=np.float32),
            language="pick up the cup",
        )
        assert ep.length == 10
        assert ep.language == "pick up the cup"

    def test_empty_episode(self):
        """Empty episode should have length 0."""
        ep = NeonEpisode(images=[], actions=np.zeros((0, 14), dtype=np.float32))
        assert ep.length == 0


class TestDataSourceConfig:
    """Test data source configuration."""

    def test_default_config(self):
        """Default config should have reasonable values."""
        config = DataSourceConfig(name="test", type="lerobot", path="test/dataset")
        assert config.weight == 1.0
        assert config.image_size == 224
        assert config.fps == 4

    def test_voice_commands_config(self):
        """Voice commands config should support category filtering."""
        config = DataSourceConfig(
            name="voice",
            type="voice_commands",
            path="cagataydev/vlm-voice-commands",
            weight=0.5,
            category_filter=["pick_place", "manipulation"],
            max_samples=1000,
        )
        assert config.type == "voice_commands"
        assert config.category_filter == ["pick_place", "manipulation"]
        assert config.max_samples == 1000

    def test_stereo4d_config(self):
        """Stereo4D config should support stereo pair options."""
        config = DataSourceConfig(
            name="stereo",
            type="stereo4d",
            path="cagataydev/strands-kitchen-stereo4d",
            weight=1.0,
            use_stereo_pair=True,
            compute_depth=False,
        )
        assert config.type == "stereo4d"
        assert config.use_stereo_pair is True
        assert config.compute_depth is False

    def test_cosmos_dreamgen_config(self):
        """Cosmos DreamGen config should support relative actions."""
        config = DataSourceConfig(
            name="cosmos-bridge",
            type="cosmos_dreamgen",
            path="/data/bridge",
            weight=2.0,
            use_relative_actions=True,
            action_scaler=20.0,
            camera_ids=["head", "hand_0", "hand_1"],
        )
        assert config.type == "cosmos_dreamgen"
        assert config.use_relative_actions is True
        assert config.action_scaler == 20.0
        assert config.camera_ids == ["head", "hand_0", "hand_1"]


class TestNeonEpisodeFlags:
    """Test episode metadata flags for training loop."""

    def test_language_only_flag(self):
        """Language-only episodes should flag correctly."""
        ep = NeonEpisode(
            images=[],
            actions=np.zeros((1, 14), dtype=np.float32),
            language="pick up the red cup",
            metadata={"language_only": True, "type": "voice_commands"},
        )
        assert ep.metadata["language_only"] is True
        assert ep.length == 0

    def test_visual_only_flag(self):
        """Visual-only episodes should flag correctly."""
        ep = NeonEpisode(
            images=[None] * 16,
            actions=np.zeros((16, 14), dtype=np.float32),
            language="observe kitchen scene",
            metadata={"visual_only": True, "type": "stereo4d"},
        )
        assert ep.metadata["visual_only"] is True
        assert ep.length == 16

    def test_standard_episode_no_flags(self):
        """Standard episodes should not have special flags."""
        ep = NeonEpisode(
            images=[None] * 10,
            actions=np.random.randn(10, 14).astype(np.float32),
            language="move arm forward",
            metadata={"source": "agibot"},
        )
        assert ep.metadata.get("language_only", False) is False
        assert ep.metadata.get("visual_only", False) is False
