"""
Deep SORT tracker implementation
"""
import numpy as np

from .tracker_utils import cascade_match, iou_cost, min_cost_match, keypoints_to_bbox
from .kalman_filter import KalmanFilter


class Detection:
    """Detection object
    """
    def __init__(self, bbox, cls, feature):
        """
        Args:
            bbox (np.array): object bounding box in format [x_center, y_center, width, height]
            cls (int): object class
            feature (np.ndarray): object feature matrix (keypoints)
        """
        self.bbox = np.asarray(bbox, dtype=np.float32)
        self.cls = cls
        self.feature = np.asarray(feature, dtype=np.float32)
    
    def state_vector(self):
        """Convert object to state vector [x_center, y_center, area]

        Returns:
            np.array: output state vector
        """
        # object centroid
        x, y = self.bbox[:2]
        # object size
        area = np.product(self.bbox[2:])
        state = np.array([x, y, area], dtype=np.float32)
        return state


def detections_from_data(hand_boxes, hand_keypoints, handness, face_boxes, face_keypoints):
    """Convert detected hands and faces and their poses to detection objects

    Args:
        hand_boxes (list): list of detected hand bounding boxes
        hand_keypoints (list): list of hand keypoints of detected hands
        handness (list): list of hand indicators of detected hands
        face_boxes (list): list of detected face bounding boxes
        face_keypoints (list): list of face keypoints of detected faces

    Returns:
        list: list of Detection objects
    """
    detections = list()
    # for every hand object
    for box, kpts, hand in zip(hand_boxes, hand_keypoints, handness):
        # construct bounding box from keypoints
        box = keypoints_to_bbox(kpts[:21, :], scale=1.25)
        # construct detection object
        detections.append(Detection(box, 0, kpts[:21, :]))
        # if two hands are present -> create another Detection object from the second hand
        if np.all(hand):
            box = keypoints_to_bbox(kpts[21:, :], scale=1.25)
            detections.append(Detection(box, 0, kpts[21:, :]))
    # for every face object, create a Detection object
    for box, kpts in zip(face_boxes, face_keypoints):
        detections.append(Detection(box, 1, kpts))
    return detections


class TrackState:
    """Track states enumeration
    """
    Tentative = 1
    Confirmed = 2
    Deleted = 3


class Track:
    """Tracked object
    """
    # unique object id
    track_id = 0
    def __init__(self, init_measurement, cls, feature, n_init=5, max_age=10, dT=None, kalman=True):
        """
        Args:
            init_measurement (np.array): initial object state measurement
            feature (np.ndarray): initial object feature matrix
            n_init (int, optional): number of time steps a track is considered in an unconfirmed state. Defaults to 5.
            max_age (int, optional): maximum time steps in a row a confirmed track is kept without matching with a detection. Defaults to 10.
            dT (float, optional): Kalman filter time step if constant, if kept None - dT is calculated automatically. Defaults to None.
            kalman (bool, optional): flag if kalman filter should be used to filter feature matrix. Defaults to True.
        """
        self.id = Track.track_id
        Track.track_id += 1
        self.bbox = init_measurement
        self.cls = cls
        self.feature = feature
        # object kalman filter for bounding box filtering and optionally for keypoint filtering
        self.box_kalman = KalmanFilter(3, Track.bbox_to_state(init_measurement), dT)
        self.keypoint_kalman = KalmanFilter(feature.size, feature.reshape(-1), dT) if kalman else None
        
        self.n_init = n_init
        self.max_age = max_age
        # init as tentative - unconfirmed state, when hits reach n_init, track is confirmed
        self.state = TrackState.Tentative
        self.hits = 1
        self.time_since_update = 0

    def update(self, measurement):
        """Update track object with measurement

        Args:
            measurement (Detection): matched detection object
        """
        # update bbox kalman filter
        self.box_kalman.update(Track.bbox_to_state(measurement.bbox))
        # update keypoints kalman filter
        if self.keypoint_kalman is not None:
            self.keypoint_kalman.update(measurement.feature.reshape(-1))
        else:
            self.feature = measurement.feature
        # update stats and confirm track if unconfirmed and hits reach n_init
        self.hits += 1
        self.time_since_update = 0
        if self.state == TrackState.Tentative and self.hits >= self.n_init:
            self.state = TrackState.Confirmed
    
    def predict(self):
        """Predict track state

        Returns:
            np.array, np.ndarray: _description_
        """
        # update stats
        self.time_since_update += 1
        # predict bbox
        self.bbox = Track.state_to_bbox(self.box_kalman.predict())
        # predict bbox and features if possible
        if self.keypoint_kalman is not None:
            self.feature = self.keypoint_kalman.predict().reshape(-1, 2)
        return self.bbox, self.feature

    def mark_missed(self):
        """Mark track as missed in current time step
        """
        # if track is unconfirmed -> delete it
        if self.state == TrackState.Tentative:
            self.state = TrackState.Deleted
        # if track is confirmed -> check for maximum age, if crossed -> delete it
        elif self.time_since_update > self.max_age:
            self.state = TrackState.Deleted

    @staticmethod
    def bbox_to_state(bbox):
        """Convert bounding box to state vector [x_center, y_center, area]

        Returns:
            np.array: output state vector
        """
        x, y = bbox[:2]
        area = np.product(bbox[2:])
        state = np.array([x, y, area], dtype=np.float32)
        return state
    
    @staticmethod
    def state_to_bbox(state):
        """Converts measurement state vector to bounding box [x_center, y_center, width, height]

        Args:
            state (np.array): measurements state [x_center, y_center, area]

        Returns:
            np.array: output bounding box
        """
        xy = state[:2]
        # width and height from area
        # area can be negative from kalman filter update -> clip to zero to ensure stability
        wh = np.sqrt(np.clip(state[2:], a_min=0, a_max=None))
        bbox = np.concatenate((xy, wh, wh), axis=-1)
        return bbox

    def is_tentative(self):
        """Check if track is in Tentative state

        Returns:
            bool: True if track's state is Tentative
        """
        return self.state == TrackState.Tentative

    def is_confirmed(self):
        """Check if track is in Confirmed state

        Returns:
            bool: True if track's state is Confirmed
        """
        return self.state == TrackState.Confirmed

    def is_deleted(self):
        """Check if track is in Deleted state

        Returns:
            bool: True if track's state is Deleted
        """
        return self.state == TrackState.Deleted


