#!/bin/python3

from fastai.tabular import *
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, StratifiedShuffleSplit
from sklearn.utils import shuffle

print('Imports complete')

df = pd.read_csv('../../malware_dataset/complete_binarized_dataset_cleaned.csv', index_col=0)
df.head()

procs = [FillMissing, Categorify, Normalize]
sss = StratifiedShuffleSplit(n_splits=10, test_size=0.2, random_state=1)
print(sss)

dep_var = 'Label_binarized'
cont_names = list(set(df.columns) - set([dep_var]))

for train_idx, test_idx in sss.split(df.index, df['Label_binarized']):
    data_fold = (TabularList.from_df(df, path='../../malware_dataset/complete_binarized_dataset_cleaned.csv', cont_names=cont_names, procs=procs)
                .split_by_idxs(train_idx, test_idx)
                .label_from_df(cols='Label_binarized')
                .databunch())
    model = tabular_learner(data_fold, layers=[50], metrics=[accuracy], callback_fns=ShowGraph, ps=0.5)
    model.fit_one_cycle(cyc_len=10)

