import numpy as np
import math
import logging as log
import sys
from tqdm import tqdm
from common.feature_distance import calc_features_similarity
from common.common_objects import DetectedObject, validate_detected_object, Bbox
from common.common_objects import get_bbox_center, get_dist, calc_bbox_area
from common.find_best_assignment import solve_assignment_problem
from common.annotation import AnnotationObject, AnnotationStorage

class Track:
    __next_track_id = 0
    def __init__(self, first_obj):
        self.objects = []
        self._track_id = Track.__next_track_id
        Track.__next_track_id += 1

        self.objects.append(first_obj)

    def _validate(self):
        assert len(self.objects) > 0
        for o in self.objects:
            validate_detected_object(o)
        for i in range(len(self.objects) - 1):
            self.objects[i].frame_index < self.objects[i+1].frame_index

    def add_object(self, o):
        self._validate()
        validate_detected_object(o)

        last_frame_index = self.objects[-1].frame_index
        if not last_frame_index < o.frame_index:
            raise RuntimeError("Add object={} to track with the last_frame_index={}".format(o, last_frame_index))

        self.objects.append(o)

    def last(self):
        return self.objects[-1]

    def get_id(self):
        return self._track_id

    def get_bbox_for_frame(self, cur_frame_ind):
        """Finds bbox for frame index using linear approximation"""
        self._validate()
        i_found = None
        for i, o in enumerate(self.objects):
            if o.frame_index == cur_frame_ind:
                return o.bbox
            if o.frame_index > cur_frame_ind:
                i_found = i
                break
        if i_found is None: # cur_frame_ind after the last frame_index in track
            return None

        if i_found == 0: # cur_frame_ind before the first frame_index in track
            return None
        log.debug("using linear approximation for track id={}, frame_index={}".format(self._track_id, cur_frame_ind))
        o1 = self.objects[i_found-1]
        o2 = self.objects[i_found]
        assert o1.frame_index < cur_frame_ind < o2.frame_index

        dindex = o2.frame_index - o1.frame_index
        d_cur_index1 = cur_frame_ind - o1.frame_index
        d_cur_index2 = o2.frame_index - cur_frame_ind

        bbox1 = o1.bbox
        bbox2 = o2.bbox

        res_bbox = [None, None, None, None]
        for k in range(4):
            # linear approximation for all bbox fields
            res_bbox[k] = (bbox1[k] * d_cur_index2 + bbox2[k] * d_cur_index1) / dindex
        res_bbox = Bbox(res_bbox[0], res_bbox[1], res_bbox[2], res_bbox[3])

        return res_bbox

