# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_basetrack.ipynb.

# %% auto 0
__all__ = ['TrackState', 'BaseTrack']

# %% ../nbs/02_basetrack.ipynb 3
from collections import OrderedDict
from enum import Enum
import numpy as np

# %% ../nbs/02_basetrack.ipynb 4
class TrackState(Enum):
    """
    Enum class to represent the possible states of a track.

    Attributes:
    
    - New: The initial state of the track when it's just detected.
    - Tracked: The state when the track is being actively tracked.
    - Lost: The state when the track is lost or not detected in the current frame.
    - Removed: The state when the track is removed from the tracking system.
    """
    New = 0
    Tracked = 1
    Lost = 2
    Removed = 3

# %% ../nbs/02_basetrack.ipynb 5
class BaseTrack:
    """
    Base class for a track in a tracking system.
    
    This class provides the basic functionalities and properties for a track, 
    such as its unique ID, state, history, features, and related functions.
    Derived classes should implement the activate, predict, and update methods.

    Attributes:
    
    - `_count` (int): Class-level counter to generate unique IDs for tracks.
    """
    _count = 0  # class-level counter for track IDs

    def __init__(self):
        """
        Initializes a new track with default attributes.
        """
        self.track_id = 0
        self.is_activated = False  # Flag to check if the track is activated
        self.state = TrackState.New  # Initial state for the track

        self.history = OrderedDict()  # Ordered dictionary to store track's movement history
        self.features = []  # List to store features of the track
        self.curr_feature = None  # Feature of the track in the current frame
        self.score = 0  # Confidence score of the track
        self.start_frame = 0  # Frame number when the track was first detected
        self.frame_id = 0  # Current frame number for the track
        self.time_since_update = 0  # Number of frames since the track was last updated

        # multi-camera setup: the location of the track in a global coordinate system
        self.location = (np.inf, np.inf)

    @property
    def end_frame(self
                 ) -> int: # Current frame number.
        """
        Returns the current frame number as the end frame of the track.
        """
        return self.frame_id

    @staticmethod
    def next_id() -> int: # The next unique track ID.
        """
        Generates and returns the next unique track ID.
        """
        BaseTrack._count += 1
        return BaseTrack._count

    def activate(self, *args):
        """
        Activates the track. This method should be implemented by derived classes.

        Raises:
        NotImplementedError: If the method is not implemented by derived classes.
        """
        raise NotImplementedError

    def predict(self):
        """
        Predicts the next state of the track. This method should be implemented by derived classes.

        Raises:
        NotImplementedError: If the method is not implemented by derived classes.
        """
        raise NotImplementedError

    def update(self, *args, **kwargs):
        """
        Updates the state of the track based on new observations.
        This method should be implemented by derived classes.

        Raises:
        NotImplementedError: If the method is not implemented by derived classes.
        """
        raise NotImplementedError

    def mark_lost(self):
        """
        Marks the track as lost.
        """
        self.state = TrackState.Lost

    def mark_removed(self):
        """
        Marks the track as removed.
        """
        self.state = TrackState.Removed
