"""Metadata utilities and services for evaluators and experiments.

This module provides comprehensive metadata generation functionality for the
Brisk package, including model evaluation metadata, dataset information,
and rerun configuration metadata. It serves as a centralized service for
tracking and organizing metadata throughout the machine learning pipeline.

The MetadataService provides methods to generate standardized metadata
for different types of evaluations, ensuring consistent tracking of
experiments, models, datasets, and configurations.

Classes
-------
MetadataService
    Main service class for generating and managing metadata

Examples
--------
>>> from brisk.services.metadata import MetadataService
>>> from brisk.configuration import AlgorithmCollection
>>> 
>>> # Create metadata service
>>> metadata_service = MetadataService("metadata")
>>> metadata_service.set_algorithm_config(algorithm_config)
>>> 
>>> # Generate model metadata
>>> model_metadata = metadata_service.get_model(
...     models=my_model,
...     method_name="evaluate_model",
...     is_test=False
... )
>>> 
>>> # Generate dataset metadata
>>> dataset_metadata = metadata_service.get_dataset(
...     method_name="analyze_dataset",
...     dataset_name="iris",
...     group_name="classification"
... )
"""

from typing import Dict, Any, Union, List
import datetime

from sklearn import base

from brisk.services import base as base_service
from brisk.configuration import algorithm_collection

