# -*- coding: utf-8 -*-
"""
Created on Wed Sep 20 14:21:13 2023

@author: ym
"""
import cv2

# import sys
# sys.path.append(r"D:\DeepLearning\yolov5")
# from ultralytics.utils.plotting import Annotator, colors

from .plotting import Annotator, colors

class TrackAnnotator(Annotator): 
    
    def plotting_track(self, track, names='abc'):
        """
        track：[x, y, w, h, track_id, score, cls, frame_index]
        
        boxes: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
                0   1   2   3      4        5     6       7           8
        id：跟踪id，从 1 开始计数，
        frame_index: 帧索引，从 1 开始计数
        cls：类别编号，从 0 开始计数，用作 names 的 key 值
        """
    
        id, cls = track[0, 4],  track[0, 6]
        if id >=0 and cls==0:
            color = colors(int(cls), True)
        elif id >=0 and cls!=0:
            color = colors(int(id), True)   # 不存在 id = 0，不会和上面产生冲突
        else:
            color = colors(19, True)     # 19为调色板的最后一个元素        
        
        nb = track.shape[0]
        for i in range(nb):
            if i == 0: 
                # label = f'{int(track[i, 4])}:({int(track[i, 7])})'
                label = f'ID_{int(track[i, 4])}'
            elif i == nb-1:
                label = ''
                # label = f'{int(track[i, 4])}:({int(track[i, 7])})&{int(nb)}'
            else:
                label = ''
                
            self.circle_label(track[i, :], label, color=color)
    
    
    def circle_label(self, track,  label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
        """
        绘制选定 track 的轨迹
        """
        x, y = int((track[0]+track[2])/2), int((track[1]+track[3])/2)
        
        cv2.circle(self.im, (x, y), 6, color, 2)
        # txt_color = (0,0,0)
        
        if label:
            tf = max(self.lw - 1, 1)  # font thickness
            w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0]  # text width, height
            outside = x + w <= self.im.shape[1]-3
            # p2 = x + w, y - h - 3 if outside else y + h + 3
            # cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA)  # filled
            
            cv2.putText(self.im,
                        label, (x-10 if outside else x-w+2,  y-20),
                        0,
                        # self.lw / 3,
                        self.lw/2,
                        txt_color,
                        thickness=tf,
                        lineType=cv2.LINE_AA)