import warnings

import pandas as pd


class OphysMixin:
    _df: pd.DataFrame

    """A mixin class for ophys project data"""
    def __init__(self):
        self._merge_columns()

    def _merge_columns(self):
        """Some columns such as date of acquisition are stored in both
        behavior_sessions table as well as ophys_sessions table. If a field
        is in both, then it gets suffix _behavior or _ophys.
        We select the value in the ophys_sessions table and remove the
        duplicated columns"""
        columns = self._df.columns
        to_drop = []
        for column in columns:
            if column.endswith('_behavior'):
                column = column.replace('_behavior', '')
                if f'{column}_ophys' in self._df:
                    self._check_behavior_ophys_equal(column=column)
                    self._df[column] = self._merge_column_values(column=column)
                    to_drop += [f'{column}_behavior', f'{column}_ophys']
        self._df.drop(to_drop, axis=1, inplace=True)

    def _check_behavior_ophys_equal(self, column: str):
        """Checks that the value for behavior_session and ophys_session are
        equal. If not, issues a warning

        Parameters
        ----------
        column
            Column to check
        """
        mask = ~self._df[f'{column}_ophys'].isna()

        if not self._df[f'{column}_ophys'][mask].equals(
                self._df[f'{column}_behavior'][mask]):
            warnings.warn("BehaviorSession and OphysSession "
                          f"{column} do not agree. This is "
                          "likely due to issues with the data in "
                          f"LIMS.")

    def _merge_column_values(self, column: str) -> pd.Series:
        """Takes the non-null values from ophys and merges with behavior
        values

        Parameters
        ----------
        column
            Column to merge

        Returns
        -------
        pd.Series
            Merged Series

        """
        values = self._df[f'{column}_ophys']
        values.loc[values.isna()] = \
            self._df[f'{column}_behavior'][values.isna()]
        return values