class Tracker:
    def __init__(self, num_frames_to_remove_track, num_objects_to_make_track_valid, affinity_threshold):
        self.tracks = []
        self.track_archive = []
        self.num_frames_to_remove_track = num_frames_to_remove_track
        self.num_objects_to_make_track_valid = num_objects_to_make_track_valid
        self.affinity_threshold = affinity_threshold

    def add_objects(self, det_objs):
        log.debug("begin: handling {} objects".format(len(det_objs)))
        if len(det_objs) == 0:
            return

        frame_index = det_objs[0].frame_index
        assert all(o.frame_index == frame_index for o in det_objs), "All det_objs should have the same frame_index"

        affinity_matrix = self._build_affinity_matrix(det_objs)
        self._validate_affinity_matrix(affinity_matrix, len(self.tracks), len(det_objs))

        self._log_affinity_matrix(affinity_matrix)

        decision, best_affinity = self._solve_assignment_problem(affinity_matrix)
        self._log_decision(decision, best_affinity, det_objs, frame_index)

        self._apply_decision(decision, det_objs, frame_index)
        self._move_obsolete_tracks_to_archive(frame_index)
        log.debug("end: handling {} objects".format(len(det_objs)))

    @staticmethod
    def _validate_affinity_matrix(affinity_matrix, num_tracks, num_det_objs):
        assert isinstance(affinity_matrix, list)
        assert len(affinity_matrix) == num_tracks
        for affinity_row in affinity_matrix:
            assert isinstance(affinity_row, list)
            assert len(affinity_row) == num_det_objs
            assert all(isinstance(v, float) for v in affinity_row)
            assert all(v >= 0 for v in affinity_row)

    def _build_affinity_matrix(self, det_objs):
        affinity_matrix = []
        for t in self.tracks:
            affinity_row = []
            for o in det_objs:
                cur_affinity = self._calc_affinity(t, o)
                affinity_row.append(cur_affinity)

            affinity_matrix.append(affinity_row)

        return affinity_matrix

    def _calc_affinity(self, track, obj):
        affinity_appearance = self._calc_affinity_appearance(track, obj)
        affinity_position = self._calc_affinity_position(track, obj)
        affinity_shape = self._calc_affinity_shape(track, obj)
        return affinity_appearance * affinity_position * affinity_shape

    def _calc_affinity_appearance(self, track, obj):
        raise NotImplementedError("The function _calc_affinity_appearance  is not implemented -- implement it by yourself")

    def _calc_affinity_position(self, track, obj):
        raise NotImplementedError("The function _calc_affinity_position is not implemented -- implement it by yourself")

    def _calc_affinity_shape(self, track, obj):
        raise NotImplementedError("The function _calc_affinity_shape is not implemented -- implement it by yourself")

    @staticmethod
    def _log_affinity_matrix(affinity_matrix):
        with np.printoptions(precision=2, suppress=True, threshold=sys.maxsize, linewidth=sys.maxsize):
            log.debug("Affinity matrix =\n{}".format(np.array(affinity_matrix)))

    def _solve_assignment_problem(self, affinity_matrix):
        decision, best_affinity = solve_assignment_problem(affinity_matrix, self.affinity_threshold)
        return decision, best_affinity

    def _log_decision(self, decision, best_affinity, det_objs, frame_index):
        log.debug("Logging decision for frame index={}".format(frame_index))
        num_tracks = len(self.tracks)
        for track_index in range(num_tracks):
            assert track_index in decision
            obj_index = decision[track_index] # index of the object assigned to the track
            if obj_index is not None:
                assert 0 <= obj_index < len(det_objs)
                obj_bbox = det_objs[obj_index].bbox
            else:
                obj_bbox = None

            cur_best_affinity = best_affinity[track_index]
            if cur_best_affinity is not None:
                best_affinity_str = "{:.3f}".format(cur_best_affinity)
            else:
                best_affinity_str = str(cur_best_affinity)

            log.debug("track_index={}, track id={}, last_bbox={}, decision={}, best_affinity={} => {}".format(
                      track_index, self.tracks[track_index].get_id(),
                      self.tracks[track_index].last().bbox,
                      decision[track_index],
                      best_affinity_str,
                      obj_bbox))

    def _apply_decision(self, decision, det_objs, frame_index):
        set_updated_tracks_indexes = set()
        num_det_objs = len(det_objs)
        num_tracks = len(self.tracks)
        object_indexes_not_mapped_to_tracks = set(range(num_det_objs)) # all indexes from 0 to num_det_objs-1
        for track_index in range(num_tracks):
            assert track_index in decision

            obj_index = decision[track_index] # index of the object assigned to the track
            if obj_index is None:
                # no objects are mapped for this track
                continue

            assert 0 <= obj_index < num_det_objs
            if obj_index not in object_indexes_not_mapped_to_tracks:
                raise RuntimeError("ERROR: Algorithm assigned the object {} to several tracks".format(obj_index))

            object_indexes_not_mapped_to_tracks.remove(obj_index)

            o = det_objs[obj_index]
            self.tracks[track_index].add_object(o)

        # create new tracks for all the objects not mapped to tracks
        for obj_index in object_indexes_not_mapped_to_tracks:
            o = det_objs[obj_index]
            self._create_new_track(o)

    def _create_new_track(self, o):
        new_track = Track(o)
        self.tracks.append(new_track)
        log.debug("created new track: id={} object: frame_index={}, {}".format(
                  new_track.get_id(), o.frame_index, o.bbox))

    def _move_obsolete_tracks_to_archive(self, frame_index):
        new_tracks = []
        for t in self.tracks:
            last_frame_index = t.last().frame_index
            if frame_index - last_frame_index >= self.num_frames_to_remove_track:
                log.debug("Move the track id={} to archive: the current frame_index={}, "
                          "the last frame_index in track={}".format(
                              t.get_id(), frame_index, last_frame_index))
                self.track_archive.append(t)
            else:
                new_tracks.append(t)

        self.tracks = new_tracks

    def is_track_valid(self, track):
        assert isinstance(track, Track)
        return len(track.objects) > self.num_objects_to_make_track_valid

    def get_all_valid_tracks(self):
        res = []
        for t in self.track_archive:
            if self.is_track_valid(t):
                res.append(t)

        for t in self.tracks:
            if self.is_track_valid(t):
                res.append(t)

        return res

