# -*- coding: utf-8 -*-
  
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn import metrics
from sklearn.externals import joblib
from sklearn.model_selection import GridSearchCV
import os
import xlrd
from imblearn.over_sampling import RandomOverSampler
from sklearn.cross_validation import train_test_split

import warnings
from sklearn.model_selection import cross_val_score,cross_val_predict


def train():
    black_df = pd.read_csv('black_raw.csv')
    black_df['label'] = 'Black'
    white_df = pd.read_csv('white_raw.csv')
    white_df['label'] = 'White'
    data = pd.concat([black_df,white_df])
    X,y = loadTrain(data)
    return X,y

def loadTrain(data):
    data.dropna(inplace=True)
    extraKeys = ['clientPort', 'recList', 'sendList', 'clientIP','serverIP','label','serverPort']
    all_keys = data.columns
    feature_keys = list(set(all_keys)-set(extraKeys))    
    label = data['label']
    return data[feature_keys],label


x,y = train()
train_X, test_X, train_y, test_y = train_test_split(x,y,test_size=0.3,random_state=0)
rf = RandomForestClassifier()
rf.fit(train_X,train_y)
predicted = rf.predict(test_X)
print(metrics.accuracy_score(test_y, predicted))
print(metrics.confusion_matrix(test_y, predicted))
