import logging
from typing import List, Optional

import pandas as pd
from allensdk.brain_observatory.behavior.behavior_project_cache.tables.ophys_mixin import (  # noqa: E501
    OphysMixin,
)
from allensdk.brain_observatory.behavior.behavior_project_cache.tables.project_table import (  # noqa: E501
    ProjectTable,
)
from allensdk.brain_observatory.behavior.utils.metadata_parsers import (  # noqa: E501
    parse_num_cortical_structures,
    parse_num_depths,
)
from allensdk.core.dataframe_utils import (
        enforce_df_column_order
)
from allensdk.brain_observatory.ophys.project_constants import VBO_METADATA_COLUMN_ORDER  # noqa: E501


class BehaviorOphysSessionsTable(ProjectTable, OphysMixin):
    """Class for storing and manipulating project-level data
    at the behavior-ophys session level"""

    def __init__(
        self,
        df: pd.DataFrame,
        suppress: Optional[List[str]] = None,
        index_column: str = "ophys_session_id",
    ):
        """
        Parameters
        ----------
        df
            The behavior-ophys session-level data
        suppress
            columns to drop from table
        index_column
            See description in BehaviorProjectCache.get_session_table
        """

        self._logger = logging.getLogger(self.__class__.__name__)
        self._index_column = index_column
        ProjectTable.__init__(self, df=df, suppress=suppress)
        OphysMixin.__init__(self)
        self._df = enforce_df_column_order(
            self._df,
            VBO_METADATA_COLUMN_ORDER
        )

    def postprocess_additional(self):
        # Add ophys specific information.
        project_code_col = (
            "project_code_ophys"
            if "project_code_ophys" in self._df.columns
            else "project_code"
        )
        self._df["num_targeted_structures"] = (
            self._df[project_code_col]
            .apply(parse_num_cortical_structures)
            .astype("Int64")
        )
        self._df["num_depths_per_area"] = (
            self._df[project_code_col].apply(parse_num_depths).astype("Int64")
        )
        # Possibly explode and reindex
        self.__explode()

    def __explode(self):
        if self._index_column == "ophys_session_id":
            pass
        elif self._index_column == "ophys_experiment_id":
            self._df = (
                self._df.reset_index()
                .explode("ophys_experiment_id")
                .set_index("ophys_experiment_id")
            )
        else:
            self._logger.warning(
                f"Invalid value for `by`, '{self._index_column}', passed to "
                f"BehaviorOphysSessionsCacheTable."
                " Valid choices for `by` are 'ophys_experiment_id' and "
                "'ophys_session_id'."
            )
