from sklearn.base import BaseEstimator, TransformerMixin


class FlagRegexTransformer(BaseEstimator, TransformerMixin):
    """
    A scikit-learn compatible transformer that creates a binary flag column
    based on whether values in a specified column match a given regex pattern.

    Parameters
    ----------
    target_column : str
        The name of the column to apply the regex to.
    pattern : str
        The regex pattern to match against the target column.
    new_column : str
        The name of the new column to store the flag (0/1).
    negate : bool, default=False
        If True, inverts the flag (e.g., for marking "pseudo" instead of "real").
    """

    def __init__(self, target_column, pattern, new_column, strip_string=False, negate=False):
        self.target_column = target_column
        self.pattern = pattern
        self.new_column = new_column
        self.strip_string = strip_string
        self.negate = negate

    def fit(self, X, y=None):
        """
        Fit method, does nothing but required for sklearn compatibility.
        """
        return self

    def transform(self, X):
        """
        Apply the regex match to the target column and create the new flag column.

        Parameters
        ----------
        X : pd.DataFrame
            Input dataframe.

        Returns
        -------
        X_transformed : pd.DataFrame
            Copy of the input dataframe with an additional flag column.
        """
        X = X.copy()
        flag = self._apply_pattern(X[self.target_column], self.pattern, strip=self.strip_string)

        if self.negate:
            flag = ~flag
        X[self.new_column] = flag
        return X

    @staticmethod
    def _apply_pattern(series, pattern, strip=True):
        if strip:
            series = series.str.strip()
        return series.str.contains(pattern, na=False)
