import asyncio
import os
import logging

import numpy as np
from quamash import QEventLoop, QtWidgets
import pyqtgraph as pg
from sipyco.pipe_ipc import AsyncioChildComm
from sipyco import pyon
 

class SpectrogramWidget(pg.PlotWidget):
    def __init__(self, block_size=4096):
        super(SpectrogramWidget, self).__init__()
 
        self.block_size = block_size
 
        self.img = pg.ImageItem()
        self.addItem(self.img)
 
        self.img_array = np.zeros((100, block_size))
 
        pos = np.array([0., 1., 0.5, 0.25, 0.75])
        color = np.array([[0,255,255,255], [255,255,0,255], [0,0,0,255], (0, 0, 255, 255), (255, 0, 0, 255)], dtype=np.ubyte)
        cmap = pg.ColorMap(pos, color)
        lut = cmap.getLookupTable(0.0, 1.0, 256)
 
        self.img.setLookupTable(lut)
        self.img.setLevels([-50,40])
  
        self.show()

    def update(self, block):
        self.img_array = np.roll(self.img_array, -1, 0)
        self.img_array[-1:] = block 
        self.img.setImage(self.img_array, autoLevels=True)


class IPCClient(AsyncioChildComm):
    def set_close_cb(self, close_cb):
        self.close_cb = close_cb

    async def read_pyon(self):
        line = await self.readline()
        return pyon.decode(line.decode())

    async def listen(self, spectrogram):
        while True:
            obj = await self.read_pyon()
            try:
                action = obj["action"]
                if action == "update":
                    spectrogram.update(obj["data"])
                if action == "terminate":
                    self.close_cb()
                    return
            except:
                logging.error("error processing parent message",
                              exc_info=True)
                self.close_cb()


def main():
    app = QtWidgets.QApplication([])
    loop = QEventLoop(app)
    asyncio.set_event_loop(loop)

    try:
        ipc = IPCClient(os.getenv("NOPTICA2_IPC"))
        loop.run_until_complete(ipc.connect())
        try:
            main_widget = SpectrogramWidget()
            main_widget.show()
            ipc.set_close_cb(main_widget.close)
            asyncio.ensure_future(ipc.listen(main_widget))
            loop.run_forever()
        finally:
            ipc.close()
    finally:
        loop.close()


if __name__ == "__main__":
    main()