import sklearn as sk
from sklearn.tree import export_graphviz
import numpy as np
import csv
import sys
import pickle
import pydotplus


def load_data(filename):
    with open(filename, newline='') as csvfile:
        reader = csv.reader(csvfile)
        try:
            data = []
            for row in reader:
                data.append(row)
            print(data)
            return np.array(data)
        except csv.Error as e:
            sys.exit("file {}, line {}".format(filename, row))


def plant_trees(data):
    parameters, classes = data[:], data[:]
    tree = sk.tree.DecisionTreeClassifier()
    tree = tree.fit(parameters, classes)  # TODO: add cross-validation
    return tree


def validation(tree, data):
    confusion_matrix = np.zeros((3, 3))
    parameters = np.array(data)
    classes = data[len(data) - 1]
    for parameter, pclass in zip(parameters, classes):
        result = tree.predict(parameter)
        confusion_matrix[result, pclass] += 1
    return confusion_matrix


def report(tree, confusion_matrix):
    print(tree)
    print(confusion_matrix)


def store(tree, confusion_matrix):
    col_names = []
    with open("decisionTree.dt", "w") as tree_file:
        pickle.dump(tree, tree_file)
    with open('matrix.txt', 'w') as matrixFile:
        for row in confusion_matrix:
            matrixFile.write(' '.join([str(a) for a in row]) + '\n')
    dot_data = export_graphviz(tree,
                               feature_names=col_names,
                               out_file=None,
                               filled=True,
                               rounded=True)
    pydot_data = pydotplus.graph_from_dot_data(dot_data)
    pydot_data.write_png("decisionTree.png")


def main():
    filename = sys.argv[1]
    print("Loading file {}".format(filename))
    data = load_data(filename)
    tree = plant_trees(data)
    confusion_matrix = validation(tree)
    report(tree, confusion_matrix)
    store(tree, confusion_matrix)


if __name__ == '__main__':
    main()
