# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_strack.ipynb.

# %% auto 0
__all__ = ['STrack']

# %% ../nbs/01_strack.ipynb 3
import numpy as np
from .kalman_filter import KalmanFilter
from .basetrack import BaseTrack, TrackState

# %% ../nbs/01_strack.ipynb 4
class STrack(BaseTrack):
    """A Simple Track object"""
    # Shared Kalman Filter for all instances of STrack
    shared_kalman = KalmanFilter()

    def __init__(self, 
                 tlwh, # List or array-like containing four values (top-left x, top-left y, width, height).
                 score # The confidence score of the detection.
                ):
        """
        Initialize an STrack (Simple Track) object.
        """
        # Convert tlwh to numpy array of type float
        self._tlwh = np.asarray(tlwh, dtype=float)
        
        # Kalman filter instance for this track
        self.kalman_filter = None
        
        # Mean and covariance of the state
        self.mean, self.covariance = None, None
        
        # Flag to check if the track is activated
        self.is_activated = False
        
        # Score of the detection
        self.score = score
        
        # Length of the tracklet
        self.tracklet_len = 0

    def predict(self):
        """
        Predict the next state using Kalman filter.
        """
        mean_state = self.mean.copy()
        
        # If the track is not in Tracked state, set velocity to 0
        if self.state != TrackState.Tracked:
            mean_state[7] = 0
        
        # Predict the next state using the kalman filter
        self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)

    @staticmethod
    def multi_predict(stracks # List of STrack objects.
                     ):
        """
        Predict the state for multiple tracks simultaneously using the shared Kalman filter.
        """
        if not stracks:
            return

        # Extract mean and covariance for each track
        multi_states = [(st.mean.copy(), st.covariance) for st in stracks]
        
        # For each track, set velocity to 0 if it's not in Tracked state
        for i, st in enumerate(stracks):
            if st.state != TrackState.Tracked:
                multi_states[i][0][7] = 0
        
        # Use shared kalman filter for prediction
        multi_means, multi_covariances = zip(*multi_states)
        multi_means, multi_covariances = STrack.shared_kalman.multi_predict(np.asarray(multi_means), np.asarray(multi_covariances))

        # Update each track with the predicted mean and covariance
        for i, st in enumerate(stracks):
            st.mean, st.covariance = multi_means[i], multi_covariances[i]

    def activate(self, 
                 kalman_filter, # KalmanFilter instance.
                 frame_id # ID of the current frame.
                ):
        """
        Activate the track. 
        """
        self.kalman_filter = kalman_filter
        
        # Initialize track ID
        self.track_id = self.next_id()
        
        # Initiate mean and covariance using the kalman filter
        self.mean, self.covariance = self.kalman_filter.initiate(self.tlwh_to_xyah(self._tlwh))
        
        self.tracklet_len = 0
        self.state = TrackState.Tracked
        
        # Check if track is activated in the first frame
        self.is_activated = frame_id == 1
        self.frame_id = frame_id
        self.start_frame = frame_id

    def re_activate(self, 
                    new_track, # The new STrack object with updated details.
                    frame_id, # ID of the current frame.
                    new_id=False # Flag to determine if a new ID should be assigned.
                   ):
        """
        Reactivate a track with new details.
        """
        self._update_track(new_track, frame_id, new_id)

    def update(self, new_track, frame_id):
        """
        Update the track with new details.

        :param new_track: The new STrack object with updated details.
        :param frame_id: ID of the current frame.
        """
        self._update_track(new_track, frame_id)

    def _update_track(self, 
                      new_track, # The new STrack object with updated details.
                      frame_id, # ID of the current frame.
                      new_id=False # Flag to determine if a new ID should be assigned.
                     ):
        """
        Internal method to update track details.
        """
        self.frame_id = frame_id
        self.tracklet_len += 1
        
        # Update mean and covariance using the kalman filter
        self.mean, self.covariance = self.kalman_filter.update(self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh))
        self.state = TrackState.Tracked
        self.is_activated = True
        
        # Assign new track ID if required
        if new_id:
            self.track_id = self.next_id()
        
        # Update score
        self.score = new_track.score

    @staticmethod
    def tlwh_to_xyah(tlwh # Bounding box in tlwh format.
                    ): # Bounding box in xyah format.
        """
        Convert bounding box from (top-left x, top-left y, width, height) to (center x, center y, aspect ratio, height). 
        """
        ret = np.asarray(tlwh).copy()
        ret[:2] += ret[2:] / 2
        ret[2] /= ret[3]
        return ret

    @staticmethod
    def tlwh_to_tlbr(tlwh # Bounding box in tlwh format.
                    ): # Bounding box in tlbr format.
        """
        Convert bounding box from (top-left x, top-left y, width, height) to (top-left x, top-left y, bottom-right x, bottom-right y).
        """
        ret = np.asarray(tlwh).copy()
        ret[2:] += ret[:2]
        return ret

    @staticmethod
    def tlbr_to_tlwh(tlbr # Bounding box in tlbr format.
                    ): # Bounding box in tlwh format.
        """
        Convert bounding box from (top-left x, top-left y, bottom-right x, bottom-right y) to (top-left x, top-left y, width, height).
        """
        ret = np.asarray(tlbr).copy()
        ret[2:] -= ret[:2]
        return ret

    @property
    def tlwh(self
            ): # Bounding box in tlwh format.
        """
        Get bounding box in (top-left x, top-left y, width, height) format.
        """
        if self.mean is None:
            return self._tlwh.copy()
        ret = self.mean[:4].copy()
        ret[2] *= ret[3]
        ret[:2] -= ret[2:] / 2
        return ret

    @property
    def tlbr(self
            ): # Bounding box in tlbr format.
        """
        Get bounding box in (top-left x, top-left y, bottom-right x, bottom-right y) format.
        """
        return self.tlwh_to_tlbr(self.tlwh)

    def __repr__(self
                ): # String representation of the track.
        """
        String representation of the STrack object.
        """
        return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame, self.end_frame)
