from __future__ import annotations

import math
import sys
from pathlib import Path

import numpy as np
import pytest

TEST_DIR = Path(__file__).resolve().parent
if str(TEST_DIR) not in sys.path:
    sys.path.insert(0, str(TEST_DIR))

import ik_method_comparison as comparison


def test_slerp_uses_shortest_arc_and_returns_unit_quaternion() -> None:
    start = np.asarray([0.0, 0.0, 0.0, 1.0])
    end = -np.asarray([0.0, 0.0, math.sin(0.1), math.cos(0.1)])

    actual = comparison._slerp_quaternion(start, end, 0.5)

    assert np.linalg.norm(actual) == pytest.approx(1.0)
    assert actual == pytest.approx(
        [0.0, 0.0, math.sin(0.05), math.cos(0.05)]
    )


def test_resample_trajectory_keeps_endpoints_and_uses_requested_rate() -> None:
    trajectory = comparison.EpisodeTrajectory(
        source_path=Path("episode.hdf5"),
        times_s=np.asarray([0.0, 0.5, 1.0]),
        target_poses=np.asarray(
            [
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
                [0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
                [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            ]
        ),
        initial_joints=np.zeros(7),
    )

    actual = comparison.resample_trajectory(trajectory, 4.0)

    assert actual.times_s == pytest.approx([0.0, 0.25, 0.5, 0.75, 1.0])
    assert actual.target_poses[0] == pytest.approx(trajectory.target_poses[0])
    assert actual.target_poses[-1] == pytest.approx(trajectory.target_poses[-1])
    assert actual.target_poses[:, 0] == pytest.approx(actual.times_s)