import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin

class FeatureDropTransformer(BaseEstimator, TransformerMixin):
    """
    Drop one or more columns from a DataFrame within an sklearn pipeline.

    Optionally adds a new feature that contains, per row, the count of non-null values
    across the columns that are about to be dropped. This can be useful if the dropped
    columns encode sparse signals and you want to keep a compact summary.

    Parameters
    ----------
    features : list[str]
        Column names to drop.
    count_feature_name : str | None, optional
        If provided, a new column with this name is created that contains the per-row
        non-null count across `features`.
    """
    def __init__(self, features: list[str], count_feature_name = None):
        self.features = features
        self.count_feature_name = count_feature_name

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X = pd.DataFrame(X).copy()

        if not self.features:
            return X

        if self.count_feature_name:
            X[self.count_feature_name] = X[self.features].notna().sum(axis=1)

        X = X.drop(columns=self.features)
        return X

    def __sklearn_is_fitted__(self):
        return self.features is not None and len(self.features) > 0