# mypy: ignore-errors
"""
=========
Component
=========

A base Component class to be used to create components for use in ``vivarium``
simulations.

"""

import re
from abc import ABC
from importlib import import_module
from inspect import signature
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union

import pandas as pd
from layered_config_tree import ConfigurationError, LayeredConfigTree
from loguru._logger import Logger

from vivarium.framework.artifact import ArtifactException
from vivarium.framework.event import Event
from vivarium.framework.lookup import LookupTable
from vivarium.framework.population import PopulationError, PopulationView

if TYPE_CHECKING:
    from vivarium.framework.engine import Builder
    from vivarium.framework.population import SimulantData

DEFAULT_EVENT_PRIORITY = 5
"""The default priority at which events will be triggered."""


class Component(ABC):
    """The base class for all components used in a Vivarium simulation.

    A `Component` in a Vivarium simulation represents a distinct feature or
    aspect of the model. It encapsulates the logic and data needed for that
    feature. Components commonly interact with the rest of the simulation by
    creating and updating columns in the state table, registering pipelines,
    and registering modifiers on pipelines created by other components. Observer
    components might also register observations. All components within a
    simulation must have a unique name, which is generated by default from the
    component's class and the argument passed to its constructor.

    The `setup_component` is run by Vivarium during the setup phase and performs
    a series of operations to prepare the component for the simulation. These
    operations include setting the logger for the component, calling the
    component's custom `setup` method, setting the population view if the
    component needs one, and registering listeners for each lifecycle event if
    the component has defined a method to be triggered on that event.

    Subclasses of `Component` should override these properties as needed:

    - `sub_components`
    - `configuration_defaults`
    - `columns_created`
    - `columns_required`
    - `initialization_requirements`
    - `population_view_query`
    - `post_setup_priority`
    - `time_step_prepare_priority`
    - `time_step_priority`
    - `time_step_cleanup_priority`
    - `collect_metrics_priority`
    - `simulation_end_priority`

    Subclasses of `Component` should override these methods in order to have
    operations occur during the appropriate lifecycle phase of a simulation:

    - `setup`
    - `on_post_setup`
    - `on_initialize_simulants`
    - `on_time_step_prepare`
    - `on_time_step`
    - `on_time_step_cleanup`
    - `on_collect_metrics`
    - `on_simulation_end`

    """

    CONFIGURATION_DEFAULTS: Dict[str, Any] = {}
    """A dictionary containing the defaults for any configurations managed by this
    component. An empty dictionary indicates no managed configurations.
    """

    def __repr__(self):
        """Returns a string representation of the __init__ call made to create this
        object.

        The representation is built by retrieving the initialization parameters
        and their values. If a value is an instance of Component, its own
        __repr__() is called. The resulting string is stored in the _repr
        attribute and returned.

        IMPORTANT: this method must not be called within the `__init__`
        functions of this component or its subclasses or its value may not be
        initialized correctly.

        Returns
        -------
            A string representation of the __init__ call made to create this
            object.
        """
        if not self._repr:
            args = [
                f"{name}={value.__repr__() if isinstance(value, Component) else value}"
                for name, value in self.get_initialization_parameters().items()
            ]
            args = ", ".join(args)
            self._repr = f"{type(self).__name__}({args})"

        return self._repr

    def __str__(self):
        return self._repr

    ##############
    # Properties #
    ##############

    @property
    def name(self) -> str:
        """Returns the name of the component.

        By convention, these are in snake case with arguments of the `__init__`
        appended and separated by `.`.

        Names must be unique within a simulation.

        The name is created by first converting the name of the class to snake
        case. Then, the names of the initialization parameters are appended,
        separated by `.`. If a parameter is an instance of Component, its name
        property is used; otherwise, the string representation of the parameter
        is used. The resulting string is stored in the _name attribute and
        returned.

        IMPORTANT: this property must not be accessed within the `__init__`
        functions of this component or its subclasses or its value may not be
        initialized correctly.

        Returns
        -------
        str
            The name of the component.
        """
        if not self._name:
            base_name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", type(self).__name__)
            base_name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", base_name).lower()

            args = [
                f"'{value.name}'" if isinstance(value, Component) else str(value)
                for value in self.get_initialization_parameters().values()
            ]
            self._name = ".".join([base_name] + args)

        return self._name

    @property
    def population_view(self) -> PopulationView:
        """Provides the PopulationView for this component.

        Returns
        -------
        PopulationView
            The PopulationView for this component

        Raises
        ------
        PopulationError
            If the component does not have access to the state table.
        """
        if self._population_view is None:
            raise PopulationError(
                f"Component '{self.name}' does not have access to the state "
                "table. This is likely due to a failure to set columns_required "
                "or columns_created for this component."
            )
        return self._population_view

    @property
    def sub_components(self) -> List["Component"]:
        """Provide components managed by this component.

        Returns
        -------
        List[Component]
            The sub-components that are managed by this component.
        """
        return self._sub_components

    @property
    def configuration_defaults(self) -> Dict[str, Any]:
        """Provides a dictionary containing the defaults for any configurations
        managed by this component.

        These default values will be stored at the `component_configs` layer of the
        simulation's LayeredConfigTree.

        Returns
        -------
            A dictionary containing the defaults for any configurations managed by
            this component.
        """
        return self.CONFIGURATION_DEFAULTS

    @property
    def columns_created(self) -> List[str]:
        """Provides names of columns created by the component.

        Returns
        -------
            Names of the columns created by this component, or an empty list if
            none.
        """
        return []

    @property
    def columns_required(self) -> Optional[List[str]]:
        """Provides names of columns required by the component.

        Returns
        -------
            Names of required columns not created by this component. An empty
            list means all available columns are needed. `None` means no
            additional columns are necessary.
        """
        return None

    @property
    def initialization_requirements(self) -> Dict[str, List[str]]:
        """Provides the names of all values required by this component during
        simulant initialization.

        Returns
        -------
            A dictionary containing the additional requirements of this
            component during simulant initialization. An omitted key or an empty
            list for a key implies no requirements for that key during
            initialization.
        """
        return {
            "requires_columns": [],
            "requires_values": [],
            "requires_streams": [],
        }

    @property
    def population_view_query(self) -> Optional[str]:
        """Provides a query to use when filtering the component's `PopulationView`.

        Returns
        -------
            A pandas query string for filtering the component's `PopulationView`.
            Returns `None` if no filtering is required.
        """
        return None

    @property
    def post_setup_priority(self) -> int:
        """Provides the priority of this component's post_setup listener.

        Returns
        -------
            The priority of this component's post_setup listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    @property
    def time_step_prepare_priority(self) -> int:
        """Provides the priority of this component's time_step__prepare listener.

        Returns
        -------
            The priority of this component's time_step__prepare listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    @property
    def time_step_priority(self) -> int:
        """Provides the priority of this component's time_step listener.

        Returns
        -------
            The priority of this component's time_step listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    @property
    def time_step_cleanup_priority(self) -> int:
        """Provides the priority of this component's time_step__cleanup listener.

        Returns
        -------
            The priority of this component's time_step__cleanup listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    @property
    def collect_metrics_priority(self) -> int:
        """Provides the priority of this component's collect_metrics listener.

        Returns
        -------
            The priority of this component's collect_metrics listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    @property
    def simulation_end_priority(self) -> int:
        """Provides the priority of this component's simulation_end listener.

        Returns
        -------
            The priority of this component's simulation_end listener. This value
            can range from 0 to 9, inclusive.
        """
        return DEFAULT_EVENT_PRIORITY

    #####################
    # Lifecycle methods #
    #####################

    def __init__(self) -> None:
        """Initializes a new instance of the Component class.

        This method is the initializer for the Component class. It initializes
        logger of type Logger and population_view of type PopulationView to None.
        These attributes will be fully initialized in the setup_component method
        of this class.
        """
        self._repr: str = ""
        self._name: str = ""
        self._sub_components: List["Component"] = []
        self.logger: Optional[Logger] = None
        self.get_value_columns: Optional[
            Callable[[Union[str, pd.DataFrame]], List[str]]
        ] = None
        self.configuration: Optional[LayeredConfigTree] = None
        self._population_view: Optional[PopulationView] = None
        self.lookup_tables: Dict[str, LookupTable] = {}

    def setup_component(self, builder: "Builder") -> None:
        """Sets up the component for a Vivarium simulation.

        This method is run by Vivarium during the setup phase. It performs a series
        of operations to prepare the component for the simulation.

        It sets the logger for the component, sets up the component, sets the
        population view, and registers various listeners including post_setup,
        simulant_initializer, time_step_prepare, time_step, time_step_cleanup,
        collect_metrics, and simulation_end listeners.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        """
        self.logger = builder.logging.get_logger(self.name)
        self.get_value_columns = builder.data.value_columns()
        self.configuration = self.get_configuration(builder)
        self.build_all_lookup_tables(builder)
        self.setup(builder)
        self._set_population_view(builder)
        self._register_post_setup_listener(builder)
        self._register_simulant_initializer(builder)
        self._register_time_step_prepare_listener(builder)
        self._register_time_step_listener(builder)
        self._register_time_step_cleanup_listener(builder)
        self._register_collect_metrics_listener(builder)
        self._register_simulation_end_listener(builder)

    #######################
    # Methods to override #
    #######################

    def setup(self, builder: "Builder") -> None:
        """Defines custom actions this component needs to run during the setup
        lifecycle phase.

        This method is intended to be overridden by subclasses to perform any
        necessary setup operations specific to the component. By default, it
        does nothing.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        """
        pass

    def on_post_setup(self, event: Event) -> None:
        """Method that vivarium will run during the post_setup event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the post_setup event.

        Notes
        -----
        This method is not commonly used functionality.

        Parameters
        ----------
        event
            The event object associated with the post_setup event.
        """
        pass

    def on_initialize_simulants(self, pop_data: "SimulantData") -> None:
        """
        Method that vivarium will run during simulant initialization.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the simulant
        initialization phase.

        Parameters
        ----------
        pop_data : SimulantData
            The data associated with the simulants being initialized.

        Returns
        -------
        None
        """
        pass

    def on_time_step_prepare(self, event: Event) -> None:
        """Method that vivarium will run during the time_step__prepare event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the
        time_step__prepare event.

        Parameters
        ----------
        event
            The event object associated with the time_step__prepare event.
        """
        pass

    def on_time_step(self, event: Event) -> None:
        """Method that vivarium will run during the time_step event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the time_step event.

        Parameters
        ----------
        event
            The event object associated with the time_step event.
        """
        pass

    def on_time_step_cleanup(self, event: Event) -> None:
        """Method that vivarium will run during the time_step__cleanup event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the
        time_step__cleanup event.

        Parameters
        ----------
        event
            The event object associated with the time_step__cleanup event.
        """
        pass

    def on_collect_metrics(self, event: Event) -> None:
        """Method that vivarium will run during the collect_metrics event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the collect_metrics
        event.

        Parameters
        ----------
        event
            The event object associated with the collect_metrics event.
        """
        pass

    def on_simulation_end(self, event: Event) -> None:
        """Method that vivarium will run during the simulation_end event.

        This method is intended to be overridden by subclasses if there are
        operations they need to perform specifically during the simulation_end
        event.

        Parameters
        ----------
        event
            The event object associated with the simulation_end event.
        """
        pass

    ##################
    # Helper methods #
    ##################

    def get_initialization_parameters(self) -> Dict[str, Any]:
        """Retrieves the values of all parameters specified in the `__init__` that
        have an attribute with the same name.

        Notes
        -----
        This retrieves the value of the attribute at the time of calling,
        which is not guaranteed to be the same as the original value.

        Returns
        -------
            A dictionary where the keys are the names of the parameters used in
            the `__init__` method and the values are their current values.
        """
        return {
            parameter_name: getattr(self, parameter_name)
            for parameter_name in signature(self.__init__).parameters
            if hasattr(self, parameter_name)
        }

    def get_configuration(self, builder: "Builder") -> Optional[LayeredConfigTree]:
        """Retrieves the configuration for this component from the builder.

        This method retrieves the configuration for this component from the
        simulation's overall configuration. The configuration is retrieved using
        the name of the component as the key.

        Parameters
        ----------
        builder
            The simulation's builder object.

        Returns
        -------
            The configuration for this component, or `None` if the component has
            no configuration.
        """

        if self.name in builder.configuration:
            return builder.configuration[self.name]
        return None

    def build_all_lookup_tables(self, builder: "Builder") -> None:
        """Builds all lookup tables for this component.

        This method builds lookup tables for this component based on the data
        sources specified in the configuration. If no data sources are specified,
        no lookup tables are built.

        The created lookup tables are stored in the lookup_tables dictionary of
        the component, with the table name as the key.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        """
        if self.configuration and "data_sources" in self.configuration:
            for table_name in self.configuration.data_sources.keys():
                try:
                    self.lookup_tables[table_name] = self.build_lookup_table(
                        builder, self.configuration.data_sources[table_name]
                    )
                except ConfigurationError as e:
                    raise ConfigurationError(
                        f"Error building lookup table '{table_name}': {e}"
                    )

    def build_lookup_table(
        self,
        builder: "Builder",
        # todo: replace with LookupTableData
        data_source: Union[str, float, int, list, pd.DataFrame],
        value_columns: Optional[Iterable[str]] = None,
    ) -> LookupTable:
        """Builds a LookupTable from a data source.

        Uses `get_data` to parse the data source and retrieve the lookup table
        data. The LookupTable is built from the data source, with the value
        columns specified in the value_columns parameter. If value_columns is
        None and the data is a DataFrame, the ArtifactManager will determine
        the value columns.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        data_source
            The data source to build the LookupTable from.
        value_columns
            The columns to include in the LookupTable.

        Returns
        -------
            The LookupTable built from the data source.

        Raises
        ------
        layered_config_tree.exceptions.ConfigurationError
            If the data source is invalid.
        """
        data = self.get_data(builder, data_source)

        if isinstance(data, list):
            return builder.lookup.build_table(data, value_columns=list(value_columns))
        if isinstance(data, pd.DataFrame):
            all_columns = set(data.columns)
            if value_columns is None:
                value_columns = set(self.get_value_columns(data))
            else:
                value_columns = set(value_columns)

            potential_parameter_columns = [
                str(col).removesuffix("_start")
                for col in all_columns
                if str(col).endswith("_start")
            ]
            parameter_columns = set()
            bin_edge_columns = set()
            for column in potential_parameter_columns:
                if f"{column}_end" in all_columns:
                    parameter_columns.add(column)
                    bin_edge_columns.update([f"{column}_start", f"{column}_end"])

            key_columns = all_columns - value_columns - bin_edge_columns

            return builder.lookup.build_table(
                data=data,
                key_columns=list(key_columns),
                parameter_columns=list(parameter_columns),
                value_columns=list(value_columns),
            )

        return builder.lookup.build_table(data)

    def get_data(
        # TODO: replace with LookupTableData
        self,
        builder: "Builder",
        data_source: Union[str, float, pd.DataFrame],
    ) -> Union[float, pd.DataFrame]:
        """Retrieves data from a data source.

        If the data source is a float or a DataFrame, it is treated as the data
        itself. If the data source is a string, containing the substring '::',
        it is treated as a function to call to retrieve the data. The string to
        the left of '::' is the module to import, and the string to the right is
        the function to call. 'self' can be provided to the left of '::' to call
        a method on the component itself. If the data source is a string without
        the substring '::', it is treated as a key in the artifact.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        data_source
            The data source to retrieve data from.

        Returns
        -------
            The data retrieved from the data source.

        Raises
        ------
        layered_config_tree.exceptions.ConfigurationError
            If the data source is invalid.
        """
        if isinstance(data_source, (float, int, list, pd.DataFrame)):
            return data_source

        if "::" in data_source:
            module, method = data_source.split("::")
            try:
                if module == "self":
                    data_getter = getattr(self, method)
                else:
                    data_getter = getattr(import_module(module), method)
            except ModuleNotFoundError:
                raise ConfigurationError(f"Unable to find module '{module}'.")
            except AttributeError:
                module_string = (
                    f"component {self.name}." if module == "self" else f"module '{module}'."
                )
                raise ConfigurationError(
                    f"There is no method '{method}' for the {module_string}."
                )

            return data_getter(builder)

        try:
            return builder.data.load(data_source)
        except ArtifactException:
            raise ConfigurationError(f"Failed to find key '{data_source}' in artifact.")

    def _set_population_view(self, builder: "Builder") -> None:
        """Creates the PopulationView for this component if it needs access to
        the state table.

        The method determines the necessary columns for the PopulationView
        based on the columns required and created by this component. If no
        columns are required or created, no PopulationView is set.

        Parameters
        ----------
        builder
            The builder object used to set up the component.
        """
        if self.columns_required:
            # Get all columns created and required
            population_view_columns = self.columns_created + self.columns_required
        elif self.columns_required == []:
            # Empty list means population view needs all available columns
            population_view_columns = []
        elif self.columns_required is None and self.columns_created:
            # No additional columns required, so just get columns created
            population_view_columns = self.columns_created
        else:
            # no need for a population view if no columns created or required
            population_view_columns = None

        if population_view_columns is not None:
            self._population_view = builder.population.get_view(
                population_view_columns, self.population_view_query
            )

    def _register_post_setup_listener(self, builder: "Builder") -> None:
        """Registers a post_setup listener if this component has defined one.

        This method allows the component to respond to "post_setup" events if it
        has its own `on_post_setup` method. The listener will be registered with
        the component's post_setup priority, allowing control over the order of
        operations when multiple components are listening to the same event.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_post_setup != Component.on_post_setup:
            builder.event.register_listener(
                "post_setup",
                self.on_post_setup,
                self.post_setup_priority,
            )

    def _register_simulant_initializer(self, builder: "Builder") -> None:
        """Registers a simulant initializer if this component has defined one.

        This method allows the component to initialize simulants if it has its
        own `on_initialize_simulants` method. It registers this method with the
        builder's `PopulationManager``. It also specifies the columns that the
        component creates and any additional requirements for initialization.

        Parameters
        ----------
        builder
            The builder with which to register the initializer.
        """
        if type(self).on_initialize_simulants != Component.on_initialize_simulants:
            builder.population.initializes_simulants(
                self.on_initialize_simulants,
                creates_columns=self.columns_created,
                **self.initialization_requirements,
            )

    def _register_time_step_prepare_listener(self, builder: "Builder") -> None:
        """Registers a time_step_prepare listener if this component has defined one.

        This method allows the component to respond to "time_step_prepare" events
        if it has its own `on_time_step_prepare` method. The listener will be
        registered with the component's time_step_prepare priority.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_time_step_prepare != Component.on_time_step_prepare:
            builder.event.register_listener(
                "time_step__prepare",
                self.on_time_step_prepare,
                self.time_step_prepare_priority,
            )

    def _register_time_step_listener(self, builder: "Builder") -> None:
        """Registers a time_step listener if this component has defined one.

        This method allows the component to respond to "time_step" events
        if it has its own `on_time_step` method. The listener will be
        registered with the component's time_step priority.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_time_step != Component.on_time_step:
            builder.event.register_listener(
                "time_step",
                self.on_time_step,
                self.time_step_priority,
            )

    def _register_time_step_cleanup_listener(self, builder: "Builder") -> None:
        """Registers a time_step_cleanup listener if this component has defined one.

        This method allows the component to respond to "time_step_cleanup" events
        if it has its own `on_time_step_cleanup` method. The listener will be
        registered with the component's time_step_cleanup priority.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_time_step_cleanup != Component.on_time_step_cleanup:
            builder.event.register_listener(
                "time_step__cleanup",
                self.on_time_step_cleanup,
                self.time_step_cleanup_priority,
            )

    def _register_collect_metrics_listener(self, builder: "Builder") -> None:
        """Registers a collect_metrics listener if this component has defined one.

        This method allows the component to respond to "collect_metrics" events
        if it has its own `on_collect_metrics` method. The listener will be
        registered with the component's collect_metrics priority.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_collect_metrics != Component.on_collect_metrics:
            builder.event.register_listener(
                "collect_metrics",
                self.on_collect_metrics,
                self.collect_metrics_priority,
            )

    def _register_simulation_end_listener(self, builder: "Builder") -> None:
        """Registers a simulation_end listener if this component has defined one.

        This method allows the component to respond to "simulation_end" events
        if it has its own `on_simulation_end` method. The listener will be
        registered with the component's simulation_end priority.

        Parameters
        ----------
        builder
            The builder with which to register the listener.
        """
        if type(self).on_simulation_end != Component.on_simulation_end:
            builder.event.register_listener(
                "simulation_end",
                self.on_simulation_end,
                self.simulation_end_priority,
            )
