# -*- coding: utf-8 -*-
"""
Created on Sun Dec 31 17:14:37 2023

@author: ym
"""
import numpy as np


class Boxes:
    def __init__(self, boxes, orig_shape=None) -> None:
        """Initialize the Boxes class."""
        if boxes.ndim == 1:
            boxes = boxes[None, :]
        m, n = boxes.shape
        assert n in (6, 7), f'expected `n` in [6, 7], but got {n}'  # xyxy, track_id, conf, cls
        
        '''对每一个box进行编号，利用该编号可以索引对应 feature'''
        self.data = np.concatenate([boxes[:, :4], np.arange(m).reshape(-1, 1), boxes[:, 4:]], axis=-1)
        
        self.orig_shape = orig_shape
    
    def cpu(self):
        """Return a copy of the tensor on CPU memory."""
        return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.cpu(), self.orig_shape)
    def numpy(self):
        """Return a copy of the tensor as a numpy array."""
        return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.numpy(), self.orig_shape)
    
    @property
    def xyxy(self):
        """Return the boxes in xyxy format."""
        return self.data[:, :4]
    
    @property
    def xyxyb(self):
        """Return the boxes in xyxyb format."""    
        return self.data[:, :5]

    @property
    def conf(self):
        """Return the confidence values of the boxes."""
        return self.data[:, -2]

    @property
    def cls(self):
        """Return the class values of the boxes."""
        return self.data[:, -1]



# def boxes_add_fid(tboxes):
#     '''
#         将 bboxes 对应的帧索引添加到 boxes 最后一列
#     Return：
#         bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index]
#     '''
    
#     bboxes = np.empty((0, 8), dtype = np.float32)  
#     for tbox, f in tboxes:
#         data = tbox.numpy()
        
#         frame = f * np.ones([data.shape[0], 1])
#         bbox = np.concatenate([data, frame], axis=1)
#         bboxes = np.concatenate([bboxes, bbox], axis=0)
    
#     return bboxes


def boxes_add_fid(tboxes):
    '''
        将 bboxes 对应的帧索引添加到 boxes 最后一列
    Return：
        bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index]
    '''
    
    bboxes = np.empty((0, 8), dtype = np.float32)  
    for data, f in tboxes:
        frame = f * np.ones([data.shape[0], 1])
        bbox = np.concatenate([data, frame], axis=1)
        bboxes = np.concatenate([bboxes, bbox], axis=0)
    
    return bboxes