# Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license

from boxmot.utils import logger as LOGGER
from boxmot.utils.checks import RequirementsChecker

checker = RequirementsChecker()

UL_MODELS = ['yolov8', 'yolov9', 'yolov10', 'yolo11', 'yolo12', 'rtdetr', 'sam']


def is_ultralytics_model(yolo_name):
    return any(yolo in str(yolo_name) for yolo in UL_MODELS)


def is_yolox_model(yolo_name):
    return 'yolox' in str(yolo_name)


def default_imgsz(yolo_name):
    if is_ultralytics_model(yolo_name):
        return [640, 640]
    elif is_yolox_model(yolo_name):
        return [800, 1440]
    else:
        return [640, 640]


def get_yolo_inferer(yolo_model):

    if is_yolox_model(yolo_model):
        try:
            import yolox  # for linear_assignment
            assert yolox.__version__
        except (ImportError, AssertionError, AttributeError):
            checker.check_packages(('yolox',))
            checker.check_packages(('tabulate',))  # needed dependency
            checker.check_packages(('thop',))  # needed dependency
        from .yolox import YoloXStrategy
        return YoloXStrategy
    elif 'yolov8' in str(yolo_model):
        # ultralytics already installed when running track.py
        from .yolov8 import Yolov8Strategy
        return Yolov8Strategy
    elif 'rf-detr' in str(yolo_model):
        try:
            import rfdetr
        except (ImportError, AssertionError, AttributeError):
            checker.check_packages(('onnxruntime',))  # needed dependency
            checker.check_packages(('rfdetr',))  # needed dependency
        from .rfdetr import RFDETRStrategy
        return RFDETRStrategy
    elif 'yolo_nas' in str(yolo_model):
        try:
            import super_gradients  # for linear_assignment
            assert super_gradients.__version__
        except (ImportError, AssertionError, AttributeError):
            checker.check_packages(('super-gradients==3.1.3',))  # install
        from .yolonas import YoloNASStrategy
        return YoloNASStrategy
    else:
        LOGGER.error('Failed to infer inference mode from yolo model name')
        LOGGER.error('Your model name has to contain either yolox, yolo_nas or yolov8')
        exit()