from features.base_feature import BaseFeature


class NewFeatures(BaseFeature):

    def get_loadcols(self):
        cols = [
            "Census_ProcessorCoreCount",
            "Census_ProcessorModelIdentifier",
            "Census_PrimaryDiskTotalCapacity",
            "Census_SystemVolumeTotalCapacity",
            "Census_TotalPhysicalRAM",
            "Census_PowerPlatformRoleName",
            "CountryIdentifier",
            "CityIdentifier",
            "GeoNameIdentifier",
            "LocaleEnglishNameIdentifier",
            "IeVerIdentifier",
            "EngineVersion",
            "AppVersion",
            "AvSigVersion",
            "OsBuildLab"
        ]
        return cols

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

    def processorCoreCount(self, df):
        """
        Census_ProcessorCoreCount
        """
        df["Census_ProcessorCoreCount"] = df["Census_ProcessorCoreCount"].fillna(4).astype(int)
        df.loc[df["Census_ProcessorCoreCount"] > 12, "Census_ProcessorCoreCount"] = 0
        return df[["MachineIdentifier", "Census_ProcessorCoreCount"]]

    def clfProcessorModel(self, df):
        """
        Census_ProcessorModelIdentifier
        """
        df["Census_ProcessorModelIdentifier"] = df["Census_ProcessorModelIdentifier"].fillna(2697).astype(int)
        df["Census_ProcessorModelIdentifier"].where((df["Census_ProcessorModelIdentifier"] < 2000) | (df["Census_ProcessorModelIdentifier"] > 3300),
                                                    2000, inplace=True)
        return df[["MachineIdentifier", "Census_ProcessorModelIdentifier"]]

    def clfPrimaryDiskTotalCapacity(self, df):
        """
        Census_PrimaryDiskTotalCapacity
        """
        df["Census_PrimaryDiskTotalCapacity"] = df["Census_PrimaryDiskTotalCapacity"].astype(float).fillna(4.769400e+05)
        df["Census_PrimaryDiskTotalCapacity"].where(df["Census_PrimaryDiskTotalCapacity"] < 400000, 400000, inplace=True)
        return df[["MachineIdentifier", "Census_PrimaryDiskTotalCapacity"]]

    def clfSystemVolumeTotalCapacity(self, df):
        """
        Census_SystemVolumeTotalCapacity
        """
        df["Census_SystemVolumeTotalCapacity"] = df["Census_SystemVolumeTotalCapacity"].fillna(28542).astype(int)
        df["Census_SystemVolumeTotalCapacity"] = (df["Census_SystemVolumeTotalCapacity"] > 10000)
        return df[["MachineIdentifier", "Census_SystemVolumeTotalCapacity"]]

    def clfTotalPhysicalRAM(self, df):
        """
        Census_TotalPhysicalRAM
        """
        df["Census_TotalPhysicalRAM"] = df["Census_TotalPhysicalRAM"].fillna(4096).astype(int)
        df["Census_TotalPhysicalRAM"].where(df["Census_TotalPhysicalRAM"] < 17500, 2500, inplace=True)
        return df[["MachineIdentifier", "Census_TotalPhysicalRAM"]]

    def platformRole_LabelEncoding(self, df):
        """
        Census_PowerPlatformRoleName
        """
        def reformat(se):
            se.replace({
                "AppliancePC": "Desktop",
                "EnterpriseServer": "Server",
                "PerformanceServer": "Server",
                "SOHOServer": "Server",
                "Slate": "Server",
                "Workstation": "Server",
                "UNKNOWN": "Others",
                "Unspecified": "Others"
            }, inplace=True)
            se.fillna("Others", inplace=True)
            return se

        df["Census_PowerPlatformRoleName"] = reformat(df["Census_PowerPlatformRoleName"])
        Label = {"Desktop": 1, "Mobile": 2, "Others": 3, "Server": 4}
        df["Census_PowerPlatformRoleName"] = df["Census_PowerPlatformRoleName"].map(Label).fillna(-1)

        return df[["MachineIdentifier", "Census_PowerPlatformRoleName"]]

    def clfCountry(self, df):
        """
        CountryIdentifier
        """
        df["CountryIdentifier"] = df["CountryIdentifier"].fillna(df["CountryIdentifier"].mode()[0]).astype(int)
        df = df.assign(clf=0)
        df.loc[df["CountryIdentifier"] <= 5, "clf"] = 1
        df.loc[(10 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 15), "clf"] = 1
        df.loc[(20 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 25), "clf"] = 2
        df.loc[(35 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 40), "clf"] = 1
        df.loc[(45 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 55), "clf"] = 1
        df.loc[(60 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 65), "clf"] = 1
        df.loc[(70 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 80), "clf"] = 1
        df.loc[(85 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 95), "clf"] = 2
        df.loc[(115 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 125), "clf"] = 1
        df.loc[(130 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 135), "clf"] = 2
        df.loc[(160 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 170), "clf"] = 1
        df.loc[(175 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 180), "clf"] = 1
        df.loc[(185 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 195), "clf"] = 2
        df.loc[(210 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 215), "clf"] = 2
        df.loc[(215 < df["CountryIdentifier"]) & (df["CountryIdentifier"] <= 220), "clf"] = 1
        return df[["MachineIdentifier", "clf"]]

    def clfCity(self, df):
        """
        CityIdentifier
        """
        df["CityIdentifier"] = df["CityIdentifier"].fillna(df["CityIdentifier"].mode()[0]).astype(int)
        df = df.assign(clf=0)
        df.loc[(8000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 9000), "clf"] = 1
        df.loc[(10000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 11000), "clf"] = 1
        df.loc[(16000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 17000), "clf"] = 1
        df.loc[(22000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 23000), "clf"] = 1
        df.loc[(37000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 38000), "clf"] = 1
        df.loc[(52000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 53000), "clf"] = 1
        df.loc[(58000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 59000), "clf"] = 1
        df.loc[(61000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 62000), "clf"] = 1
        df.loc[(70000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 71000), "clf"] = 1
        df.loc[(76000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 77000), "clf"] = 1
        df.loc[(87000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 88000), "clf"] = 1
        df.loc[(90000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 91000), "clf"] = 1
        df.loc[(111000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 112000), "clf"] = 1
        df.loc[(129000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 130000), "clf"] = 1
        df.loc[(137000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 138000), "clf"] = 1
        df.loc[(142000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 143000), "clf"] = 1
        df.loc[(145000 < df["CityIdentifier"]) & (df["CityIdentifier"] <= 146000), "clf"] = 1
        return df[["MachineIdentifier", "clf"]]

    def clfGeoName(self, df):
        """
        GeoNameIdentifier
        """
        df["GeoNameIdentifier"] = df["GeoNameIdentifier"].fillna(df["GeoNameIdentifier"].mode()[0]).astype(int)
        df = df.assign(clf=0)
        df.loc[(10 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 20), "clf"] = 1
        df.loc[(40 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 50), "clf"] = 1
        df.loc[(50 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 60), "clf"] = 2
        df.loc[(70 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 80), "clf"] = 2
        df.loc[(80 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 100), "clf"] = 1
        df.loc[(110 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 120), "clf"] = 2
        df.loc[(180 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 190), "clf"] = 1
        df.loc[(280 < df["GeoNameIdentifier"]) & (df["GeoNameIdentifier"] <= 290), "clf"] = 2
        return df[["MachineIdentifier", "clf"]]

    def clfLocaleEnglishName(self, df):
        """
        LocaleEnglishNameIdentifier
        """
        df["LocaleEnglishNameIdentifier"] = df["LocaleEnglishNameIdentifier"].fillna(df["LocaleEnglishNameIdentifier"].mode()[0]).astype(int)
        df = df.assign(clf=0)
        df.loc[(80 < df["LocaleEnglishNameIdentifier"]) & (df["LocaleEnglishNameIdentifier"] <= 100), "clf"] = 1
        df.loc[(120 < df["LocaleEnglishNameIdentifier"]) & (df["LocaleEnglishNameIdentifier"] <= 160), "clf"] = 1
        df.loc[280 < df["LocaleEnglishNameIdentifier"], "clf"] = 1
        return df[["MachineIdentifier", "clf"]]

    def clfIeVer(self, df):
        """
        IeVerIdentifier
        """
        df["IeVerIdentifier"] = df["IeVerIdentifier"].fillna(df["IeVerIdentifier"].mode()[0]).astype(int)
        df = df.assign(clf=0)
        df.loc[(0 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 20), "clf"] = 1
        df.loc[(30 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 50), "clf"] = 1
        df.loc[(70 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 110), "clf"] = 1
        df.loc[(120 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 130), "clf"] = 1
        df.loc[(140 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 230), "clf"] = 1
        df.loc[(290 < df["IeVerIdentifier"]) & (df["IeVerIdentifier"] <= 330), "clf"] = 1
        df.loc[390 < df["IeVerIdentifier"], "clf"] = 1
        return df[["MachineIdentifier", "clf"]]

    def clfCountRevision(self, df):
        """
        OsBuildLab
        https://github.com/icebee16/Malware/blob/master/notebook/eda/Census_OSBuildRevisionEDA.ipynb
        """
        ratio = 0.0028  # 25000 / 8921483

        df["OsBuildLab"] = df["OsBuildLab"].fillna("17134.1.amd64fre.rs4_release.180410-1804").str.replace(r"*", r".")
        df["OsBuildLab_Revision"] = df["OsBuildLab"].str.split(".", expand=True)[1].astype(int)

        vc = df["OsBuildLab_Revision"].value_counts()
        vc_list = vc[vc > ratio * len(df)].index.tolist()

        df = df.assign(cls=0)
        df.loc[df["OsBuildLab_Revision"].isin(vc_list) & (df["OsBuildLab_Revision"] < 650), "cls"] = 1
        df.loc[df["OsBuildLab_Revision"].isin(vc_list) & (df["OsBuildLab_Revision"] >= 650), "cls"] = 2

        return df[["MachineIdentifier", "cls"]]

    def popularEngine2(self, df):
        """
        EngineVersion
        """
        ratio = 0.1

        df["EngineVersion"] = df["EngineVersion"].str.split(".", expand=True)[2]

        vc = df["EngineVersion"].value_counts()
        vc_list = vc[vc > ratio * len(df)].index.tolist()

        df = df.assign(cls=0)
        df.loc[df["EngineVersion"].isin(vc_list), "cls"] = 1

        return df[["MachineIdentifier", "cls"]]

    def App1_is_18(self, df):
        """
        AppVersion
        """
        df["AppVersion"] = df["AppVersion"].str.split(".", expand=True)[1]
        df["AppVersion"] = (df["AppVersion"] == "18")
        return df[["MachineIdentifier", "AppVersion"]]

    def clfApp2(self, df):
        """
        AppVersion
        """
        df["AppVersion"] = df["AppVersion"].str.split(".", expand=True)[2]
        df["AppVersion"] = (df["AppVersion"].astype(float) > 10000)
        return df[["MachineIdentifier", "AppVersion"]]

    def popularAvSig1(self, df):
        """
        AvSigVersion
        """
        ratio = 0.1

        df["AvSigVersion"] = df["AvSigVersion"].str.replace(r"1.2&#x17;3.1144.0", r"1.273.1144.0")
        df["AvSigVersion"] = df["AvSigVersion"].str.split(".", expand=True)[1]

        vc = df["AvSigVersion"].value_counts()
        vc_list = vc[vc > ratio * len(df)].index.tolist()

        df = df.assign(cls=0)
        df.loc[df["AvSigVersion"].isin(vc_list), "cls"] = 1

        return df[["MachineIdentifier", "cls"]]

    def Eng2_mul_App1(self, df):
        """
        AvSigVersion
        EngineVersion
        """
        df = self.load_part()
        ratio = 0.1

        df["EngineVersion"] = df["EngineVersion"].str.split(".", expand=True)[2]
        df["AppVersion"] = df["AppVersion"].str.split(".", expand=True)[1]
        df["Eng2_mul_App1"] = df["EngineVersion"].astype(int) * df["AppVersion"].astype(int)

        vc = df["Eng2_mul_App1"].value_counts()
        vc_list = vc[vc > ratio * len(df)].index.tolist()

        df = df.assign(cls=0)
        df.loc[df["Eng2_mul_App1"].isin(vc_list), "cls"] = 1

        return df[["MachineIdentifier", "cls"]]
