import matplotlib.image as pltimg
import matplotlib.pyplot as plt
import os
import pandas
import pickle
import pydotplus
from sklearn import tree
from sklearn.tree import DecisionTreeClassifier
def treelearn():
    if os.path.exists("assets/tree.pkl"):
        dtree = pickle.load(open(os.path.join('assets', "tree.pkl"), "rb"))
    else:
        df = pandas.read_csv(os.path.join('assets/data', 'data.csv'))
        columns = ['Fuel','Water','Fertalizer','Carrots','Potatoes','Wheat','X','Y','seeds']
        x = df[columns]
        y = df['back to station']
        dtree = DecisionTreeClassifier()
        dtree = dtree.fit(x, y)
        pickle.dump(dtree, open(os.path.join('assets', "tree.pkl"), "wb"))
        data = tree.export_graphviz(dtree, out_file=None, feature_names=columns)
        graph = pydotplus.graph_from_dot_data(data)
        graph.write_png(os.path.join('assets', 'mytree.png'))
        img = pltimg.imread(os.path.join('assets', 'mytree.png'))
        imgplot = plt.imshow(img)
        plt.show()
    return dtree
def make_decision(tree, Fuel,Water,Fertalizer,Carrots,Potatoes,Wheat,X,Y,seeds):
    decision = tree.predict([[Fuel, Water, Fertalizer, Carrots, Potatoes, Wheat, X, Y, seeds]])
    return decision