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


class Categorical(BaseFeature):

    def get_loadcols(self):
        cols = [
            "ProductName",
            "EngineVersion",
            "AppVersion",
            "AvSigVersion",
            "RtpStateBitfield",
            "DefaultBrowsersIdentifier",
            "AVProductStatesIdentifier",
            "AVProductsInstalled",
            "AVProductsEnabled",
            "CountryIdentifier",
            "CityIdentifier",
            "OrganizationIdentifier",
            "GeoNameIdentifier",
            "LocaleEnglishNameIdentifier",
            "Platform",
            "Processor",
            "OsVer",
            "OsBuild",
            "OsSuite",
            "OsPlatformSubRelease",
            "OsBuildLab",
            "SkuEdition",
            "PuaMode",
            "IeVerIdentifier",
            "SmartScreen",
            "UacLuaenable",
            "Census_MDC2FormFactor",
            "Census_DeviceFamily",
            "Census_OEMNameIdentifier",
            "Census_OEMModelIdentifier",
            "Census_ProcessorManufacturerIdentifier",
            "Census_ProcessorModelIdentifier",
            "Census_ProcessorClass",
            "Census_PrimaryDiskTypeName",
            "Census_ChassisTypeName",
            "Census_PowerPlatformRoleName",
            "Census_InternalBatteryType",
            "Census_OSVersion",
            "Census_OSArchitecture",
            "Census_OSBranch",
            "Census_OSBuildNumber",
            "Census_OSBuildRevision",
            "Census_OSEdition",
            "Census_OSSkuName",
            "Census_OSInstallTypeName",
            "Census_OSInstallLanguageIdentifier",
            "Census_OSUILocaleIdentifier",
            "Census_OSWUAutoUpdateOptionsName",
            "Census_GenuineStateName",
            "Census_ActivationChannel",
            "Census_FlightRing",
            "Census_ThresholdOptIn",
            "Census_FirmwareManufacturerIdentifier",
            "Census_FirmwareVersionIdentifier",
            "Wdft_RegionIdentifier"
        ]
        return cols

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

    def plain(self, df):
        return df

    def LabelEncoding(self, df):
        cols = df.columns.tolist()

        full_df = None
        for part in ["train", "validate", "test"]:
            if part == "test":
                part_path = self.dataset_path / "test.csv"
            else:
                part_path = self.dataset_path / self.valid_dir / "{}.csv".format(part)
            another_df = pd.read_csv(part_path,
                                     dtype=get_csv_dtype(),
                                     usecols=cols)
            full_df = another_df if full_df is None else pd.concat([full_df, another_df])
        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

    def CountEncoding(self, df):
        feature_col = df.columns.tolist()[1]

        value_counts = dd.read_csv(self.dataset_path / self.valid_dir / "train.csv",
                                   dtype=get_csv_dtype(),
                                   usecols=[feature_col]
                                   ).compute()[feature_col].value_counts()

        df[feature_col] = df[feature_col].map(value_counts.to_dict()).fillna(1)
        return df

    def RankEncoding(self, df):
        feature_col = df.columns.tolist()[1]

        value_counts = dd.read_csv(self.dataset_path / self.valid_dir / "train.csv",
                                   dtype=get_csv_dtype(),
                                   usecols=[feature_col]
                                   ).compute()[feature_col].value_counts()
        rank_counts = value_counts.rank(method="min", ascending=False)

        df[feature_col] = df[feature_col].map(rank_counts.to_dict()).fillna(len(rank_counts))
        return df

    def FactorizeEncoding(self, df):
        cols = df.columns.tolist()
        train_df = pd.read_csv(self.dataset_path / self.valid_dir / "train.csv",
                               dtype=get_csv_dtype(),
                               usecols=cols)
        train_df.sort_values("MachineIdentifier", inplace=True)

        labels, uniques = pd.factorize(train_df[cols[1]])
        df[cols[1]] = uniques.get_indexer(df[cols[1]])

        return df

    def FrequencyEncoding(self, df):
        feature_col = df.columns.tolist()[1]

        train_df = pd.read_csv(self.dataset_path / self.valid_dir / "train.csv",
                               dtype=get_csv_dtype(),
                               usecols=[feature_col])
        value_counts = train_df[feature_col].value_counts()
        value_freq = value_counts / len(train_df)

        df[feature_col] = df[feature_col].map(value_freq.to_dict()).fillna(1 / len(train_df))
        return df

    def TargetEncoding(self, df):
        cols = df.columns.tolist()
        train_df = pd.read_csv(self.dataset_path / self.valid_dir / "train.csv",
                               dtype=get_csv_dtype(),
                               usecols=[cols[1]] + ["HasDetections"])
        mean_df = train_df.groupby(cols[1]).mean()["HasDetections"]

        df[cols[1]] = df[cols[1]].map(mean_df.to_dict()).fillna(1 / len(train_df))
        return df

    def PopularEncoding(self, df):
        cols = df.columns.tolist()
        value_counts = df[cols[1]].value_counts()
        value_freq = value_counts / len(df)
        df[cols[1]] = df[cols[1]].map(value_freq.to_dict()).fillna(1 / len(df))
        return df[cols]
