from .base import Tracker

import numpy as np


from scipy.optimize import linear_sum_assignment as lisa


class IoUTracker(Tracker):
    def __init__(self):
        self.memory = []
        self.ids = []
        self.accumulative_id = 0
        self.iou_thesh=0.3
        self.no_box_count=0


    def __call__(self, boxes):
        """
        boxes: shape(N, 4), columns(left, top, right, bottom)
        return: an ID list corresponding to the boxes' identity
        """
        if len(boxes) == 0:

            self.no_box_count += 1
            if self.no_box_count == 100:
                self.memory = []
                self.ids = []
                self.accumulative_id=0
                self.no_box_count=0
            return []

        else:
            self.no_box_count = 0

        # if len(boxes) == 0:
        #     # no car
        #     self.memory = []
        #     self.ids = []
        #     self.accumulative_id = 0
        #     return self.ids

        if len(self.memory):
            # subsequent calls
            reference_ids = []
            iou = iou_func(boxes, self.memory)

            iou_max=np.max(iou,1)
            iou_max_idex=np.argmax(iou,1)

            #ti, mi = lisa(-iou)
            
            num_box = len(boxes)
            
            new_ids = np.zeros(num_box, dtype=np.int32)-1
            
            #new_ids[ti] = self.ids[mi]
            
            new_ids[iou_max>self.iou_thesh]=self.ids[iou_max_idex[iou_max>self.iou_thesh]]


            # assign id to each new box
            new_ti=list(np.where(new_ids==-1)[0])
            
            end = len(new_ti) + self.accumulative_id

            new_ids[new_ti] = np.arange(self.accumulative_id, end, dtype=np.int64)
            

            self.accumulative_id = end
            self.ids = new_ids
            self.memory = boxes
        else:
            # the first call
            self.memory = boxes
            end = len(boxes) + self.accumulative_id
            self.ids = np.arange(self.accumulative_id, end, dtype=np.int64)
            self.accumulative_id = end

        return self.ids


def iou_func(boxes1, boxes2, asymmetric=False):
    area1 = (boxes1[:, 3] - boxes1[:, 1]) * (boxes1[:, 2] - boxes1[:, 0])
    area2 = (boxes2[:, 3] - boxes2[:, 1]) * (boxes2[:, 2] - boxes2[:, 0])
    left = np.maximum(boxes1[:, np.newaxis, 0], boxes2[np.newaxis, :, 0])
    top = np.maximum(boxes1[:, np.newaxis, 1], boxes2[np.newaxis, :, 1])
    right = np.minimum(boxes1[:, np.newaxis, 2], boxes2[np.newaxis, :, 2])
    bottom = np.minimum(boxes1[:, np.newaxis, 3], boxes2[np.newaxis, :, 3])
    intersection = np.maximum(right - left, 0) * np.maximum(bottom - top, 0)
    if asymmetric:
        union = area2[np.newaxis]
    else:
        union = area1[:, np.newaxis] + area2[np.newaxis] - intersection
    return intersection / union
