import sys
from class_vis2 import prettyPicture

import matplotlib.pyplot as plt
import numpy as np
import pylab as pl

# features_train, labels_train, features_test, labels_test = makeTerrainData()

from data2 import *

features_train = X
labels_train = y
features_test = X_test
labels_test = y_test


########################## DECISION TREE #################################


### your code goes here--now create 2 decision tree classifiers,
### one with min_samples_split=2 and one with min_samples_split=50
### compute the accuracies on the testing data and store
### the accuracy numbers to acc_min_samples_split_2 and
### acc_min_samples_split_50, respectively

from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score

clf = DecisionTreeClassifier(min_samples_split=2)
clf.fit(features_train, labels_train)
pred2 = clf.predict(features_test)
acc_min_samples_split_2 = accuracy_score(labels_test, pred2)

prettyPicture(clf, features_test, labels_test, "test2.2.png")

clf = DecisionTreeClassifier(min_samples_split=50)
clf.fit(features_train, labels_train)
pred50 = clf.predict(features_test)
acc_min_samples_split_50 = accuracy_score(labels_test, pred50)

prettyPicture(clf, features_test, labels_test, "test2.50.png")

def submitAccuracies():
  return {"acc_min_samples_split_2":round(acc_min_samples_split_2,3),
          "acc_min_samples_split_50":round(acc_min_samples_split_50,3)}