"""
YOLOv7 model definitions.
"""

from collections import OrderedDict
from typing import Optional

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


class YOLOv7_Set(ModelInfoSet):
    """Configuration set for YOLOv7 models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7",
                "filename": "yolov7.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
                "anchors": [
                    [12, 16, 19, 36, 40, 28],
                    [36, 75, 76, 55, 72, 146],
                    [142, 110, 192, 243, 459, 401],
                ],
            }
        ),
    )

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


class YOLOv7d6_Set(ModelInfoSet):
    """Configuration set for YOLOv7d6 models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7d6",
                "filename": "yolov7d6.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [1280, 1280],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "anchors": [
                    [19, 27, 44, 40, 38, 94],  # P3/8
                    [96, 68, 86, 152, 180, 137],  # P4/16
                    [140, 301, 303, 264, 238, 542],  # P5/32
                    [436, 615, 739, 380, 925, 792],  # P6/64
                ],
            }
        ),
    )

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


class YOLOv7e6_Set(ModelInfoSet):
    """Configuration set for YOLOv7e6 models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7e6",
                "filename": "yolov7e6.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [1280, 1280],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "anchors": [
                    [19, 27, 44, 40, 38, 94],  # P3/8
                    [96, 68, 86, 152, 180, 137],  # P4/16
                    [140, 301, 303, 264, 238, 542],  # P5/32
                    [436, 615, 739, 380, 925, 792],  # P6/64
                ],
            }
        ),
    )

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


class YOLOv7e6e_Set(ModelInfoSet):
    """Configuration set for YOLOv7e6e models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7e6e",
                "filename": "yolov7e6e.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [1280, 1280],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "anchors": [
                    [19, 27, 44, 40, 38, 94],  # P3/8
                    [96, 68, 86, 152, 180, 137],  # P4/16
                    [140, 301, 303, 264, 238, 542],  # P5/32
                    [436, 615, 739, 380, 925, 792],  # P6/64
                ],
            }
        ),
    )

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


class YOLOv7w6_Set(ModelInfoSet):
    """Configuration set for YOLOv7w6 models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7w6",
                "filename": "yolov7w6.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "numpy",
                },
                "LetterBox": {
                    "img_size": [1280, 1280],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "cv",
                },
            }
        ),
        post_cfg=OrderedDict(
            {
                "task": "object_detection",
                "nc": 80,  # Number of classes
                "anchors": [
                    [19, 27, 44, 40, 38, 94],  # P3/8
                    [96, 68, 86, 152, 180, 137],  # P4/16
                    [140, 301, 303, 264, 238, 542],  # P5/32
                    [436, 615, 739, 380, 925, 792],  # P6/64
                ],
            }
        ),
    )

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


class YOLOv7x_Set(ModelInfoSet):
    """Configuration set for YOLOv7x models."""

    DEFAULT = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/YOLOv7x",
                "filename": "yolov7x.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
                "anchors": [
                    [12, 16, 19, 36, 40, 28],
                    [36, 75, 76, 55, 72, 146],
                    [142, 110, 192, 243, 459, 401],
                ],
            }
        ),
    )

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


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7 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(
            YOLOv7_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7d6 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(
            YOLOv7d6_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7e6 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(
            YOLOv7e6_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7e6e 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(
            YOLOv7e6e_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7w6 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(
            YOLOv7w6_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)


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

    def __init__(
        self,
        local_path: Optional[str] = None,
        model_type: str = "DEFAULT",
        infer_mode: str = "global8",
        product: str = "aries",
    ):
        """Initializes the YOLOv7x 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(
            YOLOv7x_Set,
            local_path=local_path,
            model_type=model_type,
            infer_mode=infer_mode,
            product=product,
        )
        super().__init__(model_cfg, pre_cfg, post_cfg)
