import cv2
import numpy as np
import pandas as pd
from scipy.spatial import distance as dist


class MultiTracker(object):
    """
    Multi Object Tracker based on OpenCV algorithms. Mantains
    object identities and assigns ids to newly detected objects.
    """

    def __init__(self, tracker_name, maxDetections=30):
        """
        Args:
        --------------------
        tracker_name: (string) must be one of the available
                    options in self._setup_tracker.
        maxDetections: (int) number of failed trackings an
                    object must get through for it to be
                    considered gone.
        """
        self.tracker_name = tracker_name
        self.maxDetections = maxDetections
       # Fundamental data
        self.data = pd.DataFrame(
            columns=["tracker", "detections", "history", "failures"]
        )

    def _setup_tracker(self):
        """
        (helper function) Defines the tracker that will be used
        in all the multitracking process.
        """
        object_trackers = {'csrt': cv2.TrackerCSRT_create,
                           'kcf': cv2.TrackerKCF_create,
                           'boosting': cv2.TrackerBoosting_create,
                           'mil': cv2.TrackerMIL_create,
                           'tld': cv2.TrackerTLD_create,
                           'medianflow': cv2.TrackerMedianFlow_create,
                           'mosse': cv2.TrackerMOSSE_create}
        return object_trackers[self.tracker_name]()

    def register(self, frame, bbox):
        """
        Registers ids to newly detected objects.

        Args:
        -------------------
        frame: video frame in which new objects will be registered.
        gtruth: groundtruth boxes used for tracker initialization.
        """
        # Initialize tracker
        tracker = self._setup_tracker()
        tracker.init(frame, bbox)
        # Set properties
        new_object = pd.Series({
            "tracker": tracker,
            "detections": bbox,
            "history": np.nan,
            "failures": 0
        })
        self.data = self.data.append(new_object, ignore_index=True)

    def deregister(self, _id):
        """
        Deletes objects that are gone or unable to being tracked.
        """
        self.data = self.data.drop(_id, axis=0)

    def update(self, frame):
        ids = self.data.index
        trackers = self.data.tracker
        for _id, tracker in zip(ids, trackers):
            success, newbox = tracker.update(frame)
            newbox = tuple(int(coord) for coord in newbox)
            if not success:
                failures = self.data.loc[_id]["failures"]
                failures += 1
                self.data.loc[_id]["failures"] = failures
                if failures > self.maxDetections:
                    self.deregister(_id)
            else:
                self.data.loc[_id]["history"] = self.data.loc[_id]["detections"]
                self.data.loc[_id]["detections"] = newbox
        return self.data["detections"]


class IDTracker(MultiTracker):

    def _create_centroid(self, bboxes):
        xmins, ymins, xmaxs, ymaxs = np.transpose(bboxes)
        centroids = np.c_[(
            (xmins + xmaxs) / 2,
            (ymins + ymaxs) / 2)].astype(np.int64)
        return centroids

    def assign_ids(self, frame, detections):
        """
        Mantains object identities as the video goes on, and assigns
        ids to newly detected objects.

        Args:
        --------------------
        frame: video frame in which new objects will be registered.
        detections: bounding boxes generated by an object tracker.
        """
        if len(detections) == 0:
            for _id in self.data.index:
                failures = self.data.loc[_id]["failures"]
                failures += 1
                self.data.loc[_id]["failures"] = failures
                if failures > self.maxDetections:
                    self.deregister(_id)
            return self.data["detections"]

        elif len(self.data["detections"]) == 0:
            for coords in detections:
                self.register(frame, tuple(coords))
            return self.data["detections"]

        inputCentroids = self._create_centroid(detections)
        objectCentroids = self._create_centroid(
            self.data["detections"].values.tolist()
        )
        _ids = self.data.index
        D = dist.cdist(objectCentroids, inputCentroids)
        rows = D.min(axis=1).argsort()
        cols = D.argmin(axis=1)[rows]
        usedRows = set()
        usedCols = set()
        for row, col in zip(rows, cols):
            if row in usedRows or col in usedCols:
                continue
            _id = _ids[row]
            self.data.loc[_id]["detections"] = tuple(detections[col])
            self.data.loc[_id]["failures"] = 0
            usedRows.add(row)
            usedCols.add(col)
        unusedCols = set(range(0, D.shape[1])).difference(usedCols)
        if D.shape[0] < D.shape[1]:
            for col in unusedCols:
                self.register(frame, detections[col])
        return self.data["detections"]
