# Copyright ServiceNow, Inc. 2021 – 2022
# This source code is licensed under the Apache 2.0 license found in the LICENSE file
# in the root directory of this source tree.
from typing import List, Optional, Set, cast

import numpy as np
from datasets import Dataset
from tqdm import tqdm

from azimuth.config import ModelContractConfig, PipelineDefinition
from azimuth.dataset_split_manager import DatasetSplitManager, PredictionTableKey
from azimuth.modules.base_classes import ArtifactManager, ConfigScope, DaskModule
from azimuth.modules.base_classes.dask_module import Worker
from azimuth.types import DatasetColumn, DatasetSplitName, ModuleOptions, ModuleResponse
from azimuth.types.general.module_arguments import ModuleEffectiveArguments
from azimuth.utils.conversion import md5_hash
from azimuth.utils.exclude_fields_from_cache import exclude_fields_from_cache
from azimuth.utils.validation import assert_not_none


class Module(DaskModule[ConfigScope]):
    """Abstract class to define functions used by all modules for computing results and interfacing
    with the dataset and the pipelines."""

    required_mod_options: Set[str] = set()
    optional_mod_options: Set[str] = set()

    def __init__(
        self,
        dataset_split_name: DatasetSplitName,
        config: ConfigScope,
        mod_options: Optional[ModuleOptions] = None,
    ):
        mod_options = mod_options or ModuleOptions()
        self.mod_options = mod_options
        defined_mod_options = set(self.mod_options.dict(exclude_defaults=True).keys())
        if not self.required_mod_options.issubset(defined_mod_options):
            raise ValueError(f"{self.__class__.__name__} requires {self.required_mod_options}.")
        if diff := (defined_mod_options - self.required_mod_options - self.optional_mod_options):
            raise ValueError(f"Unexpected mod_options {diff} for {self.__class__.__name__}.")

        self.model_contract_method_name = mod_options.model_contract_method_name
        self.task_name = self.model_contract_method_name or self.__class__.__name__
        super().__init__(dataset_split_name, config)

    def get_effective_arguments(self) -> ModuleEffectiveArguments:
        """Retrieve arguments affecting the Module, i.e. module options and config scope.

        Returns:
            Module's Effective Arguments
        """

        # indices are excluded, since the cache for all indices should be in the same file.
        # model_contract_method_name are excluded too because it's already in the task_name.
        return ModuleEffectiveArguments(
            mod_options=self.mod_options.dict(
                exclude={"indices", "model_contract_method_name"},
                include=self.required_mod_options | self.optional_mod_options,
            ),
            config_scope=self.config.dict(
                exclude=exclude_fields_from_cache(self.config),
            ),
        )

    def _get_name(self) -> str:
        effective_arguments = self.get_effective_arguments()
        hash_options = md5_hash(effective_arguments.mod_options)[:5]
        hash_attributes = md5_hash(effective_arguments.config_scope)[:5]

        return f"{self.task_name}_{self.dataset_split_name}_{hash_options}_{hash_attributes}"

    def get_caching_indices(self) -> List[int]:
        return self.mod_options.indices or self.get_indices()

    def get_indices(self, name: Optional[DatasetSplitName] = None) -> List[int]:
        """Get indices from dataset_split."""
        return cast(List[int], self.get_dataset_split(name)[DatasetColumn.row_idx])

    @property
    def artifact_manager(self):
        """This is set as a property so the Module always have access to the current version of
        the ArtifactManager on the worker."""
        return ArtifactManager.instance()

    @property
    def available_dataset_splits(self) -> Set[DatasetSplitName]:
        dataset_dict = self.artifact_manager.get_dataset_dict(self.config)
        return set(dataset_dict.keys()).intersection(
            [DatasetSplitName.eval, DatasetSplitName.train]
        )

    def compute_on_dataset_split(self) -> List[ModuleResponse]:
        """Method that iterates over a dataset_split and call `compute`"""
        ds: Dataset = self.get_dataset_split()
        batch_size = self.config.batch_size
        num_rows = list(range(len(ds)))
        batches = [num_rows[i : i + batch_size] for i in range(0, len(ds), batch_size)]

        result = []
        for batch in tqdm(
            batches,
            desc=f"{self.task_name} on {self.dataset_split_name} set "
            f"for pipeline {self.mod_options.pipeline_index}",
        ):
            result += self.compute(ds.select(batch))
        return result

    def compute(self, batch: Dataset) -> List[ModuleResponse]:
        raise NotImplementedError

    def get_dataset_split(self, name: Optional[DatasetSplitName] = None) -> Dataset:
        """Get the specified dataset_split. For regular modules, not filtering or indexing is done.

        Args:
            name: Which dataset_split to get.

        Returns:
            The loaded dataset_split.

        Raises:
            ValueError: If no dataset_split name is provided or defined in the module.

        """
        if self.dataset_split_name is None and name is None:
            raise ValueError(
                "No dataset_split_name defined for this module. Specify what dataset_split to get."
            )
        dataset_split_name = name or self.dataset_split_name
        dm = self.get_dataset_split_manager(dataset_split_name)
        dataset_split: Dataset = dm.get_dataset_split(self._get_table_key())
        return dataset_split

    def get_dataset_split_manager(
        self, name: Optional[DatasetSplitName] = None
    ) -> DatasetSplitManager:
        """
        Get the DatasetSplitManager object that the module is working with.

        Args:
            name: Which dataset_split to get DatasetSplitManager for.

        Returns:
            DatasetSplitManager instance.

        Raises:
            If there is no dataset defined in the config. Or if the split is not defined.

        """
        dataset_split_name = name or self.dataset_split_name
        dataset_split_manager: DatasetSplitManager = (
            self.artifact_manager.get_dataset_split_manager(self.config, dataset_split_name)
        )
        return dataset_split_manager

    def get_model(self):
        """Load the model according to the config and module options.

        This will invoke the Python method supplied in the config.
        See our API Doc for more details. It can return a HF Pipeline, a Callable (which can be just
        a model or a custom pipeline) and a file prediction reader.

        Returns:
            Loaded model according to user spec.

        Raises:
            ValueError if no valid pipeline exists.
        """
        if self.worker != Worker.model:
            raise RuntimeError("This module cannot load the model. Modify self.worker.")
        _ = self.get_pipeline_definition()  # Validate current pipeline exists
        return self.artifact_manager.get_model(self.config, self.mod_options.pipeline_index)

    def _get_table_key(self) -> Optional[PredictionTableKey]:
        """Get the Table key for the current module.

        Returns:
            PredictionTableKey to use with DatasetSplitManager.
        """
        if self.mod_options.pipeline_index is None or not isinstance(
            self.config, ModelContractConfig
        ):
            return None
        return PredictionTableKey(
            temperature=self.get_pipeline_definition().temperature,
            threshold=self.get_threshold(),
            use_bma=self.mod_options.use_bma,
            pipeline_index=self.mod_options.pipeline_index,
        )

    def get_threshold(self) -> Optional[float]:
        # The default is None so we have to handle it this way.
        thresh = self.mod_options.threshold
        if thresh is None:
            thresh = self.get_pipeline_definition().threshold
        return thresh

    def get_pipeline_definition(self) -> PipelineDefinition:
        """Get the pipeline definition from the config according to the pipeline_index.

        Returns:
            Definition of the pipeline from the config.

        Raises:
            ValueError if no valid pipeline exists.
        """
        # TODO: Could use single dispatch instead?
        if not isinstance(self.config, ModelContractConfig):
            raise ValueError(
                "This module doesn't have access to pipelines as it doesn't use ModelContractConfig"
            )
        if self.config.pipelines is None:
            raise ValueError("No pipelines configured.")
        if self.mod_options.pipeline_index is None:
            raise ValueError(
                f"`pipeline_index` is None, expected one of {np.arange(len(self.config.pipelines))}"
            )
        pipelines = assert_not_none(self.config.pipelines)
        pipeline_index = assert_not_none(self.mod_options.pipeline_index)
        current_pipeline = pipelines[pipeline_index]
        return current_pipeline
