# 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 for evaluating mlcvzoo models"""

import argparse
import logging
import os.path
from abc import ABC
from typing import Any, Dict, Generic, List, Optional, Tuple, Type, cast

import mlflow
from config_builder import ConfigBuilder
from mlcvzoo_base.api.data.annotation import BaseAnnotation
from mlcvzoo_base.api.data.annotation_class_mapper import AnnotationClassMapper
from mlcvzoo_base.api.interfaces import Classifiable, NetBased
from mlcvzoo_base.api.model import Model
from mlcvzoo_base.configuration.annotation_handler_config import AnnotationHandlerConfig
from mlcvzoo_base.configuration.mlfow_config import MLFlowConfig
from mlcvzoo_base.configuration.model_config import ModelConfig
from mlcvzoo_base.configuration.structs import MLFlowExperimentTypes
from mlcvzoo_base.data_preparation.annotation_handler import AnnotationHandler
from mlcvzoo_base.metrics.mlflow.mlflow_runner import MLFLowRunner
from mlcvzoo_base.models.model_registry import ModelRegistry
from mlcvzoo_base.utils import get_file_list

from mlcvzoo_util.logger import Logger
from mlcvzoo_util.mlcvzoo_cli_tool import MLCVZooCLITool, configure_model_argparse
from mlcvzoo_util.model_evaluator.configuration import (
    CheckpointConfig,
    ModelEvaluatorCLIConfig,
    ModelEvaluatorConfig,
)
from mlcvzoo_util.model_evaluator.metric_factory import (
    EvaluationMetricType,
    MetricFactory,
    ModelType,
    get_factory,
)
from mlcvzoo_util.model_evaluator.structs import CheckpointInfo, CheckpointLoggingModes

logger = logging.getLogger(__name__)


