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

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