from __future__ import annotations

import torch

from scripts.fetchreach_benchmark import _obs_to_vec
from scripts.her_comparison import _replace_desired


def test_replace_desired_is_pure_and_replaces_only_goal() -> None:
    obs = {
        "observation": torch.tensor([1.0, 2.0], dtype=torch.float32),
        "achieved_goal": torch.tensor([0.5, 0.0], dtype=torch.float32),
        "desired_goal": torch.tensor([9.0, 9.0], dtype=torch.float32),
        "extra": torch.tensor([7.0], dtype=torch.float32),
    }
    new_goal = torch.tensor([3.0, 4.0], dtype=torch.float32)
    out = _replace_desired(obs, new_goal)

    # input not mutated
    assert torch.allclose(obs["desired_goal"], torch.tensor([9.0, 9.0]))
    # output has replacement
    assert torch.allclose(out["desired_goal"], new_goal)
    # non-goal keys preserved
    assert torch.allclose(out["observation"], obs["observation"])  # type: ignore[index]
    assert torch.allclose(out["achieved_goal"], obs["achieved_goal"])  # type: ignore[index]
    assert torch.allclose(out["extra"], obs["extra"])  # type: ignore[index]


def test_obs_vec_reflects_relabelled_goal_segment() -> None:
    obs = {
        "observation": torch.tensor([1.0, 2.0], dtype=torch.float32),
        "achieved_goal": torch.tensor([0.5, 0.0], dtype=torch.float32),
        "desired_goal": torch.tensor([9.0, 9.0], dtype=torch.float32),
    }
    new_goal = torch.tensor([3.0, 4.0], dtype=torch.float32)
    out = _replace_desired(obs, new_goal)
    v = _obs_to_vec(out)
    # Order is observation, achieved_goal, desired_goal
    assert torch.allclose(v[-2:], new_goal)
