from sklearn import tree
from sklearn import model_selection
from sklearn import metrics
import numpy as np
import csv
import sys
import pickle
import pydotplus


def conv(str):
    try:
        return float(str)
    except ValueError:
        return str


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


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


def validation(tree_clf, data):
    confusion_matrix = np.zeros((3, 3))
    parameters, classes = data[:, 1:10], data[:, 10:]
    for parameter, pclass in zip(parameters, classes):
        result = tree_clf.predict([parameter])
        confusion_matrix[int(float(result))-1, int(float(pclass))-1] += 1
    print(metrics.classification_report(classes,tree_clf.predict(parameters)))
    return confusion_matrix


def store(tree_clf, confusion_matrix):
    col_names = ["size",
                 "uniqueItems",
                 "minLength",
                 "avgLength",
                 "maxLength",
                 "minRepeatibility",
                 "avgRepeatibility",
                 "maxRepeatibility",
                 "avgRepeatedItems"]
    with open("decisionTree.dt", "wb") as tree_file:
        pickle.dump(tree_clf, 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 = tree.export_graphviz(tree_clf,
                                    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))
    train_data, test_data = load_data(filename)
    tree_clf = plant_trees(train_data)
    confusion_matrix = validation(tree_clf, test_data)
    store(tree_clf, confusion_matrix)


if __name__ == '__main__':
    main()
