"""Tests for the relative action representation (ported from Cosmos-Predict2.5)."""

import math

import numpy as np
import pytest

from neon.data.relative_actions import (
    euler2rotm,
    rotm2euler,
    rotm2quat,
    states_to_relative_actions,
    get_action_sequence_from_states,
)


class TestEuler2Rotm:
    """Test Euler angle to rotation matrix conversion."""

    def test_identity(self):
        """Zero angles should give identity matrix."""
        R = euler2rotm(np.array([0.0, 0.0, 0.0]))
        np.testing.assert_allclose(R, np.eye(3), atol=1e-10)

    def test_90_degree_yaw(self):
        """90° yaw (Z-axis) rotation."""
        R = euler2rotm(np.array([0.0, 0.0, np.pi / 2]))
        expected = np.array([
            [0, -1, 0],
            [1, 0, 0],
            [0, 0, 1],
        ], dtype=float)
        np.testing.assert_allclose(R, expected, atol=1e-10)

    def test_is_rotation_matrix(self):
        """Output should be a valid rotation matrix (det=1, R.T @ R = I)."""
        for _ in range(10):
            angles = np.random.uniform(-np.pi, np.pi, size=3)
            R = euler2rotm(angles)
            np.testing.assert_allclose(np.linalg.det(R), 1.0, atol=1e-10)
            np.testing.assert_allclose(R.T @ R, np.eye(3), atol=1e-10)


class TestRotm2Euler:
    """Test rotation matrix to Euler angles roundtrip."""

    def test_roundtrip(self):
        """euler2rotm → rotm2euler should recover original angles."""
        for _ in range(20):
            original = np.random.uniform(-np.pi / 2, np.pi / 2, size=3)
            R = euler2rotm(original)
            recovered = rotm2euler(R)
            R_recovered = euler2rotm(recovered)
            # Compare rotation matrices (more robust than comparing angles)
            np.testing.assert_allclose(R, R_recovered, atol=1e-8)


class TestRotm2Quat:
    """Test rotation matrix to quaternion."""

    def test_identity(self):
        """Identity rotation → unit quaternion [1, 0, 0, 0]."""
        q = rotm2quat(np.eye(3))
        np.testing.assert_allclose(np.abs(q[0]), 1.0, atol=1e-10)  # w = ±1
        np.testing.assert_allclose(q[1:], 0.0, atol=1e-10)

    def test_unit_quaternion(self):
        """Output quaternion should have unit norm."""
        for _ in range(10):
            angles = np.random.uniform(-np.pi, np.pi, size=3)
            R = euler2rotm(angles)
            q = rotm2quat(R)
            np.testing.assert_allclose(np.linalg.norm(q), 1.0, atol=1e-10)


class TestStatesToRelativeActions:
    """Test the Cosmos-style relative action computation."""

    def test_stationary_robot(self):
        """Stationary robot should produce zero actions."""
        T = 10
        arm_states = np.tile(np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0]), (T, 1))
        gripper_states = np.zeros(T)

        actions = states_to_relative_actions(arm_states, gripper_states, action_scaler=1.0)
        assert actions.shape == (T - 1, 7)
        np.testing.assert_allclose(actions[:, :6], 0.0, atol=1e-10)

    def test_pure_translation(self):
        """Pure X translation should produce nonzero rel_x only."""
        T = 5
        arm_states = np.zeros((T, 6))
        arm_states[:, 0] = np.arange(T) * 0.01  # Move 1cm per step in X
        gripper_states = np.zeros(T)

        actions = states_to_relative_actions(arm_states, gripper_states, action_scaler=1.0)
        assert actions.shape == (T - 1, 7)
        # Since rotation is zero, gripper frame = world frame
        # rel_x should be ~0.01 each step
        np.testing.assert_allclose(actions[:, 0], 0.01, atol=1e-10)
        np.testing.assert_allclose(actions[:, 1:6], 0.0, atol=1e-10)

    def test_output_shape_with_scaler(self):
        """Action scaler should multiply the output."""
        T = 3
        arm_states = np.zeros((T, 6))
        arm_states[:, 0] = [0.0, 0.1, 0.2]
        gripper_states = np.array([0.0, 0.5, 1.0])

        actions_s1 = states_to_relative_actions(arm_states, gripper_states, action_scaler=1.0)
        actions_s20 = states_to_relative_actions(arm_states, gripper_states, action_scaler=20.0)

        # XYZ and rotation should be 20× larger
        np.testing.assert_allclose(actions_s20[:, :6], actions_s1[:, :6] * 20.0, atol=1e-10)

    def test_quaternion_mode(self):
        """use_quat=True should produce 8-dim actions."""
        T = 5
        arm_states = np.random.randn(T, 6).astype(np.float32) * 0.1
        gripper_states = np.random.rand(T).astype(np.float32)

        actions = states_to_relative_actions(
            arm_states, gripper_states, action_scaler=1.0, use_quat=True
        )
        assert actions.shape == (T - 1, 8)

    def test_gripper_preserved(self):
        """Gripper state should be passed through directly."""
        T = 4
        arm_states = np.zeros((T, 6))
        gripper_states = np.array([0.0, 0.3, 0.7, 1.0])

        actions = states_to_relative_actions(
            arm_states, gripper_states, action_scaler=1.0, gripper_scale=1.0
        )
        np.testing.assert_allclose(actions[:, -1], [0.3, 0.7, 1.0], atol=1e-10)


class TestGetActionSequenceFromStates:
    """Test the Cosmos-compatible dict-based API."""

    def test_bridge_format(self):
        """Should work with Bridge dataset format."""
        data = {
            "state": [[0, 0, 0, 0, 0, 0]] * 10,
            "continuous_gripper_state": [0.0] * 10,
        }
        actions = get_action_sequence_from_states(data)
        assert actions.shape == (9, 7)

    def test_fps_downsample(self):
        """FPS downsampling should reduce sequence length."""
        data = {
            "state": [[0, 0, 0, 0, 0, 0]] * 20,
            "continuous_gripper_state": [0.0] * 20,
        }
        actions_1x = get_action_sequence_from_states(data, fps_downsample_ratio=1)
        actions_2x = get_action_sequence_from_states(data, fps_downsample_ratio=2)

        assert actions_1x.shape[0] == 19  # 20 - 1
        assert actions_2x.shape[0] == 9   # 10 - 1
