# Copyright (c) Facebook, Inc. and its affiliates.
import copy
import numpy as np
import unittest
from typing import Dict
import torch

from detectron2.config import CfgNode as CfgNode_
from detectron2.config import instantiate
from detectron2.structures import Boxes, Instances
from detectron2.tracking.base_tracker import build_tracker_head
from detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker import (  # noqa
    IOUWeightedHungarianBBoxIOUTracker,
)


class TestIOUWeightedHungarianBBoxIOUTracker(unittest.TestCase):
    def setUp(self):
        self._img_size = np.array([600, 800])
        self._prev_boxes = np.array(
            [
                [101, 101, 200, 200],
                [301, 301, 450, 450],
            ]
        ).astype(np.float32)
        self._prev_scores = np.array([0.9, 0.9])
        self._prev_classes = np.array([1, 1])
        self._prev_masks = np.ones((2, 600, 800)).astype("uint8")
        self._curr_boxes = np.array(
            [
                [302, 303, 451, 452],
                [101, 102, 201, 203],
            ]
        ).astype(np.float32)
        self._curr_scores = np.array([0.95, 0.85])
        self._curr_classes = np.array([1, 1])
        self._curr_masks = np.ones((2, 600, 800)).astype("uint8")

        self._prev_instances = {
            "image_size": self._img_size,
            "pred_boxes": self._prev_boxes,
            "scores": self._prev_scores,
            "pred_classes": self._prev_classes,
            "pred_masks": self._prev_masks,
        }
        self._prev_instances = self._convertDictPredictionToInstance(self._prev_instances)
        self._curr_instances = {
            "image_size": self._img_size,
            "pred_boxes": self._curr_boxes,
            "scores": self._curr_scores,
            "pred_classes": self._curr_classes,
            "pred_masks": self._curr_masks,
        }
        self._curr_instances = self._convertDictPredictionToInstance(self._curr_instances)

        self._max_num_instances = 10
        self._max_lost_frame_count = 3
        self._min_box_rel_dim = 0.02
        self._min_instance_period = 1
        self._track_iou_threshold = 0.5

    def _convertDictPredictionToInstance(self, prediction: Dict) -> Instances:
        """
        convert prediction from Dict to D2 Instances format
        """
        res = Instances(
            image_size=torch.IntTensor(prediction["image_size"]),
            pred_boxes=Boxes(torch.FloatTensor(prediction["pred_boxes"])),
            pred_masks=torch.IntTensor(prediction["pred_masks"]),
            pred_classes=torch.IntTensor(prediction["pred_classes"]),
            scores=torch.FloatTensor(prediction["scores"]),
        )
        return res

    def test_init(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        self.assertTrue(tracker._video_height == self._img_size[0])

    def test_from_config(self):
        cfg = CfgNode_()
        cfg.TRACKER_HEADS = CfgNode_()
        cfg.TRACKER_HEADS.TRACKER_NAME = "IOUWeightedHungarianBBoxIOUTracker"
        cfg.TRACKER_HEADS.VIDEO_HEIGHT = int(self._img_size[0])
        cfg.TRACKER_HEADS.VIDEO_WIDTH = int(self._img_size[1])
        cfg.TRACKER_HEADS.MAX_NUM_INSTANCES = self._max_num_instances
        cfg.TRACKER_HEADS.MAX_LOST_FRAME_COUNT = self._max_lost_frame_count
        cfg.TRACKER_HEADS.MIN_BOX_REL_DIM = self._min_box_rel_dim
        cfg.TRACKER_HEADS.MIN_INSTANCE_PERIOD = self._min_instance_period
        cfg.TRACKER_HEADS.TRACK_IOU_THRESHOLD = self._track_iou_threshold
        tracker = build_tracker_head(cfg)
        self.assertTrue(tracker._video_height == self._img_size[0])

    def test_initialize_extra_fields(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        instances = tracker._initialize_extra_fields(self._curr_instances)
        self.assertTrue(instances.has("ID"))
        self.assertTrue(instances.has("ID_period"))
        self.assertTrue(instances.has("lost_frame_count"))

    def test_process_matched_idx(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        prev_instances = tracker._initialize_extra_fields(self._prev_instances)
        tracker._prev_instances = prev_instances
        curr_instances = tracker._initialize_extra_fields(self._curr_instances)
        matched_idx = np.array([0])
        matched_prev_idx = np.array([1])
        curr_instances = tracker._process_matched_idx(curr_instances, matched_idx, matched_prev_idx)
        self.assertTrue(curr_instances.ID[0] == 1)

    def test_process_unmatched_idx(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        prev_instances = tracker._initialize_extra_fields(self._prev_instances)
        tracker._prev_instances = prev_instances
        curr_instances = tracker._initialize_extra_fields(self._curr_instances)
        matched_idx = np.array([0])
        matched_prev_idx = np.array([1])
        curr_instances = tracker._process_matched_idx(curr_instances, matched_idx, matched_prev_idx)
        curr_instances = tracker._process_unmatched_idx(curr_instances, matched_idx)
        self.assertTrue(curr_instances.ID[1] == 2)

    def test_process_unmatched_prev_idx(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        prev_instances = tracker._initialize_extra_fields(self._prev_instances)
        prev_instances.ID_period = [3, 3]
        tracker._prev_instances = prev_instances
        curr_instances = tracker._initialize_extra_fields(self._curr_instances)
        matched_idx = np.array([0])
        matched_prev_idx = np.array([1])
        curr_instances = tracker._process_matched_idx(curr_instances, matched_idx, matched_prev_idx)
        curr_instances = tracker._process_unmatched_idx(curr_instances, matched_idx)
        curr_instances = tracker._process_unmatched_prev_idx(curr_instances, matched_prev_idx)
        self.assertTrue(curr_instances.ID[2] == 0)

    def test_assign_cost_matrix_values(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        pair1 = {"idx": 0, "prev_idx": 1, "IoU": 0.6}
        pair2 = {"idx": 1, "prev_idx": 0, "IoU": 0.8}
        bbox_pairs = [pair1, pair2]
        cost_matrix = np.full((2, 2), np.inf)
        target_matrix = copy.deepcopy(cost_matrix)
        target_matrix[0, 1] = -0.6
        target_matrix[1, 0] = -0.8
        cost_matrix = tracker.assign_cost_matrix_values(cost_matrix, bbox_pairs)
        self.assertTrue(np.allclose(cost_matrix, target_matrix))

    def test_update(self):
        cfg = {
            "_target_": "detectron2.tracking.iou_weighted_hungarian_bbox_iou_tracker.IOUWeightedHungarianBBoxIOUTracker",  # noqa
            "video_height": self._img_size[0],
            "video_width": self._img_size[1],
            "max_num_instances": self._max_num_instances,
            "max_lost_frame_count": self._max_lost_frame_count,
            "min_box_rel_dim": self._min_box_rel_dim,
            "min_instance_period": self._min_instance_period,
            "track_iou_threshold": self._track_iou_threshold,
        }
        tracker = instantiate(cfg)
        _ = tracker.update(self._prev_instances)
        curr_instances = tracker.update(self._curr_instances)
        self.assertTrue(curr_instances.ID[0] == 1)
        self.assertTrue(curr_instances.ID[1] == 0)


if __name__ == "__main__":
    unittest.main()
