# Copyright Open Logistics Foundation
#
# Licensed under the Open Logistics Foundation License 1.3.
# For details on the licensing terms, see the LICENSE file.
# SPDX-License-Identifier: OLFL-1.3

"""Module to provide a registry for model and configuration constructors"""

import copy
import inspect
import logging
import os
from typing import Any, Dict, List, Optional, Type, cast

from attr import define
from config_builder import BaseConfigClass, ConfigBuilder
from related import ChildField, SequenceField

from mlcvzoo_base.api.model import ConfigurationType, DataType, Model, PredictionType
from mlcvzoo_base.api.registry import MLCVZooRegistry, RegistryEntry
from mlcvzoo_base.configuration.config_registry import ConfigRegistry
from mlcvzoo_base.configuration.model_config import ModelConfig
from mlcvzoo_base.configuration.replacement_config import (
    STRING_REPLACEMENT_MAP_KEY,
    ReplacementConfig,
)

logger = logging.getLogger(__name__)


@define
class ModelRegistryEntry(BaseConfigClass):
    model_entry: RegistryEntry = ChildField(cls=RegistryEntry)
    config_entry: RegistryEntry = ChildField(cls=RegistryEntry)


@define
class ModelRegistryConfig(BaseConfigClass):
    registry_entries: List[ModelRegistryEntry] = SequenceField(cls=ModelRegistryEntry)


_default_registry_config: ModelRegistryConfig = ModelRegistryConfig(
    registry_entries=[
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="read_from_file_classification",
                module_constructor="ReadFromFileClassificationModel",
                package_name="mlcvzoo_base.models.read_from_file.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="read_from_file_config",
                module_constructor="ReadFromFileConfig",
                package_name="mlcvzoo_base.models.read_from_file.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="read_from_file_object_detection",
                module_constructor="ReadFromFileObjectDetectionModel",
                package_name="mlcvzoo_base.models.read_from_file.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="read_from_file_config",
                module_constructor="ReadFromFileConfig",
                package_name="mlcvzoo_base.models.read_from_file.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="read_from_file_segmentation",
                module_constructor="ReadFromFileSegmentationModel",
                package_name="mlcvzoo_base.models.read_from_file.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="read_from_file_config",
                module_constructor="ReadFromFileConfig",
                package_name="mlcvzoo_base.models.read_from_file.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="yolox",
                module_constructor="YOLOXModel",
                package_name="mlcvzoo_yolox.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="yolox_config",
                module_constructor="YOLOXConfig",
                package_name="mlcvzoo_yolox.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="darknet_object_detection",
                module_constructor="DarknetDetectionModel",
                package_name="mlcvzoo_darknet.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="darknet_config",
                module_constructor="DarknetConfig",
                package_name="mlcvzoo_darknet.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmdetection_object_detection",
                module_constructor="MMObjectDetectionModel",
                package_name="mlcvzoo_mmdetection.object_detection_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmdet_config",
                module_constructor="MMDetectionConfig",
                package_name="mlcvzoo_mmdetection.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmdetection_segmentation",
                module_constructor="MMSegmentationModel",
                package_name="mlcvzoo_mmdetection.segmentation_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmdet_config",
                module_constructor="MMDetectionConfig",
                package_name="mlcvzoo_mmdetection.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmrotate",
                module_constructor="MMRotateModel",
                package_name="mlcvzoo_mmrotate.rotated_object_detection_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmdet_config",
                module_constructor="MMDetectionConfig",
                package_name="mlcvzoo_mmdetection.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmpretrain",
                module_constructor="MMPretrainModel",
                package_name="mlcvzoo_mmpretrain.classification_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmpretrain_config",
                module_constructor="MMPretrainConfig",
                package_name="mlcvzoo_mmpretrain.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmocr_text_detection",
                module_constructor="MMOCRTextDetectionModel",
                package_name="mlcvzoo_mmocr.text_detection_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmocr_config",
                module_constructor="MMOCRConfig",
                package_name="mlcvzoo_mmocr.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="mmocr_text_recognition",
                module_constructor="MMOCRTextRecognitionModel",
                package_name="mlcvzoo_mmocr.text_recognition_model",
            ),
            config_entry=RegistryEntry(
                module_type_name="mmocr_config",
                module_constructor="MMOCRConfig",
                package_name="mlcvzoo_mmocr.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="tf_classification_custom_block",
                module_constructor="CustomBlockModel",
                package_name="mlcvzoo_tf_classification.custom_block.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="tf_classification_custom_block_config",
                module_constructor="CustomBlockConfig",
                package_name="mlcvzoo_tf_classification.custom_block.configuration",
            ),
        ),
        ModelRegistryEntry(
            model_entry=RegistryEntry(
                module_type_name="tf_classification_xception",
                module_constructor="XceptionModel",
                package_name="mlcvzoo_tf_classification.xception.model",
            ),
            config_entry=RegistryEntry(
                module_type_name="tf_classification_xception_config",
                module_constructor="XceptionConfig",
                package_name="mlcvzoo_tf_classification.xception.configuration",
            ),
        ),
    ]
)


