import numpy as np
import pytest

from kinematic_tracker import NdKkfTracker


def test_get_creation_ids_one(tracker1: NdKkfTracker) -> None:
    ids = tracker1.get_creation_ids(0)
    assert ids == pytest.approx([0])


def test_get_creation_ids_two(tracker1: NdKkfTracker) -> None:
    det_rz = np.linspace(1.0, 12, 12).reshape(2, 6)
    tracker1.advance(2000_000_000, det_rz)
    ids = tracker1.get_creation_ids(0)
    assert ids == pytest.approx([0, 1])


def test_multi_class(tc2: NdKkfTracker) -> None:
    ids0 = tc2.get_creation_ids(0)
    assert ids0.shape == (1,)
    assert ids0.tolist() == [0]

    ids1 = tc2.get_creation_ids(1)
    assert ids1.shape == (2,)
    assert ids1.tolist() == [1, 2]
