from typing import Dict

from transformers.configuration_utils import PretrainedConfig
from optimum.exporters.onnx.model_configs import ViTOnnxConfig

MODEL_NAMES = [
        'efficientnet_b0',
        'efficientnet_b1',
        'efficientnet_b2',
        'efficientnet_b3',
        'efficientnet_b4',
        'efficientnet_b5',
        'efficientnet_b6',
        'efficientnet_b7',
        'efficientnet_b8',
        'efficientnet_l2'
    ]


class EfficientNetConfig(PretrainedConfig):
    model_type = 'efficientnet'

    def __init__(
        self,
        model_name: str = 'efficientnet_b0',
        pretrained: bool = False,
        num_classes: int = 1000,
        global_pool: str = 'avg',
        **kwargs,
    ): 
        if model_name not in MODEL_NAMES:
            raise ValueError(f'`model_name` must be one of these: {MODEL_NAMES}, but got {model_name}')
        
        self.model_name = model_name
        self.pretrained = pretrained
        self.num_classes = num_classes
        self.global_pool = global_pool
        super().__init__(**kwargs)


class EfficientNetOnnxConfig(ViTOnnxConfig):
    @property
    def outputs(self) -> Dict[str, Dict[int, str]]:
        common_outputs = super().outputs 

        if self.task == "image-classification":
            common_outputs["logits"] = {0: "batch_size", 1: "num_classes"}
        
        return common_outputs
    

__all__ = [
    'EfficientNetConfig',
    'EfficientNetOnnxConfig'
]