
import numpy
from matplotlib import pyplot

import hark


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

        self.isFirst = True
        self.specgram = []
        self.winlen = 300
        pyplot.ion()
        self.count = 0

    def forward(self, **kwargs):
        self.FFT = kwargs['FFT']
        self.specgram.append(abs(self.FFT[0, :]))

        if self.count % 50 == 0:
            if len(self.specgram) > self.winlen:
                self.specgram = self.specgram[len(self.specgram)-self.winlen:]

            pyplot.clf()
            pyplot.imshow(numpy.flipud(numpy.array(self.specgram).transpose()))

            if self.isFirst:
                pyplot.colorbar()
                self.isFirst = False
            pyplot.xlim([0, self.winlen])
            pyplot.xticks(range(0, self.winlen, 50),
                         range(self.count - self.winlen, self.count, 50))
            pyplot.xlabel("Time [frame]")
            pyplot.ylabel("Frequency bin")
            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()
