import joblib
import matplotlib.pyplot as plt
import pandas
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree
from Sectors_types import Sectors_types


decisions = ["decision"]
attributes = ["marking", "weekend", "c_popularity", "payment_delay", "payed", "shipping_method"]


class DecisionTree:

    def __init__(self, marking, weekend, c_popularity, payment_delay, payed, shipping_method):
        self.decision = self.decision(marking, weekend, c_popularity, payment_delay, payed, shipping_method)

    def tree(self):
        dataset = pandas.read_csv('./decision_tree/csv_file.csv')
        x = dataset[attributes]
        y = dataset[decisions]
        decision_tree = DecisionTreeClassifier()
        decision_tree = decision_tree.fit(x.values, y)

        return decision_tree

    # return decision made from tree and attributes
    def decision(self, marking, weekend, c_popularity, payment_delay, payed, shipping_method):
        decision_tree = self.tree()
        decision = decision_tree.predict(
            [[marking, weekend, c_popularity, payment_delay, payed, shipping_method]])
        if decision == 1:
            decision = Sectors_types.fragile
        elif decision == 2:
            decision = Sectors_types.normal
        elif decision == 3:
            decision = Sectors_types.shipping_tomorrow
        elif decision == 4:
            decision = Sectors_types.shipping_today
        return decision

    def tree_as_txt(self, decision_tree):
        with open('tree.txt', "w") as file:
            file.write(export_text(decision_tree))

    def tree_to_png(self, decision_tree):
        fig = plt.figure()
        plot_tree(decision_tree, feature_names=attributes, filled=True)
        plt.title("Decision tree")
        # plt.show()
        fig.savefig('tree.png')

    # def tree_to_structure(self,decision_tree):
    #     joblib.dump(decision_tree, 'tree_model')
    #
    # def tree_from_structure(self, file):
    #     return joblib.load(file)
    #
    #
    #
    # drzewo = tree()
    # # tree_as_txt(drzewo)
    # tree_to_png(drzewo)
    # # tree_to_structure(drzewo)