from dataclasses import dataclass
from typing import Sequence, Callable, Generator

from app.camera import Detection
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.
    """

    detections: Sequence[Detection]
    """
    Detections used to compute the affinity matrix.
    """

    indices_T: Sequence[int]
    indices_D: Sequence[int]

    def tracking_detections(self) -> Generator[tuple[Tracking, Detection]]:
        for t, d in zip(self.indices_T, self.indices_D):
            yield (self.trackings[t], self.detections[d])