from sklearn import tree
from sklearn.model_selection import train_test_split
from sklearn.ensemble import AdaBoostClassifier, RandomForestClassifier,BaggingClassifier


#Splitting the data
data2_np = np.array(data2)
train, test = train_test_split(data2_np, test_size=0.2, random_state=217)
X_train = train[:,0:-2]
Y_train = train[:,-1]
X_test = test[:,0:-2]
Y_test = test[:, -1]

#Function for doing predictions with the sklearn models
# 'classifier': The classifier to use e.g. tree.DecisionTreeClassifier, AdaBoostClassifier, RandomForestClassifier etc.
def classifier(classifier):
    clafi =classifier()
    clafi.fit(X_train,Y_train) 
    predictions = clafi.predict(X_test)
    return predictions
	

#Example of applying the function to AdaBoost
predictions = classifier(AdaBoostClassifier)


#Altered version of the evaluation function from monday
def evaluate_pred(pred_class, real_class):
    N = [[0,0], [0,0]]   # Make a list of lists (i.e. matrix) for counting successes/failures.
    for i in np.arange(len(real_class)):
        if (pred_class[i] == 0 and real_class[i] == 0) : N[0][0] += 1
        if (pred_class[i] == 0 and real_class[i] == 1) : N[0][1] += 1
        if (pred_class[i] == 1 and real_class[i] == 0) : N[1][0] += 1
        if (pred_class[i] == 1 and real_class[i] == 1) : N[1][1] += 1
    fracWrong = float(N[0][1]+N[1][0])/float(len(real_class))
    return N, fracWrong

#Evaluate the result
N,fracWrong = evaluate_pred(predictions, Y_test)
print("  First number is my estimate, second is the MC truth:")
print("  True-Negative (0,0)  = ", N[0][0])
print("  False-Negative (0,1) = ", N[0][1])
print("  False-Positive (1,0) = ", N[1][0])
print("  True-Positive (1,1)  = ", N[1][1])
print("    Fraction wrong = ( (0,1) + (1,0) ) / sum = ", fracWrong)