import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin


class ValueThresholdToNaNTransformer(BaseEstimator, TransformerMixin):
    """
    Replace values with NaN based on per-column threshold rules.

    For each configured column, values are set to NaN if they violate a threshold
    according to `mode`. The transformation can be applied in-place (overwrite the
    original column) or written to a new column via `new_column_suffix`.

    Parameters
    ----------
    thresholds : dict
        Mapping of column name -> threshold value.
    mode : str, default="min"
        - "min": values <= threshold are set to NaN
        - "max": values >= threshold are set to NaN
    drop_old_column : bool, default=True
        If True, overwrite the original column. If False, write to a new column.
    new_column_suffix : str, default="_cleaned"
        Suffix appended to the original name when `drop_old_column` is False.
    """
    def __init__(self, thresholds, mode="min", drop_old_column=True, new_column_suffix="_cleaned"):
        """
        thresholds: dict mapping feature names to threshold values.
                    Example: {"age": 18, "income": 1000}
        mode: "min" -> values <= threshold set to NaN
              "max" -> values >= threshold set to NaN
        new_column_suffix: string to append to feature names for new columns.
        drop_old_column: if True, original columns are removed after transformation.
        """
        self.thresholds = thresholds
        self.mode = mode
        self.new_column_suffix = new_column_suffix
        self.drop_old_column = drop_old_column

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

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

        for col, threshold in self.thresholds.items():
            if col not in X.columns:
                raise ValueError(f"Column '{col}' is not a feature inside this DataFrame")

            if self.mode == "min":
                transformed = X[col].where(X[col] > threshold, np.nan)
            elif self.mode == "max":
                transformed = X[col].where(X[col] < threshold, np.nan)
            else:
                raise ValueError("mode must be 'min' or 'max'")

            if self.drop_old_column:
                X[col] = transformed
            else:
                X[col + self.new_column_suffix] = transformed

        return X
