
import numpy
from matplotlib import pyplot

import hark


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

        self.countdown = 10
        #self.isFirst = True
        self.im = None
        self.cb = None
        self.spectrum = []
        self.winlen = 100
        pyplot.ion()
        self.count = 0

    def forward(self, **kwargs):
        tmp = kwargs['SPECTRUM']
        if(numpy.sum(tmp) == 0):
            return {'OUTPUT': None};
        self.spectrum.append(tmp)
        self.countdown -= 1

        if self.countdown == 0:
            self.countdown = 10
            if len(self.spectrum) > self.winlen:
                self.spectrum = self.spectrum[len(self.spectrum)-self.winlen:]
            arr = numpy.flipud(numpy.array(self.spectrum).transpose())
            if self.im is None:
                pyplot.clf()
                self.im = pyplot.imshow(arr, aspect="auto", interpolation="nearest")
                self.cb = pyplot.colorbar()
                pyplot.xlim([0, self.winlen])
                pyplot.xlabel("Time")
                pyplot.ylabel("Azimuth")
            else:
                self.im.set_array(arr)
                ext = self.im.get_extent()
                self.im.set_extent((ext[0], arr.shape[1], ext[2], ext[3]))
                self.cb.set_clim(vmin=arr.min(), vmax=arr.max())
                self.cb.draw_all()
            pyplot.pause(0.01)
        
        self.count += 1
        return {'OUTPUT': numpy.array([1.0, 1.0, 1.0], dtype=numpy.float32)}