class MetadataService(base_service.BaseService):
    """Metadata generation service for evaluators and experiments.
    
    This service provides comprehensive metadata generation functionality for
    the Brisk package, including model evaluation metadata, dataset information,
    and rerun configuration metadata. It ensures consistent tracking and
    organization of metadata throughout the machine learning pipeline.
    
    The service generates standardized metadata dictionaries that include
    timestamps, method names, and type-specific information for different
    evaluation contexts.
    
    Attributes
    ----------
    algorithm_config : Optional[algorithm_collection.AlgorithmCollection]
        The algorithm configuration used for model metadata generation
        
    Notes
    -----
    The service requires an algorithm configuration to be set before generating
    model metadata. This can be done using the `set_algorithm_config()` method.
    
    Examples
    --------
    >>> from brisk.services.metadata import MetadataService
    >>> from brisk.configuration import AlgorithmCollection
    >>> 
    >>> # Create metadata service
    >>> metadata_service = MetadataService("metadata")
    >>> metadata_service.set_algorithm_config(algorithm_config)
    >>> 
    >>> # Generate different types of metadata
    >>> model_meta = metadata_service.get_model(models, "evaluate", is_test=False)
    >>> dataset_meta = metadata_service.get_dataset("analyze", "iris", "classification")
    >>> rerun_meta = metadata_service.get_rerun("save_config")
    """
    def __init__(self, name: str) -> None:
        """Initialize the metadata service.
        
        This constructor sets up the metadata service with the specified name
        and initializes the algorithm configuration to None. The algorithm
        configuration must be set separately using `set_algorithm_config()`.
        
        Parameters
        ----------
        name : str
            The name identifier for this service
            
        Notes
        -----
        The algorithm configuration is required for generating model metadata
        but not for dataset or rerun metadata.
        """
        super().__init__(name)
        self.algorithm_config = None

    def _get_base(self, method_name: str) -> Dict[str, Any]:
        """Generate base metadata information common to all metadata types.
        
        This private method creates the foundational metadata that is included
        in all metadata dictionaries generated by the service. It provides
        consistent timestamp and method tracking across all metadata types.
        
        Parameters
        ----------
        method_name : str
            The name of the calling method that is generating the metadata

        Returns
        -------
        Dict[str, Any]
            Base metadata dictionary containing timestamp and method name
            
        Notes
        -----
        The timestamp is generated at the time of method call and formatted
        as "YYYY-MM-DD HH:MM:SS" for consistent logging and tracking.
        """
        return {
            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "method": method_name,
        }

    def get_model(
        self,
        models: Union[base.BaseEstimator, List[base.BaseEstimator]],
        method_name: str,
        is_test: bool = False
    ) -> Dict[str, Any]:
        """Generate metadata for a model evaluation.

        This method creates comprehensive metadata for model evaluations,
        including information about the models used, evaluation context,
        and timing information. It extracts algorithm wrapper information
        from the configured algorithm collection.

        Parameters
        ----------
        models : Union[base.BaseEstimator, List[base.BaseEstimator]]
            The model(s) to include in metadata. Can be a single model or
            a list of models
        method_name : str
            The name of the calling method that is performing the evaluation
        is_test : bool, default=False
            Whether the evaluation is being performed on test data

        Returns
        -------
        Dict[str, Any]
            Metadata dictionary containing:
            - timestamp: When the evaluation was performed
            - method: Name of the calling method
            - type: "model" (indicates this is model metadata)
            - models: Dictionary mapping wrapper names to display names
            - is_test: String representation of whether this is test data

        Raises
        ------
        AttributeError
            If algorithm_config is not set or models don't have wrapper_name
        KeyError
            If model wrapper names are not found in algorithm_config

        Examples
        --------
        >>> metadata_service = MetadataService("metadata")
        >>> metadata_service.set_algorithm_config(algorithm_config)
        >>> 
        >>> # Single model
        >>> model_meta = metadata_service.get_model(
        ...     models=my_model,
        ...     method_name="evaluate_model",
        ...     is_test=False
        ... )
        >>> 
        >>> # Multiple models
        >>> models_meta = metadata_service.get_model(
        ...     models=[model1, model2],
        ...     method_name="compare_models",
        ...     is_test=True
        ... )
        """
        metadata = self._get_base(method_name)
        metadata["type"] = "model"
        metadata["models"] = {}
        metadata["is_test"] = str(is_test)

        if not isinstance(models, list):
            models = [models]

        for model in models:
            wrapper = self.algorithm_config[model.wrapper_name]
            metadata["models"][wrapper.name] = wrapper.display_name

        return metadata

    def get_dataset(
        self,
        method_name: str,
        dataset_name: str,
        group_name: str
    ) -> Dict[str, Any]:
        """Generate metadata for a dataset evaluation.

        This method creates metadata for dataset-related evaluations,
        including information about the dataset being analyzed and
        the group context. This is typically used for dataset analysis
        and visualization operations.

        Parameters
        ----------
        method_name : str
            The name of the calling method that is performing the analysis
        dataset_name : str
            The name of the dataset being analyzed
        group_name : str
            The name of the group or experiment group the dataset belongs to

        Returns
        -------
        Dict[str, Any]
            Metadata dictionary containing:
            - timestamp: When the analysis was performed
            - method: Name of the calling method
            - type: "dataset" (indicates this is dataset metadata)
            - dataset: Name of the dataset
            - group: Name of the group

        Examples
        --------
        >>> metadata_service = MetadataService("metadata")
        >>> 
        >>> # Generate dataset metadata
        >>> dataset_meta = metadata_service.get_dataset(
        ...     method_name="analyze_dataset",
        ...     dataset_name="iris",
        ...     group_name="classification"
        ... )
        >>> 
        >>> # For visualization metadata
        >>> viz_meta = metadata_service.get_dataset(
        ...     method_name="create_plots",
        ...     dataset_name="housing",
        ...     group_name="regression"
        ... )
        """
        metadata = self._get_base(method_name)
        metadata["type"] = "dataset"
        metadata["dataset"] = dataset_name
        metadata["group"] = group_name
        return metadata

    def get_rerun(self, method_name: str) -> Dict[str, Any]:
        """Generate metadata for rerun configuration files.

        This method creates metadata for rerun configuration operations,
        which are used to track and manage experiment reruns and
        configuration persistence.

        Parameters
        ----------
        method_name : str
            The name of the calling method that is handling rerun configuration

        Returns
        -------
        Dict[str, Any]
            Metadata dictionary containing:
            - timestamp: When the rerun operation was performed
            - method: Name of the calling method
            - type: "rerun_config" (indicates this is rerun metadata)

        Examples
        --------
        >>> metadata_service = MetadataService("metadata")
        >>> 
        >>> # Generate rerun metadata
        >>> rerun_meta = metadata_service.get_rerun("save_rerun_config")
        >>> 
        >>> # For loading rerun configuration
        >>> load_meta = metadata_service.get_rerun("load_rerun_config")
        """
        metadata = self._get_base(method_name)
        metadata["type"] = "rerun_config"
        return metadata

    def set_algorithm_config(
        self,
        algorithm_config: algorithm_collection.AlgorithmCollection
    ) -> None:
        """Set the algorithm configuration for model metadata generation.

        This method sets the algorithm configuration that is used when
        generating model metadata. The configuration is required to extract
        algorithm wrapper information and display names for models.

        Parameters
        ----------
        algorithm_config : algorithm_collection.AlgorithmCollection
            The algorithm configuration containing wrapper information

        Notes
        -----
        This method must be called before using `get_model()` to generate
        model metadata. The algorithm configuration is not required for
        dataset or rerun metadata generation.

        Examples
        --------
        >>> from brisk.configuration import AlgorithmCollection
        >>> metadata_service = MetadataService("metadata")
        >>> 
        >>> # Set algorithm configuration
        >>> metadata_service.set_algorithm_config(algorithm_config)
        >>> 
        >>> # Now model metadata can be generated
        >>> model_meta = metadata_service.get_model(models, "evaluate")
        """
        self.algorithm_config = algorithm_config