def convert_tracks_to_annotation_storage(tracks):
    ann_objects_by_frame_index = {}
    for cur_track in tqdm(tracks, desc="Converting"):
        track_id = cur_track.get_id()

        first_frame_index = cur_track.objects[0].frame_index
        last_frame_index = cur_track.objects[-1].frame_index

        for frame_index in range(first_frame_index, last_frame_index+1):
            bbox = cur_track.get_bbox_for_frame(frame_index)
            tl_x = math.floor(bbox.tl_x)
            tl_y = math.floor(bbox.tl_y)
            br_x = math.ceil(bbox.br_x)
            br_y = math.ceil(bbox.br_y)
            detect_obj = DetectedObject(frame_index=frame_index,
                                        bbox=Bbox(tl_x, tl_y, br_x, br_y),
                                        appearance_feature=[])
            ann_obj = AnnotationObject(detect_obj=detect_obj,
                                       track_id=track_id)
            if frame_index not in ann_objects_by_frame_index:
                ann_objects_by_frame_index[frame_index] = {}

            ann_objects_by_frame_index[frame_index][track_id] = ann_obj

    annotation_objects = []
    for frame_index in sorted(ann_objects_by_frame_index.keys()):
        cur_ann_objects = ann_objects_by_frame_index[frame_index]
        for track_id in sorted(cur_ann_objects.keys()):
            annotation_objects.append(cur_ann_objects[track_id])

    annotation_storage = AnnotationStorage.create_annotation_storage_from_list(annotation_objects)
    return annotation_storage

[((243, 21, 243, 52), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((271, 25, 271, 98), 'common.annotation.AnnotationStorage.create_annotation_storage_from_list', 'AnnotationStorage.create_annotation_storage_from_list', ({(271, 79, 271, 97): 'annotation_objects'}, {}), '(annotation_objects)', False, 'from common.annotation import AnnotationObject, AnnotationStorage\n'), ((30, 8, 30, 35), 'common.common_objects.validate_detected_object', 'validate_detected_object', ({(30, 33, 30, 34): 'o'}, {}), '(o)', False, 'from common.common_objects import DetectedObject, validate_detected_object, Bbox\n'), ((75, 19, 75, 75), 'common.common_objects.Bbox', 'Bbox', ({(75, 24, 75, 35): 'res_bbox[0]', (75, 37, 75, 48): 'res_bbox[1]', (75, 50, 75, 61): 'res_bbox[2]', (75, 63, 75, 74): 'res_bbox[3]'}, {}), '(res_bbox[0], res_bbox[1], res_bbox[2], res_bbox[3])', False, 'from common.common_objects import DetectedObject, validate_detected_object, Bbox\n'), ((150, 34, 150, 100), 'common.find_best_assignment.solve_assignment_problem', 'solve_assignment_problem', ({(150, 59, 150, 74): 'affinity_matrix', (150, 76, 150, 99): 'self.affinity_threshold'}, {}), '(affinity_matrix, self.affinity_threshold)', False, 'from common.find_best_assignment import solve_assignment_problem\n'), ((24, 12, 24, 39), 'common.common_objects.validate_detected_object', 'validate_detected_object', ({(24, 37, 24, 38): 'o'}, {}), '(o)', False, 'from common.common_objects import DetectedObject, validate_detected_object, Bbox\n'), ((146, 13, 146, 102), 'numpy.printoptions', 'np.printoptions', (), '', True, 'import numpy as np\n'), ((251, 19, 251, 40), 'math.floor', 'math.floor', ({(251, 30, 251, 39): 'bbox.tl_x'}, {}), '(bbox.tl_x)', False, 'import math\n'), ((252, 19, 252, 40), 'math.floor', 'math.floor', ({(252, 30, 252, 39): 'bbox.tl_y'}, {}), '(bbox.tl_y)', False, 'import math\n'), ((253, 19, 253, 39), 'math.ceil', 'math.ceil', ({(253, 29, 253, 38): 'bbox.br_x'}, {}), '(bbox.br_x)', False, 'import math\n'), ((254, 19, 254, 39), 'math.ceil', 'math.ceil', ({(254, 29, 254, 38): 'bbox.br_y'}, {}), '(bbox.br_y)', False, 'import math\n'), ((258, 22, 259, 57), 'common.annotation.AnnotationObject', 'AnnotationObject', (), '', False, 'from common.annotation import AnnotationObject, AnnotationStorage\n'), ((147, 53, 147, 78), 'numpy.array', 'np.array', ({(147, 62, 147, 77): 'affinity_matrix'}, {}), '(affinity_matrix)', True, 'import numpy as np\n'), ((256, 45, 256, 73), 'common.common_objects.Bbox', 'Bbox', ({(256, 50, 256, 54): 'tl_x', (256, 56, 256, 60): 'tl_y', (256, 62, 256, 66): 'br_x', (256, 68, 256, 72): 'br_y'}, {}), '(tl_x, tl_y, br_x, br_y)', False, 'from common.common_objects import DetectedObject, validate_detected_object, Bbox\n')]