from features.base_feature import BaseFeature
import pandas as pd
import dask.dataframe as dd
from features.get_dtype import get_csv_dtype


class Defender(BaseFeature):

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

    def plain(self, df):
        return df

    def LabelEncoding(self, df):
        cols = df.columns.tolist()
        another_part = "test" if self.part == "train" else "train"

        part_df = dd.from_pandas(df, 2)
        another_df = dd.read_csv(self.input_path / "{}.csv".format(another_part),
                                 dtype=get_csv_dtype(),
                                 usecols=cols)
        full_df = dd.concat([part_df, another_df]).compute()

        full_df.sort_values("MachineIdentifier", inplace=True)
        labels, uniques = pd.factorize(full_df[cols[1]])
        full_df[cols[1]] = labels

        M_id = df["MachineIdentifier"].values.tolist()
        back_df = full_df[full_df["MachineIdentifier"].isin(M_id)]
        return back_df
