from sklearn.base import BaseEstimator, TransformerMixin


class VectorizeCleanCols(BaseEstimator, TransformerMixin):
    def __init__(self):
        self.COLS_TO_DROP = ['customerID']
        self.COLS_TO_CONVERT_TO_FLOAT = ['tenure', 'MonthlyCharges', 'TotalCharges']

    def fit(self, df, y=None):
        self.YES_NO_COLS = self.get_yes_no_cols(df)
        return self

    def transform(self, df):
        df = df[sorted(list(df.columns))]
        df = self.drop_cols(df)
        df = self.convert_to_float_cols(df)
        df = self.vectorize_yes_no_cols(df)
        df.loc[:, 'gender'] = df['gender'].map({'Male': 1, 'Female': 0})
        return df

    def drop_cols(self, df):
        print('Dropping cols', self.COLS_TO_DROP)
        for col in self.COLS_TO_DROP:
            df.drop(col, axis=1, inplace=True)
        return df

    def convert_to_float_cols(self, df):
        def convert_str_to_float(x):
            try:
                return float(x)
            except:  # empty
                return 0.
        print('Converting cols to float', self.COLS_TO_CONVERT_TO_FLOAT)
        for col in self.COLS_TO_CONVERT_TO_FLOAT:
            df.loc[:, col] = df[col].apply(convert_str_to_float)
        return df

    def vectorize_yes_no_cols(self, df):
        print('Converting Yes/No cols to int', self.YES_NO_COLS)
        for col in self.YES_NO_COLS:
            df.loc[:, col] = df[col].map({'Yes': 1, 'No': 0})
        return df

    @staticmethod
    def get_yes_no_cols(df):
        yes_no_cols = []
        for col in df.columns:
            unique_val_set = set(df[col].unique())
            if ('Yes' in unique_val_set or 'No' in unique_val_set) and unique_val_set <= {'Yes', 'No'}:
                yes_no_cols.append(col)
        return yes_no_cols
