Source code for audri

'''Provides functionality for automated driving.
AUDRI is able to receive feature vectors, learn from them, and then decide upon
an action to perform when presented with another feature vector
'''
# import pydot
# from io import StringIO
from sklearn import tree

class Agent():
    '''A learning agent who will learn how to drive from data it is given,
    and is able to use its training to command a car'''

    def __init__(self):
        self._model = tree.DecisionTreeClassifier(max_depth=3, max_leaf_nodes=4)

    def action(self, stateVector):
        '''Return the predicted label for the given state vector

        :param stateVector: (:class:`dict`) The unlabeled state vector
        '''
        pred = self._model.predict(
            [[stateVector[k] for k in sorted(stateVector)]])[0]
        return int(pred)

    def train(self, data):
        '''Train the agent using a set of data

        :param data: (:class:`list`) A list of :class:`dict` state vectors
        '''
        y = list(map(lambda row: row.pop('action'), data))
        data = list(map(lambda row: list(row.values()), data))
        self._model = self._model.fit(data, y)
        # with open("action_classifier.txt", "w") as f:
        #     out = StringIO()
        #     tree.export_graphviz(self._model, out_file=out,
        #             feature_names=['aheadDistance', 'currentLane'],
        #             class_names=['noAction', 'moveLeft', 'moveRight'])
        #     tree.export_graphviz(self._model, out_file=out)
        #     tree.
        #     pydot.graph_from_dot_data(out.getvalue())[0].write_pdf("test.pdf")

    def load(self, name):
        # TODO
        pass

    def save(self, name):
        # TODO
        pass