from features.base_feature import BaseFeature
import pandas as pd
from features.get_dtype import get_csv_dtype


class OneHotEncoding(BaseFeature):

    def get_loadcols(self):
        cols = [
            "AVProductStatesIdentifier",
            "Processor",
            "SmartScreen",
            "Census_PowerPlatformRoleName"
        ]
        return cols

    def extract(self):
        return super().extract()

    def OneHotEncoding(self, df, category=None):
        feature = df.columns.tolist()[1]
        if category is None:
            df[feature] = ~(df[feature] == df[feature])
        else:
            df[feature] = (df[feature] == category)
        return df
