from typing import List, Dict
import cv2
import numpy as np
from models.engine.visualizer import BaseVisualizer

class Visualizer(BaseVisualizer):

    def __init__(self, fps: int=-1, min_width: int=-1):
        """ Visualizer class for visualization (track_results + count_results).

        Args:
            class_map_ids (Dict): class mapping dictionary to map model's class to original class. Eg {0: 1, 1: 0, 2: 2, 3: 3} mean we swap class ID between 0 and 1.
            fps (int): FPS for output video. If fps = -1, it will have same fps as input video.
            min_width (int): minimum width for output video (height will be scaled to keep aspect ratio as input video). If min_width = -1, it will have same resolution as input video.
        """
        class_names = ['pedestrian']
        super().__init__(class_names, fps, min_width)

    def visualize(self, img: np.ndarray, dettrack_at_frame_id: List[Dict]=None, show_conf: bool=True):
        """ Function to visualize (track_results + count_results) a frame.

        Args:
            img (np.ndarray): image need to be visualized.
            dettrack_at_frame_id (List[Dict]): batch of track results which can be obtained from Tracker class.
            show_conf (bool): Visualize confidence of track results or not.
        """
        
         # Draw tracking
        if (dettrack_at_frame_id):
            boxes = dettrack_at_frame_id["boxes"]
            # classes = dettrack_at_frame_id["labels"]
            ids = dettrack_at_frame_id["ids"]
            for bbox, id_ in zip(boxes, ids): 
                id_ = int(id_)
                score = bbox[4]
                color = self.get_color(id_)
                label = f'{id_}' + (f' {score:.2f}' if (show_conf) else '')
                tl, tf = 2, 1
                c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
                img = cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
                t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
                c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
                img = cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
                img = cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

        if (img.shape[0] != self.height or img.shape[1] != self.width):
            img = cv2.resize(img, (self.width, self.height))
        self.writer.write(img)

        return img