import sys
import numpy as np
from pathlib import Path

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]  # root directory
BOXMOT = ROOT / 'boxmot'
EXAMPLES = ROOT / 'examples'
WEIGHTS = ROOT / 'examples' / 'weights'
REQUIREMENTS = ROOT / 'requirements.txt'

# global logger
from loguru import logger
logger.remove()
logger.add(sys.stderr, colorize=True, level="INFO")


class PerClassDecorator:
    def __init__(self, method):
        self.update = method
    def __get__(self, instance, owner):
        def wrapper(*args, **kwargs):
            modified_args = list(args)
            dets = modified_args[0]
            im = modified_args[1]

            # input one class of detections at a time in order to not mix them up
            if instance.per_class is True and dets.size != 0:
                dets_dict = {class_id: np.array([det for det in dets if det[5] == class_id]) for class_id in set(det[5] for det in dets)}
                # get unique classes in predictions
                detected_classes = set(dets_dict.keys())  
                # get unque classes with active trackers
                active_classes = set([tracker.cls for tracker in instance.trackers])
                # get tracks that are both active and in the current detections
                relevant_classes = active_classes.union(detected_classes)
                
                mc_dets = np.empty(shape=(0, 7))
                for class_id in relevant_classes:
                    modified_args[0] = np.array(dets_dict.get(int(class_id), np.empty((0, 6))))
                    logger.debug(f'Feeding class {int(class_id)}: {modified_args[0].shape}')
                    dets = self.update(instance, modified_args[0], im)
                    if dets.size != 0:
                        mc_dets = np.append(mc_dets, dets, axis=0)
            else:
                mc_dets = self.update(instance, dets, im)
            logger.debug(f'Per class updates output: {mc_dets.shape}')
            return mc_dets
        return wrapper