import pandas as pd
import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn import metrics
from sklearn.metrics import accuracy_score
from sklearn import tree
import matplotlib as plt
from sklearn.tree import export_graphviz
import pydotplus
from IPython.display import Image

mushData=pd.read_csv('mushrooms.csv')
X=mushData.iloc[:,1:22]
y=mushData['class']
#need confusion matrix, training accurady, test accuracy, class tree,
#top 3 features, classify mushy
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size=0.30,random_state=1)
dt=DecisionTreeClassifier(max_depth=6,random_state=12)
print(X_train)
print(y_train)

dt.fit(X_train,y_train)
y_pred=dt.predict(X_test)
print(y_pred)

confusionmatrix=metrics.confusion_matrix(y_test,y_pred)
print(confusionmatrix)

acctest=accuracy_score(y_pred,y_test)
print("Test classification accuracy: ",acctest)
acctrain=accuracy_score(y_pred,y_train)
print("Train classification accuracy: ".acctrain)
print(y_pred)

dectree=tree.export_graphviz(dt)
graph=pydotplus.graph_from_dot_data(dectree)
graph.write_pdf("tree.pdf")