"""
==================
Population Manager
==================

"""
from __future__ import annotations

from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, overload

import pandas as pd

import vivarium.framework.population.utilities as pop_utils
from vivarium.component import Component
from vivarium.framework.event import Event
from vivarium.framework.lifecycle import lifecycle_states
from vivarium.framework.population.exceptions import PopulationError
from vivarium.framework.population.population_view import PopulationView
from vivarium.framework.resource import Resource
from vivarium.manager import Manager

if TYPE_CHECKING:
    from vivarium.framework.engine import Builder
    from vivarium.types import ClockStepSize, ClockTime

from collections import defaultdict


@dataclass
class SimulantData:
    """Data to help components initialize simulants.

    Any time simulants are added to the simulation, each initializer is called
    with this structure containing information relevant to their initialization.

    """

    index: pd.Index[int]
    """The index representing the new simulants being added to the simulation."""
    user_data: dict[str, Any]
    """A dictionary of extra data passed in by the component creating the population."""
    creation_time: ClockTime
    """The time when the simulants enter the simulation."""
    creation_window: ClockStepSize
    """The span of time over which the simulants are created. Useful for, e.g., distributing 
    ages over the window."""


class PopulationManager(Manager):
    """Manages the population state table."""

    # TODO: Move the configuration for initial population creation to
    #  user components.
    CONFIGURATION_DEFAULTS = {
        "population": {
            "population_size": 100,
        },
    }

    @property
    def name(self) -> str:
        """The name of this component."""
        return "population_manager"

    @property
    def private_columns(self) -> pd.DataFrame:
        """The dataframe of all population private columns.

        Notes
        -----
        Critically, the private columns dataframe not only contains all private
        columns created for the simulation, but also serves as the simulant
        index for the entire population. Even if no private columns are created,
        this dataframe will exist and all simulants will be represented by its index.
        """
        if self._private_columns is None:
            raise PopulationError("Population has not been initialized.")
        return self._private_columns

    ############################
    # Normal Component Methods #
    ############################

    def __init__(self) -> None:
        self._private_columns: pd.DataFrame | None = None
        self._private_column_metadata: defaultdict[str, list[str]] = defaultdict(list)
        self._registered_initializers: list[Callable[[SimulantData], None]] = []
        self.creating_initial_population = False
        self.adding_simulants = False
        self._last_id = -1
        self.tracked_queries: list[str] = []
        self.pipeline_evaluation_depth: int = 0

    def setup(self, builder: Builder) -> None:
        """Registers the population manager with other vivarium systems."""
        super().setup(builder)
        self.logger = builder.logging.get_logger(self.name)
        self.clock = builder.time.clock()
        self.step_size = builder.time.step_size()
        self.resources = builder.resources
        self._add_constraint = builder.lifecycle.add_constraint
        self._get_attribute_pipelines = builder.value.get_attribute_pipelines()
        self._register_attribute_producer = builder.value.register_attribute_producer
        self._get_current_component_or_manager = (
            builder.components.get_current_component_or_manager
        )
        self.get_current_state = builder.lifecycle.current_state()

        builder.lifecycle.add_constraint(
            self.get_view,
            allow_during=[
                lifecycle_states.SETUP,
                lifecycle_states.POST_SETUP,
                lifecycle_states.POPULATION_CREATION,
                lifecycle_states.SIMULATION_END,
                lifecycle_states.REPORT,
            ],
        )
        builder.lifecycle.add_constraint(
            self.get_simulant_creator, allow_during=[lifecycle_states.SETUP]
        )
        builder.lifecycle.add_constraint(
            self.register_initializer, allow_during=[lifecycle_states.SETUP]
        )
        self._add_constraint(
            self.get_population,
            restrict_during=[
                lifecycle_states.SETUP,
                lifecycle_states.POST_SETUP,
            ],
        )

        builder.event.register_listener(lifecycle_states.POST_SETUP, self.on_post_setup)

    def on_post_setup(self, event: Event) -> None:
        # All pipelines are registered during setup and so exist at this point.
        self._attribute_pipelines = self._get_attribute_pipelines()

    def __repr__(self) -> str:
        return "PopulationManager()"

    ###########################
    # Builder API and helpers #
    ###########################

    def register_tracked_query(self, query: str) -> None:
        """Updates list of registered tracked queries with the provided query.

        Parameters
        ----------
        query
            The new query to add to the running list of tracked queries.

        Notes
        -----
        While we log a warning if the same query is registered multiple times,
        we make no attempt to de-duplicate functionally-equivalent queries that
        are syntactically different, e.g. "x > 5" and "5 < x". In such cases,
        duplicate queries will be applied which is not optimal but will not
        affect correctness.
        """
        if query in self.tracked_queries:
            self.logger.warning(
                f"The tracked query '{query}' has already been registered. "
                "Duplicate registrations are ignored."
            )
            return
        self.tracked_queries.append(query)

    def get_private_column_names(self, component_name: str) -> list[str]:
        """Gets the names of private columns created by a given component.

        Parameters
        ----------
        component_name
            The name of the component whose private column names are to be retrieved.

        Returns
        -------
            The list of private column names created by the specified component.
            If the component has not created any private columns, an empty list is returned.
        """
        return self._private_column_metadata[component_name]

    @overload
    def get_private_columns(
        self,
        component: Component | Manager,
        index: pd.Index[int] | None = None,
        columns: str = ...,
    ) -> pd.Series[Any]:
        ...

    @overload
    def get_private_columns(
        self,
        component: Component | Manager,
        index: pd.Index[int] | None = None,
        columns: list[str] | tuple[str, ...] = ...,
    ) -> pd.DataFrame:
        ...

    @overload
    def get_private_columns(
        self,
        component: Component | Manager,
        index: pd.Index[int] | None = None,
        columns: None = None,
    ) -> pd.Series[Any] | pd.DataFrame:
        ...

    def get_private_columns(
        self,
        component: Component | Manager,
        index: pd.Index[int] | None = None,
        columns: str | list[str] | tuple[str, ...] | None = None,
    ) -> pd.DataFrame | pd.Series[Any]:
        """Gets the private columns for a given component.

        While the ``private_columns`` property provides a dataframe of all private
        columns in population, this method returns only the private columns created
        by the specified component. If no component is specified, then no columns
        are returned.

        Parameters
        ----------
        component
            The component whose private columns are to be retrieved. If None,
            no columns are returned.
        index
            The index of simulants to include in the returned dataframe. If None,
            all simulants are included.
        columns
            The specific column(s) to include. If None, all columns created by the
            component are included.

        Raises
        ------
        PopulationError
            If ``columns`` are requested during initial population creation
            (when no columns yet exist) or if the provided ``component`` does not
            create one or more of them.

        Returns
        -------
            The private column(s) created by the specified component. Will return
            a Series if a single column is requested or a Dataframe otherwise.
        """

        if self.creating_initial_population:
            if columns:
                raise PopulationError(
                    "Cannot get private columns during initial population "
                    "creation when no columns yet exist."
                )
            returned_cols = []
            squeeze = False  # does not really matter (will return an empty df anyway)
        else:
            all_private_columns = self._private_column_metadata.get(component.name, [])
            if columns is None:
                returned_cols = all_private_columns
                squeeze = True
            else:
                if isinstance(columns, str):
                    columns = [columns]
                    squeeze = True
                else:
                    columns = list(columns)
                    squeeze = False
                missing_cols = set(columns).difference(set(all_private_columns))
                if missing_cols:
                    raise PopulationError(
                        f"Component {component.name} is requesting the following "
                        f"private columns to which it does not have access: {missing_cols}."
                    )
                returned_cols = columns
        private_columns = self.private_columns[returned_cols]
        if squeeze:
            private_columns = private_columns.squeeze(axis=1)
        return private_columns.loc[index] if index is not None else private_columns

    def get_population_index(self) -> pd.Index[int]:
        """Gets the index of the current population."""
        return self.private_columns.index

    def get_view(self, component: Component | None = None) -> PopulationView:
        """Gets a time-varying view of the population state table.

        The requested population view can be used to view the current state or
        to update the state with new values.

        Parameters
        ----------
        component
            The component requesting this view. If None, the view will provide
            read-only access.

        Returns
        -------
            A view of the requested private columns of the population state table.

        """
        view = self._get_view(component)
        self._add_constraint(
            view.get,
            restrict_during=[
                lifecycle_states.INITIALIZATION,
                lifecycle_states.SETUP,
                lifecycle_states.POST_SETUP,
            ],
        )
        self._add_constraint(
            view.update,
            restrict_during=[
                lifecycle_states.INITIALIZATION,
                lifecycle_states.SETUP,
                lifecycle_states.POST_SETUP,
                lifecycle_states.SIMULATION_END,
                lifecycle_states.REPORT,
            ],
        )
        return view

    def _get_view(self, component: Component | None) -> PopulationView:
        self._last_id += 1
        view = PopulationView(self, component, self._last_id)
        return view

    def get_simulant_creator(self) -> Callable[[int, dict[str, Any] | None], pd.Index[int]]:
        """Gets a function that can generate new simulants.

        The creator function takes the number of simulants to be created as its
        first argument and a population configuration dict that will be available
        to simulant initializers as its second argument. It generates the new rows
        in the population state table and then calls each initializer registered
        with the population system with a data object containing the state table
        index of the new simulants, the configuration info passed to the creator,
        the current simulation time, and the size of the next time step.

        Returns
        -------
           The simulant creator function.
        """
        return self._create_simulants

    def _create_simulants(
        self, count: int, population_configuration: dict[str, Any] | None = None
    ) -> pd.Index[int]:
        population_configuration = (
            population_configuration if population_configuration else {}
        )
        if self._private_columns is None:
            self.creating_initial_population = True
            self._private_columns = pd.DataFrame()

        new_index = range(len(self._private_columns) + count)
        new_population = self._private_columns.reindex(new_index)
        index = new_population.index.difference(self._private_columns.index)
        self._private_columns = new_population
        self.adding_simulants = True
        for initializer in self.resources.get_population_initializers():
            initializer(
                SimulantData(index, population_configuration, self.clock(), self.step_size())
            )
        self.creating_initial_population = False
        self.adding_simulants = False

        missing = {}
        for component, cols_created in self._private_column_metadata.items():
            missing_cols = [col for col in cols_created if col not in self._private_columns]
            if missing_cols:
                missing[component] = missing_cols
        if missing:
            raise PopulationError(
                "The following components registered initializers to create columns "
                f"that were not actually created: {missing}."
            )

        return index

    def register_initializer(
        self,
        initializer: Callable[[SimulantData], None],
        columns: str | Sequence[str] | None,
        required_resources: Sequence[str | Resource] = (),
    ) -> None:
        """Registers a component's initializers and any (private) columns created by them.

        This does three primary things:
        1. Registers each private column's corresponding attribute producer.
        2. Records metadata about which component created which private columns.
        3. Registers the initializer as a resource.

        A `columns` value of None indicates that no private columns are being registered.
        This is useful when a component or manager needs to register an initializer
        that does not create any private columns.

        Parameters
        ----------
        initializer
            A function that will be called to initialize the state of new simulants.
        columns
            The private columns that the given initializer provides the initial state
            information for.
        required_resources
            The resources that the initializer requires to run. Strings are interpreted
            as attributes.

        Raises
        ------
        PopulationError
            If this initializer has already been registered or if the columns being
            created by this initializer overlap with columns created by another initializer.
        """

        if initializer in self._registered_initializers:
            raise PopulationError(
                f"The initializer '{initializer.__qualname__}' has already been registered. "
                "Each initializer may only be registered once."
            )

        component = self._get_current_component_or_manager()

        if columns is None:
            columns = []
        elif isinstance(columns, str):
            columns = [columns]
        for column_name in columns:
            # Check for duplicate registration
            for component_name, columns_list in self._private_column_metadata.items():
                if column_name in columns_list:
                    raise PopulationError(
                        f"Component '{component.name}' is attempting to register "
                        f"private column '{column_name}' but it is already registered "
                        f"by component '{component_name}'."
                    )
            # Register each private column's attribute producer
            self._register_attribute_producer(
                column_name,
                source=[column_name],
                source_is_private_column=True,
            )

        # Register private column metadata
        self._private_column_metadata[component.name].extend(columns)

        # Track the initializer to prevent duplicate registration
        self._registered_initializers.append(initializer)

        # Register the initializer as a resource
        self.resources.add_private_columns(
            initializer=initializer,
            columns=columns,
            required_resources=required_resources,
        )

    ###############
    # Context API #
    ###############

    def get_all_attribute_names(self) -> list[str]:
        """Gets the names of all attributes in the population.

        Returns
        -------
            A list of all attribute names in the population.
        """
        return list(self._attribute_pipelines.keys())

    @overload
    def get_population(
        self,
        attributes: list[str] | tuple[str, ...] | Literal["all"],
        index: pd.Index[int] | None = None,
        query: str = "",
        squeeze: Literal[True] = True,
        skip_post_processor: Literal[False] = False,
    ) -> pd.Series[Any] | pd.DataFrame:
        ...

    @overload
    def get_population(
        self,
        attributes: list[str] | tuple[str, ...] | Literal["all"],
        index: pd.Index[int] | None = None,
        query: str = "",
        squeeze: Literal[False] = ...,
        skip_post_processor: Literal[False] = False,
    ) -> pd.DataFrame:
        ...

    @overload
    def get_population(
        self,
        attributes: list[str] | tuple[str, ...] | Literal["all"],
        index: pd.Index[int] | None = None,
        query: str = "",
        squeeze: Literal[True, False] = True,
        skip_post_processor: Literal[True] = ...,
    ) -> Any:
        ...

    def get_population(
        self,
        attributes: list[str] | tuple[str, ...] | Literal["all"],
        index: pd.Index[int] | None = None,
        query: str = "",
        squeeze: Literal[True, False] = True,
        skip_post_processor: Literal[True, False] = False,
    ) -> Any:
        """Provides a copy of the population state table.

        Parameters
        ----------
        attributes
            The attributes to include as the state table. If "all", all attributes are included.
        index
            The index of simulants to include in the returned population. If None,
            all simulants are included.
        query
            Additional conditions used to filter the index.
        squeeze
            Whether or not to attempt to squeeze a multi-level column into a single-level
            column and/or a single-column dataframe into a series.
        skip_post_processor
            Whether we should invoke the post-processor on the combined
            source and mutator output or return without post-processing.
            This is useful when the post-processor acts as some sort of final
            unit conversion (e.g. the rescale post processor).

        Notes
        -----
        If ``skip_post_processor`` is True, the returned data will not be squeezed
        regardless of the ``squeeze`` argument passed.

        Returns
        -------
            A copy of the population state table.

        Raises
        ------
        TypeError
            If ``attributes`` is not a list or tuple of strings or "all".
        PopulationError
            - If any of the requested attributes do not exist in the state table.
            - If a required column for querying is missing from the state table.
            - If the population has not yet been initialized.
        ValueError
            If multiple attributes are requested when ``skip_post_processor`` is True.
        """

        if self._private_columns is None:
            return pd.DataFrame()

        if isinstance(attributes, str) and attributes != "all":
            raise TypeError(
                f"Attributes must be a list of strings or 'all'; got '{attributes}'."
            )

        if attributes == "all":
            requested_attributes = self.get_all_attribute_names()
        else:
            attributes = list(attributes)
            # check for duplicate request
            if len(attributes) != len(set(attributes)):
                # deduplicate while preserving order
                requested_attributes = list(dict.fromkeys(attributes))
                self.logger.warning(
                    f"Duplicate attributes requested: {set(attributes) - set(requested_attributes)}\n"
                    "Only returning one instance of each of these duplicate requests."
                )
            else:
                requested_attributes = attributes

        non_existent_attributes = set(requested_attributes) - set(
            self._attribute_pipelines.keys()
        )
        if non_existent_attributes:
            raise PopulationError(
                f"Requested attribute(s) {non_existent_attributes} not in population state table. "
                "This is likely due to a failure to require some columns, randomness "
                "streams, or pipelines when registering a simulant initializer, an attribute "
                "producer, or an attribute modifier. NOTE: It is possible for a run to "
                "succeed even if resource requirements were not properly specified in "
                "the simulant initializers or pipeline creation/modification calls. This "
                "success depends on component initialization order which may change in "
                "different run settings."
            )

        idx = index if index is not None else self._private_columns.index

        # Filter the index based on the query
        columns_to_get = set(requested_attributes)
        if query:
            query_columns = pop_utils.extract_columns_from_query(query)
            # We can remove these query columns from requested columns (and will fetch later)
            columns_to_get = columns_to_get.difference(query_columns)
            missing_query_columns = query_columns.difference(set(self._attribute_pipelines))
            if missing_query_columns:
                raise PopulationError(
                    "Columns used for querying missing from population state table:\n"
                    f"Missing columns: {missing_query_columns}\n"
                    f"Query: {query}"
                )
            query_df = self._get_attributes(idx, list(query_columns))
            query_df = query_df.query(query)
            idx = query_df.index

        data = self._get_attributes(
            idx,
            requested_attributes if skip_post_processor else list(columns_to_get),
            skip_post_processor,
        )
        if skip_post_processor:
            # NOTE: This correctly returns the requested attribute even when it
            # overlaps with query columns because we pass `requested_attributes`
            # (not `columns_to_get`) above when `skip_post_processor` is True.
            return data

        # Add on any query columns that are actually requested to be returned
        requested_query_columns = (
            query_columns.intersection(set(requested_attributes)) if query else set()
        )
        if requested_query_columns:
            requested_query_df = query_df[list(requested_query_columns)]
            if isinstance(data.columns, pd.MultiIndex):
                # Make the query df multi-index to prevent converting columns from
                # multi-index to single index w/ tuples for column names
                requested_query_df.columns = pd.MultiIndex.from_product(
                    [requested_query_df.columns, [""]]
                )
            data = pd.concat([data, requested_query_df], axis=1)

        # Maintain column ordering
        data = data[requested_attributes]

        if squeeze:
            if (
                isinstance(data.columns, pd.MultiIndex)
                and len(set(data.columns.get_level_values(0))) == 1
            ):
                # If multi-index columns with a single outer level, drop the outer level
                data = data.droplevel(0, axis=1)
            if len(data.columns) == 1:
                # If single column df, squeeze to series
                data = data.squeeze(axis=1)

        return data

    def get_tracked_query(self) -> str:
        """Gets the combined tracked query for the population.

        Returns
        -------
            A query string combining all registered tracked queries with "and" operators.
        """
        return " and ".join(self.tracked_queries)

    @overload
    def _get_attributes(
        self,
        idx: pd.Index[int],
        requested_attributes: Sequence[str],
        skip_post_processor: Literal[False] = ...,
    ) -> pd.DataFrame:
        ...

    @overload
    def _get_attributes(
        self,
        idx: pd.Index[int],
        requested_attributes: Sequence[str],
        skip_post_processor: Literal[True] = ...,
    ) -> Any:
        ...

    def _get_attributes(
        self,
        idx: pd.Index[int],
        requested_attributes: Sequence[str],
        skip_post_processor: Literal[True, False] = False,
    ) -> Any:
        """Get the population for a given index and requested attributes.

        While evaluating attribute pipelines, we increment ``pipeline_evaluation_depth``
        so that nested calls to ``PopulationView.get`` (which may be
        triggered by pipeline sources or mutators) do not automatically re-apply
        tracked queries. The index passed to each pipeline has already been filtered
        appropriately by the enclosing ``get_population`` call.

        Note that only tracked queries are suppressed. Explicit ``query`` arguments
        passed by the pipeline source/mutator are supported.
        """

        self.pipeline_evaluation_depth += 1
        try:
            return self.__get_attributes(idx, requested_attributes, skip_post_processor)
        finally:
            self.pipeline_evaluation_depth -= 1

    def __get_attributes(
        self,
        idx: pd.Index[int],
        requested_attributes: Sequence[str],
        skip_post_processor: Literal[True, False] = False,
    ) -> Any:
        """Core implementation of ``_get_attributes``."""

        if skip_post_processor:
            if len(requested_attributes) != 1:
                raise ValueError(
                    "When skip_post_processor is True, a single attribute must "
                    f"be requested. You requested {requested_attributes}."
                )
            return self._attribute_pipelines[requested_attributes[0]](
                idx, skip_post_processor=skip_post_processor
            )

        attributes_list: list[pd.Series[Any] | pd.DataFrame] = []

        # batch simple attributes and directly leverage private column backing dataframe
        simple_attributes = [
            name
            for name, pipeline in self._attribute_pipelines.items()
            if name in requested_attributes and pipeline.is_simple
        ]
        if simple_attributes:
            if self._private_columns is None:
                raise PopulationError("Population has not been initialized.")
            attributes_list.append(self._private_columns.loc[idx, simple_attributes])

        # handle remaining non-simple attributes one by one
        remaining_attributes = [
            attribute
            for attribute in requested_attributes
            if attribute not in simple_attributes
        ]
        contains_column_multi_index = False
        for name in remaining_attributes:
            values = self._attribute_pipelines[name](idx)

            # Handle column names
            if isinstance(values, pd.Series):
                if values.name is not None and values.name != name:
                    self.logger.warning(
                        f"The '{name}' attribute pipeline returned a pd.Series with a "
                        f"different name '{values.name}'. For the column being added to the "
                        f"population state table, we will use '{name}'."
                    )
                values.name = name
            else:
                # Must be a dataframe. Coerce the columns to multi-index and set the
                # attribute name as the outer level.
                if isinstance(values.columns, pd.MultiIndex):
                    # FIXME [MIC-6645]
                    raise NotImplementedError(
                        f"The '{name}' attribute pipeline returned a DataFrame with multi-level "
                        f"columns (nlevels={values.columns.nlevels}). Multi-level columns in "
                        "attribute pipeline outputs are not supported."
                    )
                values.columns = pd.MultiIndex.from_product([[name], values.columns])
                contains_column_multi_index = True
            attributes_list.append(values)

        # Make sure all items of the list have consistent column levels
        if contains_column_multi_index:
            for i, item in enumerate(attributes_list):
                if isinstance(item, pd.Series):
                    item_df = item.to_frame()
                    item_df.columns = pd.MultiIndex.from_tuples([(item.name, "")])
                    attributes_list[i] = item_df
                if isinstance(item, pd.DataFrame) and item.columns.nlevels == 1:
                    item.columns = pd.MultiIndex.from_product([item.columns, [""]])
        df = (
            pd.concat(attributes_list, axis=1) if attributes_list else pd.DataFrame(index=idx)
        )

        return df

    def update(self, update: pd.DataFrame) -> None:
        self.private_columns[update.columns] = update
