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


class SequentialImputer(BaseEstimator, TransformerMixin):
    """
    Stage-wise imputer.
    Applies each imputer sequentially and can add a suffix for filled columns.
    Returns a DataFrame with proper column names.
    """

    def __init__(self, stages):
        """
        stages : list of dicts
            [
                {
                    "model": SimpleImputer(...),
                    "column_mask": ["A", "B"],
                    "new_column_suffix": "_filled"  # optional
                },
                ...
            ]
        """
        self.stages = stages

        self.fitted_stages_ = None

    def fit(self, X, y=None):
        self.stages = self.aggregate_stages(self.stages)

        X_work = X.copy()
        self.fitted_stages_ = []

        for stage in self.stages:
            model = stage["model"]
            cols = stage.get("column_mask", X_work.columns)
            suffix = stage.get("new_column_suffix", "")

            model.fit(X_work[cols])
            X_work = self._apply_stage(X_work, model, cols, suffix)

            self.fitted_stages_.append((model, cols, suffix))

        return self

    @staticmethod
    def aggregate_stages(stages):
        normalized_stages = []
        for stage in stages:
            if isinstance(stage, dict):
                normalized_stages.append(stage)
            else:
                normalized_stages.append({"model": stage})
        return normalized_stages

    def transform(self, X):
        X_work = X.copy()

        for model, cols, suffix in self.fitted_stages_:
            X_work = self._apply_stage(X_work, model, cols, suffix)

        return X_work

    @staticmethod
    def _apply_stage(X, model, cols, suffix=""):
        transformed = model.transform(X[cols])
        if transformed.ndim == 1:
            transformed = transformed.reshape(-1, 1)

        transformed_feature_names = model.get_feature_names_out(cols)

        df_transformed = pd.DataFrame(transformed, columns=transformed_feature_names, index=X.index)

        df_transformed = SequentialImputer._drop_unchanged(df_transformed, X)
        if suffix:
            df_transformed.columns = df_transformed.columns + suffix
            X_out = pd.concat([X, df_transformed], axis=1)
        else:
            X_out = X.copy()
            for col in df_transformed:
                X_out[col] = df_transformed[col]

        return X_out

    @staticmethod
    def _drop_unchanged(df_transformed, X, prefix="missing_indicator_"):
        cols_to_drop = []

        for col in df_transformed.columns:
            if col in X.columns:
                if X[col].isna().any():
                    continue
                if df_transformed[col].equals(X[col]):
                    cols_to_drop.append(col)

                    indicator_col = f"{prefix}{col}"
                    if indicator_col in df_transformed.columns:
                        cols_to_drop.append(indicator_col)

        df_transformed = df_transformed.drop(columns=cols_to_drop)
        return df_transformed