class Tracker:
    """Object tracker object
    """
    def __init__(self, classes=2, max_age=30, n_init=5, matching_threshold=0.4, max_iou_distance=0.5, dT=None, kalman_classes=(0, )):
        """
        Args:
            classes (int, optional): number of tracked classes. Defaults to 2.
            max_age (int, optional): maximum time steps in a row a confirmed track is kept without matching with a detection. Defaults to 30.
            n_init (int, optional): number of time steps a track is considered in an unconfirmed state. Defaults to 5.
            matching_threshold (float, optional): maximum cost distance of a valid assignment. Defaults to 0.4.
            max_iou_distance (float, optional): maximum IOU distance of a valid IOU assignment. Defaults to 0.5.
            dT (float, optional): Kalman filter time step if constant, if kept None - dT is calculated automatically. Defaults to None.
            kalman_classes (tuple, optional): list of classes which have features filtered by kalman filter. Defaults to (0, ).
        """
        self.classes = classes
        self.max_age = max_age
        self.n_init = n_init
        self.matching_threshold = matching_threshold
        self.max_iou_distance = max_iou_distance
        self.dt = dT
        self.kalman_classes = kalman_classes
        # list of tracks
        self.tracks = list()

    def match(self, tracks, detections):
        """Match tracks with detections

        Args:
            tracks (list): list of tracks to match
            detections (list): list of detections to match

        Returns:
            list, list, list: lists of matched indices, unmatched track indices and unmatched detection indices
        """
        # split track set into confirmed and unconfirmed tracks.
        confirmed_tracks = [i for i, track in enumerate(tracks) if track.is_confirmed()]
        unconfirmed_tracks = [i for i, track in enumerate(tracks) if not track.is_confirmed()]
        # associate confirmed tracks using appearance features.
        matches_cascade, unmatched_tracks_cascade, unmatched_detections = cascade_match(
                self.matching_threshold, self.max_age, tracks, detections, confirmed_tracks, np.arange(len(detections)))
        # associate remaining tracks together with unconfirmed tracks using IOU.
        iou_track_candidates = unconfirmed_tracks + [i for i in unmatched_tracks_cascade if tracks[i].time_since_update == 1]
        unmatched_tracks_cascade = [i for i in unmatched_tracks_cascade if tracks[i].time_since_update != 1]
        matches_iou, unmatched_tracks_iou, unmatched_detections = min_cost_match(
                iou_cost, self.max_iou_distance, tracks, detections, iou_track_candidates, unmatched_detections)
        # combine results
        matches = matches_cascade + matches_iou
        unmatched_tracks = list(set(unmatched_tracks_cascade + unmatched_tracks_iou))
        return matches, unmatched_tracks, unmatched_detections

    def update(self, detections):
        """Update tracker with detections

        Args:
            detections (list): list of detected objects
        """
        # assign objects of classes separately
        for cls in range(self.classes):
            # detections and tracks of current class
            detections_cls = [detection for detection in detections if detection.cls == cls]
            tracks_cls = [track for track in self.tracks if track.cls == cls]
            # match tracks and detections
            matches, unmatched_tracks, unmatched_detections = self.match(tracks_cls, detections_cls)
            # update matched tracks
            for track_idx, detection_idx in matches:
                tracks_cls[track_idx].update(detections_cls[detection_idx])
            # mark unmatched tracks as missed
            for track_idx in unmatched_tracks:
                tracks_cls[track_idx].mark_missed()
            # create new track for each unmatched detection
            for detection_idx in unmatched_detections:
                detection = detections_cls[detection_idx]
                track = Track(detection.bbox, detection.cls, detection.feature, n_init=self.n_init, dT=self.dt, kalman=detection.cls in self.kalman_classes)
                self.tracks.append(track)
        # delete tracks masked as Deleted
        self.tracks = [t for t in self.tracks if not t.is_deleted()]

    def predict(self):
        """Predict tracked object states and features

        Returns:
            list, list, list: list of predicted bounding boxes, feature matrices and track ages
        """
        boxes, features, ages = list(), list(), list()
        # for each class
        for cls in range(self.classes):
            boxes_cls, features_cls, ages_cls = list(), list(), list()
            # tracks of current class
            tracks = (track for track in self.tracks if track.cls == cls)
            # for each track
            for track in tracks:
                # predict bbox and keypoints
                box, keypoints = track.predict()
                # output only if track is confirmed
                if track.is_confirmed():
                    # if any value of bounding box is NAN -> skip
                    if np.isnan(box).any():
                        continue
                    boxes_cls.append(box)
                    features_cls.append(keypoints)
                    ages_cls.append(track.hits)
            boxes.append(np.array(boxes_cls))
            features.append(np.array(features_cls))
            ages.append(np.array(ages_cls))
        return boxes, features, ages
