import cv2
import numpy as np
import zmq
import cupy as cp

from Msg import BMMsg, ImageArgMsg, KillMsg
from config import socket1, video_width, video_height
from nodes.Node import Node


class Beamformer(Node):
    topics = [ImageArgMsg]

    def __init__(self):
        super(Beamformer, self).__init__()
        self.arg = ImageArgMsg('', 1400)

    def process(self, data: bytes):
        if data is None:
            return
        b = np.frombuffer(data, dtype=np.int16)
        b = b.reshape((256, 1502)).astype(np.float32)
        b = b[:, :self.arg.v1]
        b -= b.min()
        b /= b.max()
        b *= 255
        b = b.astype(np.uint8)
        b = cv2.rotate(b, cv2.ROTATE_90_CLOCKWISE)
        b = cv2.resize(b, (video_width, video_height))
        b = cv2.cvtColor(b, cv2.COLOR_GRAY2RGBA, b)
        self.send(BMMsg(0, b.tobytes()))

    def loop(self):
        s2 = self.context.socket(zmq.PULL)
        s2.setsockopt(zmq.CONFLATE, 1)
        s2.connect(f"tcp://{socket1}")
        self.c.poller.register(s2, zmq.POLLIN)
        buffer = None
        while True:
            socks = dict(self.c.poller.poll())
            if s2 in socks and socks[s2] == zmq.POLLIN:
                buffer = s2.recv()
            if self.c.sub in socks and socks[self.c.sub] == zmq.POLLIN:
                r = self.recv()
                if isinstance(r, KillMsg):
                    if r.name == '':
                        return
                if isinstance(r, ImageArgMsg):
                    self.arg = r
            self.process(buffer)