class ModelEvaluator(
    MLCVZooCLITool[ModelEvaluatorConfig], ABC, Generic[EvaluationMetricType, ModelType]
):
    """
    Tool for evaluating NetBased MLCVZoo models. It gathers all checkpoints that have
    been generated by training run and determines the checkpoint with regard to the type
    of the respective Model.
    There are two kind of usages:
    1. Use it programmatically
    2. Use it as a commandline tool

    Additionally, the training parameters, the calculated metrics and selected checkpoints
    can be logged with mlflow.
    """

    @staticmethod
    def create_configuration(
        yaml_config_path: str,
        string_replacement_map: Optional[Dict[str, str]] = None,
        no_checks: bool = False,
    ) -> ModelEvaluatorConfig:
        """
        Create a ModelEvaluatorConfig

        Args:
            yaml_config_path:  A yaml filepath where to build the configuration
                               object from
            string_replacement_map: A dictionary that defines placeholders which can be used
                                    while parsing the file. They can be understood as variables
                                    that can be used to define configs that are valid across
                                    multiple devices.
            no_checks: Whether the configuration object should be checked for mutual exclusiveness
                       and the "check_values" method for each attribute of the supertype
                       "BaseConfigClass" should be called
        Returns:
            The created configuration
        """

        return cast(
            ModelEvaluatorConfig,
            ConfigBuilder(
                class_type=ModelEvaluatorConfig,
                yaml_config_path=yaml_config_path,
                string_replacement_map=string_replacement_map,
                no_checks=no_checks,
            ).configuration,
        )

    @staticmethod
    def create_cli_configuration(
        string_replacement_map: Optional[Dict[str, str]] = None,
    ) -> ModelEvaluatorCLIConfig:
        """
        Create a ModelEvaluatorConfig

        Args:
            string_replacement_map: A dictionary that defines placeholders which can be used
                                    while parsing the file. They can be understood as variables
                                    that can be used to define configs that are valid across
                                    multiple devices.
        Returns:
            The created configuration
        """

        return MLCVZooCLITool._create_cli_configuration(
            configuration_class=ModelEvaluatorCLIConfig,
            string_replacement_map=string_replacement_map,
            configure_argparse=ModelEvaluator.configure_argparse,
        )

    @staticmethod
    def configure_argparse(
        parser: argparse.ArgumentParser,
    ) -> None:
        parser.description = (
            "Runs an evaluation on the configured models "
            "using the configuration evaluation data"
        )
        MLCVZooCLITool.configure_argparse(parser)
        configure_model_argparse(parser=parser)
        parser.add_argument(
            "--single-mode",
            default=False,
            action="store_true",
            help="Whether to use the checkpoint-config of the ModelEvaluator configuration "
            "or the configured inference checkpoint of the model for determining the"
            "checkpoints that should be evaluated.",
        )

    @staticmethod
    def create_model(
        model_config: ModelConfig,
        string_replacement_map: Optional[Dict[str, str]] = None,
    ) -> ModelType:
        """
        Creates a model based on the given config

        Args:
            model_config: The model config that should be used
                          for creating a model
            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.
        Returns:
            The created model
        """
        # In order to be able to load the checkpoints that
        # should be evaluated, we need an inference model
        model_config.set_inference(inference=True)
        return ModelRegistry().init_model(  # type: ignore[return-value]
            model_config=model_config, string_replacement_map=string_replacement_map
        )

    @staticmethod
    def generate_evaluation_annotations(
        annotation_handler_config: AnnotationHandlerConfig,
        mapper: Optional[AnnotationClassMapper] = None,
    ) -> List[BaseAnnotation]:
        gt_annotations: List[BaseAnnotation] = AnnotationHandler(
            mapper=mapper,
            configuration=annotation_handler_config,
        ).parse_inference_annotations()

        if mapper is not None:
            gt_annotations = AnnotationHandler.reduce_annotations(
                annotations=gt_annotations, mapper=mapper
            )

        return gt_annotations

    @staticmethod
    def select_evaluation_checkpoints(
        checkpoint_config: CheckpointConfig,
    ) -> List[str]:
        checkpoint_paths = get_file_list(
            input_dir=checkpoint_config.checkpoint_dir,
            file_extension=checkpoint_config.checkpoint_filename_suffix,
        )
        # Sort the checkpoints assuming that they will be sorted in ascending order
        # by epoch
        checkpoint_paths.sort()

        return [c for c in checkpoint_paths if c not in checkpoint_config.ignore]

    @staticmethod
    def _restore_checkpoint(checkpoint_path: str, model: Model) -> None:  # type: ignore[type-arg]
        if not os.path.exists(checkpoint_path):
            raise ValueError(
                f"The given checkpoint path '{checkpoint_path}' does not exist! "
                "Only absolute paths of checkpoints are allowed!"
            )

        if not isinstance(model, NetBased):
            raise ValueError(
                "This evaluation can only be used with models that "
                "inherit from 'mlcvzoo.api.interfaces.NetBased'"
            )

        model.restore(checkpoint_path=checkpoint_path)

    def _post_evaluation_step(
        self,
        checkpoint_path: str,
        evaluated_checkpoint_metrics: Dict[str, EvaluationMetricType],
        inference_model: ModelType,
    ) -> None:
        pass

    def evaluate_checkpoints(
        self,
        model_evaluator_config: ModelEvaluatorConfig,
        inference_model: ModelType,
        gt_annotations: List[BaseAnnotation],
        checkpoint_paths: Optional[List[str]],
        log_results: bool = False,
    ) -> Tuple[Dict[str, EvaluationMetricType], CheckpointInfo]:
        """
        Evaluate all given checkpoints according to the given configuration.
        If specified in the model-evaluator-config log selected parameters,
        checkpoints and metrics with mlflow.

        Args:
            model_evaluator_config: ModelEvaluatorConfig for this evaluation run
            inference_model: ObjectDetectionModel instantiated for inference
            gt_annotations: Ground truth annotations created with the mapper of the
                            given inference-model
            checkpoint_paths: List of checkpoints that should be evaluated
            log_results: Whether to log the results to mlflow and tensorboard

        Returns:
            Evaluated checkpoints and their calculated metrics, info about best checkpoint
        """

        metric_factory = get_factory(inference_model=inference_model)

        if metric_factory is None:
            raise ValueError(
                "Could not determine a metric-factory for "
                f"model-type='{type(inference_model)}'"
            )

        evaluated_checkpoint_metrics: Dict[str, EvaluationMetricType] = {}
        if checkpoint_paths:
            for ckpt in checkpoint_paths:
                ModelEvaluator._restore_checkpoint(
                    checkpoint_path=ckpt, model=inference_model
                )

                evaluated_checkpoint_metrics[ckpt] = metric_factory.compute_metrics(
                    inference_model=inference_model,
                    gt_annotations=gt_annotations,
                    model_evaluator_config=model_evaluator_config,
                )

                self._post_evaluation_step(
                    checkpoint_path=ckpt,
                    evaluated_checkpoint_metrics=evaluated_checkpoint_metrics,
                    inference_model=inference_model,
                )
        else:
            # TODO: Use interface method for getting the checkpoint-path of the model
            evaluated_checkpoint_metrics[MetricFactory.MODEL_STATE_INDICATOR] = (
                metric_factory.compute_metrics(
                    inference_model=inference_model,
                    gt_annotations=gt_annotations,
                    model_evaluator_config=model_evaluator_config,
                )
            )

            self._post_evaluation_step(
                checkpoint_path=MetricFactory.MODEL_STATE_INDICATOR,
                evaluated_checkpoint_metrics=evaluated_checkpoint_metrics,
                inference_model=inference_model,
            )

        best_checkpoint_info = metric_factory.determine_best_checkpoint(
            evaluated_checkpoint_metrics=evaluated_checkpoint_metrics,
        )

        logger.info(
            "Evaluation finished, best checkpoint:\n" "\t- ckpt='%s'\n" "\t- ap='%s'\n",
            best_checkpoint_info.path,
            best_checkpoint_info.score,
        )

        if log_results:
            if model_evaluator_config.mlflow_config is not None:
                ModelEvaluator._start_mlflow(
                    mlflow_config=model_evaluator_config.mlflow_config.config,
                    model_name=inference_model.unique_name,
                )
                checkpoint_log_mode: str = (
                    model_evaluator_config.mlflow_config.checkpoint_log_mode
                )
            else:
                checkpoint_log_mode = CheckpointLoggingModes.BEST.value

            metric_factory.log_results(
                checkpoint_log_mode=checkpoint_log_mode,
                evaluated_checkpoint_metrics=evaluated_checkpoint_metrics,
                best_checkpoint=best_checkpoint_info,
                logging_configs=(
                    [model_evaluator_config.tensorboard_logging_config]
                    if model_evaluator_config.tensorboard_logging_config
                    else []
                ),
            )
            if model_evaluator_config.mlflow_config is not None:
                mlflow.end_run()

        return evaluated_checkpoint_metrics, best_checkpoint_info

    @staticmethod
    def _start_mlflow(mlflow_config: MLFlowConfig, model_name: str) -> None:
        runner = MLFLowRunner(configuration=mlflow_config)
        runner.create_mlflow_experiments()

        runner.start_mlflow_run(
            experiment_name=MLFlowExperimentTypes.EVAL,
            run_name=model_name,
        )


