# %% - - visual studio code cell - -

import pandas as pd
from matplotlib import pyplot as plt
from sklearn import tree


# %% - - visual studio code cell - -

hep = pd.read_csv('hep_filled.csv')

x = hep.iloc[:, 1:]
y = hep['survival']

clf = tree.DecisionTreeClassifier(max_depth=4)
clf = clf.fit(x,y)


# %% - - visual studio code cell - -

plt.figure(figsize = (12,12))
tree.plot_tree(clf, fontsize=8, filled=True)
plt.savefig('tree.png', bbox_inches='tight')


# %% - - visual studio code cell - -

from sklearn.metrics import confusion_matrix

clf.score(x,y)
p = clf.predict(x)
confusion_matrix(y,p)


# %% - - visual studio code cell - -