from types import SimpleNamespace

import numpy as np
import torch

from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap


def test_step_preserves_rewards_and_marks_only_truncations_as_timeouts():
    state = SimpleNamespace(
        obs=np.zeros((2, 3), dtype=np.float32),
        reward=np.array([1.25, -0.5], dtype=np.float32),
        done=np.array([True, True]),
        terminated=np.array([True, False]),
        truncated=np.array([False, True]),
        info={},
    )
    wrapper = RslrlNpEnvWrap.__new__(RslrlNpEnvWrap)
    wrapper._env = SimpleNamespace(step=lambda actions: state)
    wrapper._device = torch.device("cpu")
    wrapper._state = None
    wrapper._num_envs = 2
    wrapper.episode_length_buf = torch.zeros(2, dtype=torch.long)

    _, rewards, dones, extras = wrapper.step(torch.zeros((2, 1)))

    torch.testing.assert_close(rewards, torch.tensor([1.25, -0.5]))
    torch.testing.assert_close(dones, torch.tensor([1.0, 1.0]))
    torch.testing.assert_close(extras["time_outs"], torch.tensor([0.0, 1.0]))