from collections import OrderedDict
from typing import Optional

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


class EfficientNet_B0_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B0 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B0",
                "filename": "efficientnet_b0_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 256,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [224, 224],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B1_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B1 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B1.tv1_in1k",
                "filename": "efficientnet_b1_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 256,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [240, 240],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    IMAGENET1K_V2 = IMAGENET1K_V1.update_model_cfg(
        repo_id="mobilint/EfficientNet_B1.tv2_in1k",
        filename="efficientnet_b1_IMAGENET1K_V2.mxq",
    ).update_pre_cfg(
        Resize={
            "size": 255,
            "interpolation": "bilinear",
        },
    )
    DEFAULT = IMAGENET1K_V2


class EfficientNet_B2_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B2 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B2",
                "filename": "efficientnet_b2_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 288,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [288, 288],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B3_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B3 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B3",
                "filename": "efficientnet_b3_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 320,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [300, 300],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B4_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B4 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B4",
                "filename": "efficientnet_b4_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 384,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [380, 380],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B5_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B5 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B5",
                "filename": "efficientnet_b5_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 456,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [456, 456],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B6_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B6 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B6",
                "filename": "efficientnet_b6_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 528,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [528, 528],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


class EfficientNet_B7_Set(ModelInfoSet):
    """Configuration set for EfficientNet_B7 models."""

    IMAGENET1K_V1 = ModelInfo(
        model_cfg=OrderedDict(
            {
                "repo_id": "mobilint/EfficientNet_B7",
                "filename": "efficientnet_b7_IMAGENET1K_V1.mxq",
                "revision": "main",
            }
        ),
        pre_cfg=OrderedDict(
            {
                "Reader": {
                    "style": "pil",
                },
                "Resize": {
                    "size": 600,
                    "interpolation": "bicubic",
                },
                "CenterCrop": {
                    "size": [600, 600],
                },
                "SetOrder": {"shape": "HWC"},
                "Normalize": {
                    "style": "torch",
                },
            }
        ),
        post_cfg=OrderedDict({"task": "image_classification"}),
    )
    DEFAULT = IMAGENET1K_V1


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

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


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

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


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

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


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

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


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

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


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

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


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

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


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

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