from dataclasses import dataclass
from typing import Sequence

from playground import Tracking
from beartype.typing import Sequence, Mapping
from jaxtyping import jaxtyped, Float, Int
from jax import Array


@dataclass
class AffinityResult:
    """
    Result of affinity computation between trackings and detections.
    """

    matrix: Float[Array, "T D"]
    """
    Affinity matrix between trackings and detections.
    """

    trackings: Sequence[Tracking]
    """
    Trackings used to compute the affinity matrix.
    """

    indices_T: Sequence[int]
    """
    Indices of the trackings that were used to compute the affinity matrix.
    """

    indices_D: Sequence[int]
    """
    Indices of the detections that were used to compute the affinity matrix.
    """