import logging
import time
import traceback

import cupy as cp
import zmq

from flandre.beamformer.das import gen_pwi, TFM
from flandre.beamformer.dist import direct_dist
from flandre.beamformer.kernels import dist_mat_to_yids
from flandre import C
from flandre.nodes.Node import Node
from flandre.utils.Config import DeviceConfig
from flandre.utils.Msg import ImageArgMsg, Msg, BeamformerMsg, RfMatMsg, RfFrameMsg, MaxMsg
from flandre.utils.RfFrame import RfFrameFile
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
        self.tfm = 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)
              .crop_center(arg.t_start, arg.t_end)
              #       .time_gain_compensation([
              #    ( (1 - arg.g1 * (1.0 / 128))** 2),
              #    ( (1 - arg.g2 * (1.0 / 128))** 2),
              #    ( (1 - arg.g3 * (1.0 / 128))** 2),
              #    ( (1 - arg.g4 * (1.0 / 128))** 2),
              #    ( (1 - arg.g5 * (1.0 / 128))** 2),
              #    ( (1 - arg.g6 * (1.0 / 128))** 2),
              #    ( (1 - arg.g7 * (1.0 / 128))** 2),
              #    ( (1 - arg.g8 * (1.0 / 128))** 2),
              # ])
              .time_gain_compensation_global((1 - arg.g8 * (1.0 / 128)) ** 2)
              .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 process_tfm(self, data: RfMat, arg: ImageArgMsg, tfm):
        d2 = (data
              .dct_center(arg.dct_center, arg.dct_bandwidth)
              # .filter_max_persent(icfg.bscan_max / 1000, bid=True)
              .call(lambda m: m.astype(cp.int16))
              .call(tfm)
              .call(cp.asarray, order='C')
              .argrelextrema(axis=1)
              .conv_guass(b=arg.beta * 0.01)
              .crop_center(arg.t_start, arg.t_end)
              .time_gain_compensation_global((1 - arg.g8 * (1.0 / 128)) ** 2)
              .rotate90()
              .grey()
              .cpu()
              )
        self.send(RfMatMsg(d2))

    def loop(self):
        dc = DeviceConfig()
        pwi, _, mm = gen_pwi(direct_dist(dc), dc)
        self.send(MaxMsg(mm.item()))
        last_v2 = 5900
        last_f_rows = 0
        last_blake2b = None
        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:])
                current_frame = rf_frame_msg.rf_frame


                if isinstance(current_frame, RfFrameFile):
                    if current_frame.meta.blake2b is not None and current_frame.meta.blake2b == last_blake2b:
                        continue
                mat = RfMat.from_rf_frame(rf_frame_msg.rf_frame, 'gpu')
                # logger.info(mat.frame_meta.blake2b)

                last_blake2b = mat.frame_meta.blake2b
                if mat is None:
                    continue
                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, _, m = gen_pwi(direct_dist(dc), dc)
                    self.send(MaxMsg(m.item()))
                try:
                    if mat.seq_meta.mode == RfSequenceMeta.RfSequenceMode.PWI:
                        self.process_pwi(mat, arg_msg, pwi)
                    if mat.seq_meta.mode == RfSequenceMeta.RfSequenceMode.TFM:
                        if self.tfm is None:
                            dc = DeviceConfig(v2=arg_msg.v2, rows=arg_msg.f_rows)
                            self.tfm = TFM(dc)
                            self.tfm.load_yids(dist_mat_to_yids(direct_dist()))
                        self.process_tfm(mat, arg_msg, self.tfm)
                except Exception as e:
                    logger.warning(e)
                    traceback.print_exc()
