
import numpy
from matplotlib import pyplot
import time

import hark


class plotWaveform(hark.node.PythonNode):
    def __init__(self):
        #super().__init__(['WAVE', 'WINDOW_LENGTH', 'UPDATE_INTERVAL'], ['OUTPUT'])
        super().__init__('WAVE', 'OUTPUT')

        self.showwave = None
        #self.wavelen = 16000 * 3
        self.wavelen = 100
        # matplotlib.use('PDF')
        pyplot.ion()
        self.count = 0

    def forward(self, **kwargs):
        if self.showwave is None:
            self.showwave = kwargs['WAVE']
        else:
            self.showwave = numpy.concatenate((self.showwave, kwargs['WAVE']), axis=1)

        if self.count % 100 == 0:
            if self.showwave.shape[1] > self.wavelen:
                self.showwave = self.showwave[:, len(self.showwave)-self.wavelen:]

            pyplot.clf()
            for ch in range(self.showwave.shape[0]):
                pyplot.subplot(self.showwave.shape[0], 1, ch+1)
                pyplot.plot(self.showwave[ch, :])
                pyplot.xlim([0, self.wavelen])
                pyplot.xticks(range(0, self.wavelen, 5000),
                           range(self.count, self.count + self.wavelen, 5000))
                pyplot.ylabel("Magnitude")
                mx = max(abs(self.showwave[ch, :]))
                pyplot.ylim([-mx, mx])

            if self.count == 0:
                pyplot.xlabel("Time [frame]")
            pyplot.pause(0.01)

        self.count += 1
        return {'OUTPUT': numpy.array([1.0, 1.0, 1.0], dtype=numpy.float32)}

#    def __del__(self):
#        pyplot.close()

if __name__ == "__main__":
    import cProfile
    for i in range(5):
        calculate(i, numpy.zeros((3, 512)))
