import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import matplotlib.animation as animation


class SubplotAnimation(animation.TimedAnimation):

    def __init__(self):
        fig = plt.figure()
        ax1 = fig.add_subplot(1, 1, 1)

        ax1.set_xlabel('Time')
        ax1.set_ylabel('Energy/sec')

        self.line1 = Line2D([], [], color='black')

        ax1.add_line(self.line1)

        ax1.set_xlim(0, 20)
        ax1.set_ylim(0, 5)
        #ax1.set_aspect('equal', 'datalim')

        animation.TimedAnimation.__init__(self, fig, interval=50, blit=True)


    def _draw_frame(self, framedata):

        self.x, self.y = self.pull_data()

        self.line1.set_data(self.x, self.y)

        self._drawn_artists = [self.line1]


    def pull_data(self):

        data = open('sample_plot.txt', 'r').read()
        data_list = data.split('\n')
        x_list = []
        y_list = []
        for line in data_list:
            if len(line) > 1:

                x, y = line.split(',')
                x_list.append(int(x))
                y_list.append(int(y))

        return x_list, y_list


    def new_frame_seq(self):
        #print('new_frame_seq')
        return iter(range(1))


    def _init_draw(self):
        #print('_init_draw')
        lines =  [self.line1]
        for l in lines:
            l.set_data([], [])



ani = SubplotAnimation()
#ani.save('test_sub.mp4')
plt.show()