from __future__ import annotations

import torch

from scripts.fetchreach_benchmark import _obs_to_vec, compute_goal_distance


def test_obs_to_vec_handles_dict_and_tensors() -> None:
    obs_dict = {
        "observation": torch.tensor([1.0, 2.0], dtype=torch.float32),
        "achieved_goal": torch.tensor([0.5, -0.5, 1.0], dtype=torch.float32),
        "desired_goal": torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32),
    }
    v = _obs_to_vec(obs_dict)
    # Should concatenate all parts
    assert v.shape[0] == 2 + 3 + 3
    # Tensor passthrough
    t = torch.randn(7)
    vt = _obs_to_vec(t)
    assert torch.allclose(vt, t.float())


def test_compute_goal_distance_matches_euclidean() -> None:
    a = torch.tensor([1.0, 2.0, 3.0])
    d = torch.tensor([1.5, 1.0, 3.0])
    dist = compute_goal_distance(a, d)
    expected = float(torch.linalg.vector_norm(a - d))
    assert abs(dist - expected) < 1e-6
