"""
YOLO11 model definitions.
"""

from collections import OrderedDict
from typing import Optional

from ..utils.types import ModelInfo, ModelInfoSet
from ..wrapper import MBLT_Engine


class YOLO11n_Set(ModelInfoSet):
    """Configuration set for YOLO11n models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLO11n",
                "filename": "yolo11n.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [640, 640],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "nl": 3,  # Number of detection layers
                "reg_max": 16,
            }
        ),
    )

    TURBO = DEFAULT.update_model_cfg(revision="TURBO")


class YOLO11s_Set(ModelInfoSet):
    """Configuration set for YOLO11s models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLO11s",
                "filename": "yolo11s.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [640, 640],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "nl": 3,  # Number of detection layers
                "reg_max": 16,
            }
        ),
    )

    TURBO = DEFAULT.update_model_cfg(revision="TURBO")


class YOLO11m_Set(ModelInfoSet):
    """Configuration set for YOLO11m models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLO11m",
                "filename": "yolo11m.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [640, 640],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "nl": 3,  # Number of detection layers
                "reg_max": 16,
            }
        ),
    )

    TURBO = DEFAULT.update_model_cfg(revision="TURBO")


class YOLO11l_Set(ModelInfoSet):
    """Configuration set for YOLO11l models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLO11l",
                "filename": "yolo11l.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [640, 640],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "nl": 3,  # Number of detection layers
                "reg_max": 16,
            }
        ),
    )

    TURBO = DEFAULT.update_model_cfg(revision="TURBO")


class YOLO11x_Set(ModelInfoSet):
    """Configuration set for YOLO11x models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLO11x",
                "filename": "yolo11x.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [640, 640],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "nl": 3,  # Number of detection layers
                "reg_max": 16,
            }
        ),
    )

    TURBO = DEFAULT.update_model_cfg(revision="TURBO")


class YOLO11n(MBLT_Engine):
    """YOLO11n model engine."""

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLO11n engine.

        Args:
                local_path: Path to a local model file. Defaults to None.
                model_type: Model configuration type. Defaults to "DEFAULT".
                infer_mode: Inference execution mode. Defaults to "global8".
                product: Target hardware product. Defaults to "aries".
        """
        model_cfg, pre_cfg, post_cfg = self._get_configs(
            YOLO11n_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


class YOLO11s(MBLT_Engine):
    """YOLO11s model engine."""

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLO11s engine.

        Args:
                local_path: Path to a local model file. Defaults to None.
                model_type: Model configuration type. Defaults to "DEFAULT".
                infer_mode: Inference execution mode. Defaults to "global8".
                product: Target hardware product. Defaults to "aries".
        """
        model_cfg, pre_cfg, post_cfg = self._get_configs(
            YOLO11s_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


class YOLO11m(MBLT_Engine):
    """YOLO11m model engine."""

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLO11m engine.

        Args:
                local_path: Path to a local model file. Defaults to None.
                model_type: Model configuration type. Defaults to "DEFAULT".
                infer_mode: Inference execution mode. Defaults to "global8".
                product: Target hardware product. Defaults to "aries".
        """
        model_cfg, pre_cfg, post_cfg = self._get_configs(
            YOLO11m_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


class YOLO11l(MBLT_Engine):
    """YOLO11l model engine."""

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLO11l engine.

        Args:
                local_path: Path to a local model file. Defaults to None.
                model_type: Model configuration type. Defaults to "DEFAULT".
                infer_mode: Inference execution mode. Defaults to "global8".
                product: Target hardware product. Defaults to "aries".
        """
        model_cfg, pre_cfg, post_cfg = self._get_configs(
            YOLO11l_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


class YOLO11x(MBLT_Engine):
    """YOLO11x model engine."""

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLO11x engine.

        Args:
                local_path: Path to a local model file. Defaults to None.
                model_type: Model configuration type. Defaults to "DEFAULT".
                infer_mode: Inference execution mode. Defaults to "global8".
                product: Target hardware product. Defaults to "aries".
        """
        model_cfg, pre_cfg, post_cfg = self._get_configs(
            YOLO11x_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)
