import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin


class CategoriсalFeaturesTransformer(BaseEstimator, TransformerMixin ):
    def __init__( self, features, toDtype = 'category'):
        self._numeric_features = features
        self._dtype = toDtype
    def fit( self, X, y = None ):
        return self
    def transform( self, X, y = None ):
        all_columns = set(X.columns)
        categoriсal_features  = list(all_columns.difference(self._numeric_features | set(['HasDetections'])))
        if (self._dtype == 'category'):
            X[categoriсal_features] = X[categoriсal_features].astype(self._dtype)
        else:
            for col in categoriсal_features:
                X[col] = pd.to_numeric(X[col])
        return X

class ColumnsDropper(BaseEstimator, TransformerMixin ):
    def __init__( self, columns_to_drop):
        self._columns_to_drop = columns_to_drop     
    def fit( self, X, y = None ):
        return self 
    def transform( self, X, y = None ):
        return X.drop(self._columns_to_drop & set(X.columns), axis = 1)



class ConstantImputer(BaseEstimator, TransformerMixin ):
    def __init__( self, constant):
        self._constant = constant     
    def fit( self, X, y = None ):
        return self 
    def transform( self, X, y = None ):
        for col in X.columns:
            if X[col].dtype.name == 'category':
                if self._constant not in X[col].cat.categories:
                    X[col] = X[col].cat.add_categories(str(self._constant))
                    X[col].fillna(str(self._constant), inplace = True)
            else:
            	X[col].fillna(self._constant, inplace = True)
        return X