from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

X_train = sfix.input_tensor_via(0, X_train)
X_test = sfix.input_tensor_via(0, X_test)

y_train = sint.input_tensor_via(0, y_train)
y_test = sint.input_tensor_via(0, y_test)

# use "nearest" option for deterministic result
# otherwise the Gini coefficients vary slightly from run to run
# resulting in different trees

sfix.set_precision_from_args(program)

from Compiler.decision_tree import TreeClassifier

tree = TreeClassifier(max_depth=5, n_threads=2)

# plain training
tree.fit(X_train, y_train)

# output difference between truth and prediction
print_ln('%s', (tree.predict(X_test) - y_test.get_vector()).reveal())

# output tree
tree.output()

# training with level-wise accuracy output
tree.fit_with_testing(X_train, y_train, X_test, y_test)