class ModelRegistry(MLCVZooRegistry[Type[Model[PredictionType, ConfigurationType, DataType]]]):
    """
    Class to provide a registry for model and configuration constructors
    """

    def __init__(self, configuration_path: Optional[str] = None) -> None:
        MLCVZooRegistry.__init__(self)

        if configuration_path is None:
            configuration_path = os.getenv("MLCVZOO_MODEL_REGISTRY_PATH", None)

        registry_config: ModelRegistryConfig
        if configuration_path is not None and os.path.isfile(configuration_path):
            registry_config = cast(
                ModelRegistryConfig,
                ConfigBuilder(
                    class_type=ModelRegistryConfig, yaml_config_path=configuration_path
                ).configuration,
            )
        else:
            registry_config = _default_registry_config

        self._config_registry = ConfigRegistry()

        self._model_to_config_dict: Dict[str, str] = {}

        for registry_entry in registry_config.registry_entries:
            registered_model_type_name = self.register_external_module(
                module_type_name=registry_entry.model_entry.module_type_name,
                module_constructor=registry_entry.model_entry.module_constructor,
                package_name=registry_entry.model_entry.package_name,
                force=True,
            )
            registered_config_type_name = self._config_registry.register_external_module(
                module_type_name=registry_entry.config_entry.module_type_name,
                module_constructor=registry_entry.config_entry.module_constructor,
                package_name=registry_entry.config_entry.package_name,
                force=True,
            )
            if registered_model_type_name and registered_config_type_name:
                self._model_to_config_dict[registry_entry.model_entry.module_type_name] = (
                    registry_entry.config_entry.module_type_name
                )

    @property
    def model_registry(
        self,
    ) -> Dict[str, Type[Model[PredictionType, ConfigurationType, DataType]]]:
        return self._registry

    @property
    def config_registry(self) -> Dict[str, Type[BaseConfigClass]]:
        return self._config_registry._registry

    def determine_config_class(self, model_type_name: str) -> Type[BaseConfigClass]:
        return self.config_registry[self._model_to_config_dict[model_type_name]]

    def determine_config_class_name(self, model_type_name: str) -> str:
        return self._model_to_config_dict[model_type_name]

    def register_external_model(
        self, model_type_name: str, model_constructor: str, package_name: str
    ) -> None:
        """
        Register an external model
        Args:
            model_type_name: name of the model to register
            model_constructor: name of the constructor of the model to register
            package_name: the full package to import to call the constructor

        Returns:
            Nothing
        """
        self.register_external_module(
            module_type_name=model_type_name,
            module_constructor=model_constructor,
            package_name=package_name,
        )

    def get_registered_models(
        self,
    ) -> Dict[str, Type[Model[PredictionType, ConfigurationType, DataType]]]:
        return copy.deepcopy(self._registry)

    def get_registered_model_type_names(
        self,
    ) -> List[str]:
        return copy.deepcopy(list(self._registry.keys()))

    def get_model_type(
        self, class_type: str
    ) -> Optional[Type[Model[PredictionType, ConfigurationType, DataType]]]:
        if class_type in self._registry:
            return self._registry[class_type]

        return None

    def register_model(
        self, model_type_name: str, model_constructor: Any, force: bool = False
    ) -> None:
        self.register_module(
            module_type_name=model_type_name, module_constructor=model_constructor, force=force
        )

    def build_model_config(
        self,
        model_type_name: str,
        model_configuration_dict: Dict[str, Any],
        string_replacement_map: Optional[Dict[str, str]] = None,
    ) -> ModelConfig:
        return ModelConfig(
            class_type=model_type_name,
            constructor_parameters={
                "configuration": ConfigBuilder.build_configuration_from_dict(
                    config_dict=model_configuration_dict,
                    class_type=self.determine_config_class(model_type_name=model_type_name),
                    string_replacement_map=string_replacement_map,
                ),
            },
        )

    def init_model(
        self,
        model_config: ModelConfig,
        string_replacement_map: Optional[Dict[str, str]] = None,
    ) -> Model[PredictionType, ConfigurationType, DataType]:
        """
        Generic method for instantiating any model that is registered in the model-registry

        Args:
            model_config: The model configuration defining which model should be initialized
            string_replacement_map: (Optional) A dictionary that defines placeholders which can
                                    be used while parsing a configuration file. They can be
                                    understood as variables that can be used to define configs
                                    that are valid across multiple devices.

                                    If no string_replacement_map a default map based on the
                                    ReplacementConfig will be created and used. This allows
                                    to use the attributes of the ReplacementConfig to be
                                    replaced by os environment variables.

        Returns:
            The created model instance
        """

        model_type = self.get_model_type(class_type=model_config.class_type)
        if model_type is not None:
            model: Model[PredictionType, ConfigurationType, DataType]

            init_params: List[Any] = list(inspect.getfullargspec(model_type.__init__).args)
            # We don't need self as parameter in the configuration
            init_params.remove("self")

            if (
                STRING_REPLACEMENT_MAP_KEY not in model_config.constructor_parameters
                and STRING_REPLACEMENT_MAP_KEY in init_params
            ):
                if string_replacement_map is None:
                    string_replacement_map = ReplacementConfig().to_dict()
                model_config.constructor_parameters[STRING_REPLACEMENT_MAP_KEY] = (
                    string_replacement_map
                )

            try:
                model = model_type(**model_config.constructor_parameters)  # type: ignore[arg-type]
            except TypeError as e:
                logger.error(
                    "Please provide the parameters %s, as specified for %s",
                    init_params,
                    model_type,
                )
                raise e
        else:
            message = (
                f"The model '{model_config.class_type}' is not registered! \n"
                f"The registered models are: {self._registry.keys()}"
            )

            logger.error(message)
            raise ValueError(message)

        return model
