from sklearn.base import BaseEstimator, TransformerMixin


class ValueTypeCastTransformer(BaseEstimator, TransformerMixin):
    """
    Transformer to cast multiple columns to specified dtypes.

    Parameters
    ----------
    dtype_mapping : dict
        Dictionary mapping column names to target dtypes, e.g.,
        {'col1': 'float32', 'col2': 'category'}
    """

    def __init__(self, dtype_mapping: dict):
        self.dtype_mapping = dtype_mapping

    def fit(self, X, y=None):
        missing_cols = [col for col in self.dtype_mapping if col not in X.columns]
        if missing_cols:
            raise ValueError(f"Columns {missing_cols} are not in the DataFrame")
        return self

    def transform(self, X):
        X = X.copy()
        for col, target_dtype in self.dtype_mapping.items():
            X[col] = X[col].astype(target_dtype)
        return X

