import logging
import time

import cupy as cp
import zmq

from flandre.beamformer.das import gen_pwi
from flandre.beamformer.dist import direct_dist
from flandre.config import C
from flandre.nodes.Node import Node
from flandre.utils.Config import DeviceConfig
from flandre.utils.Msg import ImageArgMsg, Msg, BeamformerMsg, RfMatMsg, RfFrameMsg
from flandre.utils.RfMat import RfMat
from flandre.utils.RfMeta import RfSequenceMeta

logger = logging.getLogger(__name__)


class Beamformer(Node):
    topics = [BeamformerMsg]

    def __init__(self, level=logging.INFO):
        super(Beamformer, self).__init__(level=level)
        self.muxer_req_socket: zmq.Socket = None

    def custom_setup(self):
        self.muxer_req_socket: zmq.Socket = self.c.ctx.socket(zmq.REQ)
        self.muxer_req_socket.connect(C.muxer_rep_socket)
        self.c.poller.register(self.muxer_req_socket, zmq.POLLIN)

    def process_pwi(self, data: RfMat, arg: ImageArgMsg, pwi):
        if data is None:
            return
        d2 = (data
              # .dct(80, 1500)
              .dct_center(arg.dct_center, arg.dct_bandwidth)
              .call(lambda m: m.astype(cp.int16))
              .call(pwi)
              .call(cp.asarray, order='C')
              .argrelextrema()
              .conv_guass(b=arg.beta * 0.01)
              .time_gain_compensation_linear_float(10)
              .crop_center(arg.t_start, arg.t_end)
              .rotate90()
              .grey()
              .cpu()
              # .call(cv2.cvtColor, cv2.COLOR_GRAY2RGB)
              # .pseudo_color()
              )
        # time.sleep(0.01)
        self.send(RfMatMsg(d2))
        # self.send(RGB888Msg(d2.__bytes__(), d2.w, d2.h))

    def loop(self):
        dc = DeviceConfig()
        pwi, _ = gen_pwi(direct_dist(dc), dc)
        last_v2 = 5900
        last_f_rows = 0
        while True:
            self.muxer_req_socket.send(b'')
            r = dict(self.c.poller.poll())
            if self.c.sub in r:
                msg = self.recv()
            if self.muxer_req_socket in r:
                msg: BeamformerMsg = Msg.decode_msg(self.muxer_req_socket.recv())
                if msg.value == b'init':
                    time.sleep(1)
                    continue
                if msg.value == b'nop':
                    continue
                r = msg.value
                id2 = r.index(Msg.magic(), 1)
                arg_msg: ImageArgMsg = Msg.decode_msg(r[:id2])
                rf_frame_msg: RfFrameMsg = Msg.decode_msg(r[id2:])
                mat = RfMat.from_rf_frame(rf_frame_msg.rf_frame, 'gpu')
                if arg_msg.v2 != last_v2 or arg_msg.f_rows != last_f_rows:
                    last_v2 = arg_msg.v2
                    last_f_rows = arg_msg.f_rows
                    dc = DeviceConfig(v2=arg_msg.v2, rows=arg_msg.f_rows)
                    pwi, _ = gen_pwi(direct_dist(dc), dc)
                try:
                    if mat.seq_meta.mode == RfSequenceMeta.RfSequenceMode.PWI:
                        self.process_pwi(mat, arg_msg, pwi)
                except Exception as e:
                    logger.warning(e)
