"""
YOLOv9 model definitions.
"""

from collections import OrderedDict
from typing import Optional

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


class YOLOv9t_Set(ModelInfoSet):
    """Configuration set for YOLOv9t models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv9t",
                "filename": "yolov9t.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 YOLOv9s_Set(ModelInfoSet):
    """Configuration set for YOLOv9s models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv9s",
                "filename": "yolov9s.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 GELANs_Set(ModelInfoSet):
    """Configuration set for GELANs models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/GELANs",
                "filename": "gelans.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 YOLOv9m_Set(ModelInfoSet):
    """Configuration set for YOLOv9m models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv9m",
                "filename": "yolov9m.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 GELANm_Set(ModelInfoSet):
    """Configuration set for GELANm models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/GELANm",
                "filename": "gelanm.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 YOLOv9c_Set(ModelInfoSet):
    """Configuration set for YOLOv9c models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv9c",
                "filename": "yolov9c.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 GELANc_Set(ModelInfoSet):
    """Configuration set for GELANc models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/GELANc",
                "filename": "gelanc.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 YOLOv9e_Set(ModelInfoSet):
    """Configuration set for YOLOv9e models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv9e",
                "filename": "yolov9e.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 GELANe_Set(ModelInfoSet):
    """Configuration set for GELANe models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/GELANe",
                "filename": "gelane.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 YOLOv9t(MBLT_Engine):
    """YOLOv9t model engine."""

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

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


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

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

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


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

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

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


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

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

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


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

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

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


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

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

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


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

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

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


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

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

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


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

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

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