from sklearn.base import BaseEstimator, TransformerMixin

class ColumnSelectorTransformer(BaseEstimator, TransformerMixin):
    """Selects and orders columns for consistent pipeline"""

    # Updated column lists
    NUMERIC_COLS = [
        "Age",
        "DurationOfPitch",
        "MonthlyIncome",
        "NumberOfTrips",
        "NumberOfPersonVisiting",
        "NumberOfFollowups",
        "PreferredPropertyStar",
        "PitchSatisfactionScore",
        "NumberOfChildrenVisiting",
        "CityTier",
    ]

    CATEGORICAL_COLS = [
        "TypeofContact",
        "Occupation",
        "Gender",
        "MaritalStatus",
        "Passport",
        "OwnCar",
        "ProductPitched",
        "Designation",
    ]

    FEATURE_COLS = NUMERIC_COLS + CATEGORICAL_COLS

    def __init__(self):
        pass

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

    def transform(self, X):
        return X[self.FEATURE_COLS]


class CastCategoricalTransformer(BaseEstimator, TransformerMixin):
    """Handles categorical columns for LightGBM"""

    def __init__(self, categorical_cols):
        self.categorical_cols = categorical_cols

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

    def transform(self, X):
        X = X.copy()
        for col in self.categorical_cols:
            X[col] = X[col].astype("category")
        return X