def __get_model_checkpoint(model: Model) -> Optional[List[str]]:  # type: ignore
    if isinstance(model, NetBased) and isinstance(model, Model):
        try:
            # TODO: Realize this as an API method?
            return [model.configuration.inference_config.checkpoint_path]
        except AttributeError:
            pass
    else:
        logger.warning(
            "Could not determine the model-checkpoint, "
            "this feature is only available for models that "
            "inherit from 'mlcvzoo.api.model.Model' and 'mlcvzoo.api.interfaces.NetBased"
        )

    return None


def run_evaluation(
    configuration: ModelEvaluatorCLIConfig,
    single_mode: bool,
) -> Tuple[Dict[str, EvaluationMetricType], CheckpointInfo]:
    """
    Runs an generic mlcvzoo evaluation

    Args:
        configuration:
        single_mode: Whether
        model:

    Returns:
        None
    """
    model: Model[Any, Any, Any] = ModelEvaluator.create_model(
        model_config=configuration.model_config,
        string_replacement_map=configuration.string_replacement_map,
    )

    if not isinstance(model, Classifiable) or not isinstance(model, Model):
        raise ValueError(
            "This evaluation can only be used with models that "
            "inherit from 'mlcvzoo.api.model.Model' and 'mlcvzoo.api.interfaces.Classifiable"
        )

    assert isinstance(model, Classifiable) and isinstance(model, Model)
    gt_annotations: List[BaseAnnotation] = (
        ModelEvaluator.generate_evaluation_annotations(
            annotation_handler_config=configuration.annotation_handler_config,
            mapper=model.mapper,
        )
    )

    checkpoint_paths: Optional[List[str]]

    if single_mode:
        checkpoint_paths = __get_model_checkpoint(model=model)
    else:
        checkpoint_paths = ModelEvaluator.select_evaluation_checkpoints(
            checkpoint_config=configuration.checkpoint_config
        )

    if checkpoint_paths is None:
        logger.warning(
            "Could not determine any model checkpoints, therefore model"
            "will be used as is for the evaluation"
        )

    return ModelEvaluator().evaluate_checkpoints(
        model_evaluator_config=configuration,
        inference_model=model,  # type: ignore[arg-type]
        gt_annotations=gt_annotations,
        checkpoint_paths=checkpoint_paths,
        log_results=True,
    )


def main() -> None:
    """
    Entry point when using the mlcvzoo-modelevaluator command line tool.
    (See [tool.poetry.scripts] section in pyproject.toml)

    Returns:
        None
    """

    args = ConfigBuilder.setup_argparse(
        configure_argparse=ModelEvaluator.configure_argparse
    ).parse_args()

    Logger.init_logging_basic(
        log_dir=args.log_dir,
        log_file_postfix="ModelEvaluator",
        no_stdout=False,
        root_log_level=args.log_level,
    )

    configuration = ModelEvaluator.create_cli_configuration()

    configuration.model_config.update_class_type(
        args_dict=vars(args),
    )

    model_type: Optional[Type[Model[Any, Any, Any]]] = ModelRegistry().get_model_type(
        class_type=configuration.model_config.class_type
    )

    if model_type:
        configuration.model_config.update_constructor_parameters(
            args_dict=vars(args),
            model_type=model_type,
        )
    else:
        raise ValueError(
            f"Model type '{configuration.model_config.class_type}' is not registered. "
            f"Use one of '{list(ModelRegistry().get_registered_models().keys())}'",
        )

    run_evaluation(configuration=configuration, single_mode=args.single_mode)


if __name__ == "__main__":
    main()
