from types import SimpleNamespace

import numpy as np

from motrix_envs.locomotion.go1.dreamwaq import (
    DreamWaQCfg,
    DreamWaQTask,
    _command_aligned_progress,
    _sanitize_dof_pos,
)
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask


def test_sanitize_dof_pos_handles_zero_quaternion_with_nonfinite_joints():
    dof_pos = np.zeros((2, 19), dtype=np.float32)
    dof_pos[0, 3:7] = 0.0
    dof_pos[0, 7] = np.inf
    dof_pos[1, 3:7] = [0.0, 0.0, 0.5, 0.5]
    dof_pos[1, 8] = np.nan

    clean = _sanitize_dof_pos(dof_pos)

    assert np.isfinite(clean).all()
    np.testing.assert_allclose(clean[0, 3:7], [0.0, 0.0, 0.0, 1.0])
    np.testing.assert_allclose(np.linalg.norm(clean[:, 3:7], axis=1), 1.0)
    assert clean[0, 7] == 0.0
    assert clean[1, 8] == 0.0


def test_command_aligned_progress_follows_current_command_without_jitter_credit():
    local_vel = np.array(
        [[1.0, 0.0], [-1.0, 0.0], [0.0, 1.0], [-0.5, 0.0]], dtype=np.float32
    )
    commands = np.array(
        [[1.0, 0.0], [-1.0, 0.0], [1.0, 0.0], [1.0, 0.0]], dtype=np.float32
    )

    progress = _command_aligned_progress(local_vel, commands, dt=0.02)

    np.testing.assert_allclose(progress, [0.02, 0.02, 0.0, -0.01])


def test_command_aligned_progress_ignores_standing_commands():
    progress = _command_aligned_progress(
        np.array([[2.0, 0.0]], dtype=np.float32),
        np.array([[0.05, 0.0]], dtype=np.float32),
        dt=0.02,
    )

    np.testing.assert_array_equal(progress, [0.0])


def test_apply_action_advances_three_frame_action_history():
    task = Go1WalkTask.__new__(Go1WalkTask)
    task.get_dof_vel = lambda data: np.ones((1, 12), dtype=np.float32)
    task._compute_torques = lambda actions, data: actions
    state = SimpleNamespace(
        data=SimpleNamespace(actuator_ctrls=None),
        info={
            "current_actions": np.full((1, 12), 2.0, dtype=np.float32),
            "last_actions": np.full((1, 12), 1.0, dtype=np.float32),
            "last_last_actions": np.zeros((1, 12), dtype=np.float32),
        },
    )
    actions = np.full((1, 12), 3.0, dtype=np.float32)

    task.apply_action(actions, state)
    actions.fill(9.0)

    np.testing.assert_array_equal(state.info["last_last_actions"], 1.0)
    np.testing.assert_array_equal(state.info["last_actions"], 2.0)
    np.testing.assert_array_equal(state.info["current_actions"], 3.0)


def test_dreamwaq_feet_air_time_matches_upstream_contact_filtering():
    task = DreamWaQTask.__new__(DreamWaQTask)
    task._cfg = SimpleNamespace(ctrl_dt=0.02)
    task._num_envs = 1
    info = {
        "feet_air_time": np.array([[0.6, 0.0]], dtype=np.float32),
        "contacts": np.array([[True, False]]),
        "last_contacts": np.array([[False, False]]),
    }

    task.update_feet_air_time(info)
    reward = task._reward_feet_air_time(
        np.array([[1.0, 0.0, 0.0]], dtype=np.float32), info
    )

    np.testing.assert_allclose(reward, [0.12])
    np.testing.assert_allclose(info["feet_air_time"], [[0.0, 0.02]])
    np.testing.assert_array_equal(info["last_contacts"], info["contacts"])


def test_dreamwaq_config_matches_upstream_go1_defaults():
    cfg = DreamWaQCfg()

    assert cfg.reward_config.only_positive_rewards
    assert cfg.reward_config.scales["tracking_lin_vel"] == 1.0
    assert cfg.reward_config.scales["tracking_ang_vel"] == 0.5
    assert cfg.reward_config.scales["feet_air_time"] == 0.1
    assert cfg.reward_config.scales["base_height"] == -10.0
    assert cfg.init_state.default_joint_angles["FL_hip"] == 0.1
    assert cfg.init_state.default_joint_angles["RR_hip"] == -0.1