import numpy as np

from features import Feature


############################################################
# Defender state
############################################################

class EngineVersion(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['EngineVersion']
        self.test[column] = test_df['EngineVersion']


class AppVersion(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['AppVersion']
        self.test[column] = test_df['AppVersion']


class AvSigVersion(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['AvSigVersion'].replace('1.2&#x17;3.1144.0', '1.273.1144.0')
        self.test[column] = test_df['AvSigVersion']


############################################################
# Antivirus
############################################################

class AVProductStatesId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['AVProductStatesIdentifier']
        self.test[column] = test_df['AVProductStatesIdentifier']


class IsAVProductStatesId53447(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return df['AVProductStatesIdentifier'] == 53447


class AVProductsInstalled(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['AVProductsInstalled']
        self.test[column] = test_df['AVProductsInstalled']


class AVProductsEnabled(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['AVProductsEnabled']
        self.test[column] = test_df['AVProductsEnabled']


############################################################
# Geolocation
############################################################

class CountryId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['CountryIdentifier']
        self.test[column] = test_df['CountryIdentifier']


class CityId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['CityIdentifier']
        self.test[column] = test_df['CityIdentifier']


class OrganizationId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['OrganizationIdentifier']
        self.test[column] = test_df['OrganizationIdentifier']


class GeoNameId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['GeoNameIdentifier']
        self.test[column] = test_df['GeoNameIdentifier']


class LocaleEnglishNameId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['LocaleEnglishNameIdentifier']
        self.test[column] = test_df['LocaleEnglishNameIdentifier']


############################################################
# OS
############################################################

# NOTE: almost the same as Census_OSBuildNumber
# class OsBuild(Feature):


class OsSuite(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['OsSuite']
        self.test[column] = test_df['OsSuite']


# NOTE: almost the same as Census_OSBranch
#class OsPlatformSubRelease(Feature):


class OsBuildLab(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['OsBuildLab']
        self.test[column] = test_df['OsBuildLab']


class OsBuildLabSplit0(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df)
        self.test[column] = self._compute(test_df)

    def _compute(self, df):
        return df['OsBuildLab'].fillna('0.0.0.0.0-0').apply(lambda x: str(x).split('.')[0])


class OsBuildLabSplit1(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df)
        self.test[column] = self._compute(test_df)

    def _compute(self, df):
        return df['OsBuildLab'].fillna('0.0.0.0.0-0').apply(lambda x: str(x).split('.')[1])


class SkuEdition(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['SkuEdition']
        self.test[column] = test_df['SkuEdition']


############################################################
# Hardware
############################################################

# NOTE: exactly the same as Census_OSArchitecture
# class Processor(Feature):


############################################################
# misc.
############################################################

class RtpStateBitfield(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['RtpStateBitfield']
        self.test[column] = test_df['RtpStateBitfield']


class IsProtected(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['IsProtected'].fillna(1).astype('int8')
        self.test[column] = test_df['IsProtected'].fillna(1).astype('int8')


class SMode(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['SMode'].fillna(0).astype('int8')
        self.test[column] = test_df['SMode'].fillna(0).astype('int8')


class IeVerId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['IeVerIdentifier']
        self.test[column] = test_df['IeVerIdentifier']


def normalize_smartscreen(df):
    trans_dict = {
        'BLOCK': 'Block',
        'Deny': 'Off',
        'off': 'Off',
        'OFF': 'Off',
        'of': 'Off',
        'on': 'On',
        'ON': 'On',
        'requireadmin': 'RequireAdmin',
        'requireAdmin': 'RequireAdmin',
        'RequiredAdmin': 'RequireAdmin',
        'Promt': 'Prompt',
        'Promprt': 'Prompt',
        'prompt': 'Prompt',
        'warn': 'Warn',
        '0': 'NoExist',
        '00000000': 'NoExist',
        '&#x01;': 'NoExist',
        '&#x02;': 'NoExist',
        '&#x03;': 'NoExist',
        np.nan: 'NoExist'
    }
    return df.replace({'SmartScreen': trans_dict})


class SmartScreen(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df)
        self.test[column] = self._compute(test_df)

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen']


class IsSmartScreenExistsNotSet(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'ExistsNotSet'


class IsSmartScreenRequireAdmin(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'RequireAdmin'


class IsSmartScreenOff(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'Off'


class IsSmartScreenOn(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'On'


class IsSmartScreenWarn(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'Warn'


class IsSmartScreenBlock(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'Block'


class IsSmartScreenPrompt(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'Prompt'


class IsSmartScreenNoExist(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = self._compute(train_df).astype('int8')
        self.test[column] = self._compute(test_df).astype('int8')

    def _compute(self, df):
        return normalize_smartscreen(df)['SmartScreen'] == 'NoExist'


############################################################
# wdft
############################################################

class WdftIsGamer(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['Wdft_IsGamer'].fillna(0).astype('int8')
        self.test[column] = test_df['Wdft_IsGamer'].fillna(0).astype('int8')


class WdftRegionId(Feature):
    def create_features_from_df(self, train_df, test_df):
        column = self.column
        self.train[column] = train_df['Wdft_RegionIdentifier']
        self.test[column] = test_df['Wdft_RegionIdentifier']

