from pathlib import Path

import joblib
import numpy as np
import pandas as pd


dtypes = {
    'MachineIdentifier':                                    'object',
#    'ProductName':                                          'object',
    'EngineVersion':                                        'object',
    'AppVersion':                                           'object',
    'AvSigVersion':                                         'object',
#    'IsBeta':                                               'int8',
    'RtpStateBitfield':                                     'float16',
#    'IsSxsPassiveMode':                                     'int8',
#    'DefaultBrowsersIdentifier':                            'float32',
    'AVProductStatesIdentifier':                            'float32',
    'AVProductsInstalled':                                  'float16',
    'AVProductsEnabled':                                    'float16',
#    'HasTpm':                                               'int8',
    'CountryIdentifier':                                    'int16',
    'CityIdentifier':                                       'float32',
    'OrganizationIdentifier':                               'float16',
    'GeoNameIdentifier':                                    'float16',
    'LocaleEnglishNameIdentifier':                          'int16',
#    'Platform':                                             'object',
    'Processor':                                            'object',
#    'OsVer':                                                'object',
    'OsBuild':                                              'int16',
    'OsSuite':                                              'int16',
    'OsPlatformSubRelease':                                 'object',
    'OsBuildLab':                                           'object',
    'SkuEdition':                                           'object',
    'IsProtected':                                          'float16',
#    'AutoSampleOptIn':                                      'int8',
#    'PuaMode':                                              'object',
    'SMode':                                                'float16',
    'IeVerIdentifier':                                      'float16',
    'SmartScreen':                                          'object',
#    'Firewall':                                             'float16',
#    'UacLuaenable':                                         'float64',
    'Census_MDC2FormFactor':                                'object',
#    'Census_DeviceFamily':                                  'object',
    'Census_OEMNameIdentifier':                             'float32',
    'Census_OEMModelIdentifier':                            'float32',
    'Census_ProcessorCoreCount':                            'float16',
    'Census_ProcessorManufacturerIdentifier':               'float16',
    'Census_ProcessorModelIdentifier':                      'float32',
#    'Census_ProcessorClass':                                'object',
    'Census_PrimaryDiskTotalCapacity':                      'float64',
    'Census_PrimaryDiskTypeName':                           'object',
    'Census_SystemVolumeTotalCapacity':                     'float64',
#    'Census_HasOpticalDiskDrive':                           'int8',
    'Census_TotalPhysicalRAM':                              'float32',
    'Census_ChassisTypeName':                               'object',
    'Census_InternalPrimaryDiagonalDisplaySizeInInches':    'float32',
    'Census_InternalPrimaryDisplayResolutionHorizontal':    'float32',
    'Census_InternalPrimaryDisplayResolutionVertical':      'float32',
    'Census_PowerPlatformRoleName':                         'object',
#    'Census_InternalBatteryType':                           'object',
#    'Census_InternalBatteryNumberOfCharges':                'float64',
    'Census_OSVersion':                                     'object',
    'Census_OSArchitecture':                                'object',
    'Census_OSBranch':                                      'object',
    'Census_OSBuildNumber':                                 'int16',
    'Census_OSBuildRevision':                               'int32',
    'Census_OSEdition':                                     'object',
    'Census_OSSkuName':                                     'object',
    'Census_OSInstallTypeName':                             'object',
    'Census_OSInstallLanguageIdentifier':                   'float16',
    'Census_OSUILocaleIdentifier':                          'int16',
    'Census_OSWUAutoUpdateOptionsName':                     'object',
#    'Census_IsPortableOperatingSystem':                     'int8',
    'Census_GenuineStateName':                              'object',
    'Census_ActivationChannel':                             'object',
#    'Census_IsFlightingInternal':                           'float16',
#    'Census_IsFlightsDisabled':                             'float16',
#    'Census_FlightRing':                                    'object',
#    'Census_ThresholdOptIn':                                'float16',
    'Census_FirmwareManufacturerIdentifier':                'float16',
    'Census_FirmwareVersionIdentifier':                     'float32',
    'Census_IsSecureBootEnabled':                           'int8',
#    'Census_IsWIMBootEnabled':                              'float16',
#    'Census_IsVirtualDevice':                               'float16',
    'Census_IsTouchEnabled':                                'int8',
#    'Census_IsPenCapable':                                  'int8',
    'Census_IsAlwaysOnAlwaysConnectedCapable':              'float16',
    'Wdft_IsGamer':                                         'float16',
    'Wdft_RegionIdentifier':                                'float16',
    'HasDetections':                                        'float32'
}


def dump_train(data_dir):
    train_df = pd.read_csv(f'{data_dir}/train.csv.xz', dtype=dtypes,
                           low_memory=False)
    joblib.dump(train_df, f'{data_dir}/train.pkl', compress=3)
    #train_timestamps = pd.read_feather(f'{data_dir}/train_timestamps.feather')
    #joblib.dump(pd.concat([train_df, train_timestamps], axis=1),
    #            f'{data_dir}/train.pkl', compress=3)


def dump_test(data_dir):
    test_df = pd.read_csv(f'{data_dir}/test.csv.xz', dtype=dtypes,
                          low_memory=False)
    joblib.dump(test_df, f'{data_dir}/test.pkl', compress=3)
    #test_timestamps = pd.read_feather(f'{data_dir}/test_timestamps.feather')
    #joblib.dump(pd.concat([test_df, test_timestamps], axis=1),
    #            f'{data_dir}/test.pkl', compress=3)


def dump_target(data_dir, target_column):
    target = pd.read_csv(f'{data_dir}/train.csv.xz', dtype=dtypes,
                         usecols=[target_column], low_memory=False)
    joblib.dump(target, f'{data_dir}/target.pkl')


if __name__ == '__main__':
    data_dir = Path(__file__).parent.parent.joinpath('./data/input')

    dump_train(data_dir)
    dump_test(data_dir)
    dump_target(data_dir, 'HasDetections')

