from sklearn.base import BaseEstimator, TransformerMixin


class ValueRemapTransformer(BaseEstimator, TransformerMixin):
    """
    Remap values in a single DataFrame column using a mapping dictionary.

    Values not present in the mapping are left unchanged. The transformer can either
    overwrite the source column or write into a new column and optionally drop the
    original column.

    Parameters
    ----------
    column : str
        Name of the column to remap.
    mapping : dict
        Mapping of old_value -> new_value.
    new_column : str | None, optional
        Output column name. If None, overwrites `column`.
    drop_old_column : bool, default=True
        Only applies when `new_column` is not None. If True, drop the original column.
    """
    def __init__(self, column, mapping, new_column=None, drop_old_column=True):
        """
        Transformer to remap values of a DataFrame column according to a dictionary.
        Values not present in the mapping are kept unchanged.

        Parameters
        ----------
        column : str
            Name of the column to transform.
        mapping : dict
            Dictionary of old_value -> new_value.
        new_column : str or None
            Name of the new column. If None, overwrite original column.
        drop_old_column : bool
            Whether to drop the old column if new_column is set.
        """
        self.column = column
        self.mapping = mapping
        self.new_column = new_column
        self.drop_old_column = drop_old_column

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

    def transform(self, X):
        if self.column not in X.columns:
            raise ValueError(f"Column '{self.column}' is not present in the DataFrame")

        X = X.copy()

        target_col = self.new_column if self.new_column else self.column

        X[target_col] = X[self.column].replace(self.mapping)

        if self.new_column and self.drop_old_column:
            X = X.drop(columns=[self.column])

        return X
