flandre/flandre/nodes/Muxer.py

import logging
import struct
import time
from threading import Thread

import zmq

from flandre import C
from flandre.nodes.Device import Device, DeviceCmd
from flandre.nodes.Node import Node
from flandre.utils.Msg import (
    BeamformerMsg,
    DeviceEnabledMsg,
    ImageArgMsg,
    ImagingConfigNameListMsg,
    KillMsg,
    RfFrameMsg,
    RobotRtsiMsg,
    SeqMetaMsg,
    SetDeviceConfigMsg,
    SetPlayMode,
    SetSeqMetaMsg,
)
from flandre.utils.RfFrame import RfFrameMemory, b2t
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta

logger = logging.getLogger(__name__)


class Muxer(Node):
    topics = [
        SetSeqMetaMsg,
        SetPlayMode,
        SetDeviceConfigMsg,
        RfFrameMsg,
        ImageArgMsg,
        RobotRtsiMsg,
        SeqMetaMsg,
        DeviceEnabledMsg,
    ]

    def __init__(self, level=logging.INFO):
        super(Muxer, self).__init__(level=level)
        self.seq_meta = None
        self.seq_meta_live: RfSequenceMeta = None
        self.seq_meta_playback = None
        self.play_mode: str | None = None
        self.rep_socket: zmq.Socket = None
        self.req_driver_socket: zmq.Socket = None
        self.driver_pull_socket: zmq.Socket = None
        self.playback_rf_msg: RfFrameMsg | None = None
        self.device_enabled = False
        self.driver_data_raw = b""
        self.run_p_thread = True

    def custom_setup(self):
        self.rep_socket: zmq.Socket = self.c.ctx.socket(zmq.REP)
        self.rep_socket.bind(f"tcp://localhost:{C.muxer_rep_port}")
        self.req_driver_socket: zmq.Socket = self.c.ctx.socket(zmq.REQ)
        # self.driver_pull_socket = self.c.ctx.socket(zmq.PULL)
        # self.driver_pull_socket.connect(C.live_push_socket)
        # self.req_driver_socket.connect(C.driver_rep_socket)
        self.req_driver_socket.connect(C.live_rep_socket)
        self.c.poller.register(self.rep_socket, zmq.POLLIN)

    def p_thread(self):
        while self.run_p_thread:
            if self.play_mode == "live":
                # ii = self.driver_pull_socket.poll(timeout=1000)
                # if ii > 0:
                #     self.driver_data_raw = self.driver_pull_socket.recv()
                self.req_driver_socket.send(
                    struct.pack("i", Device.magic)
                    + struct.pack("i", DeviceCmd.GetData.value)
                )
                self.driver_data_raw = self.req_driver_socket.recv()
            else:
                time.sleep(1)

    def handle_rep_socket(self):
        self.rep_socket.recv()
        if self.play_mode is None:
            self.rep_socket.send(BeamformerMsg(b"init").encode_msg())
            return
        match self.play_mode:
            case "playback":
                # logger.warning(f'test, {self.playback_rf_msg}')
                if self.playback_rf_msg is None:
                    self.rep_socket.send(BeamformerMsg(b"nop").encode_msg())
                    return
                data_msg = self.playback_rf_msg
            case "live":
                if not self.device_enabled:
                    self.rep_socket.send(BeamformerMsg(b"init").encode_msg())
                    logger.warning("Device not enabled")
                    return
                # self.req_driver_socket.send(b'')
                # self.driver_data_raw = self.req_driver_socket.recv()
                if self.driver_data_raw == b"":
                    # todo fixit driver no empty
                    self.rep_socket.send(BeamformerMsg(b"nop").encode_msg())
                    return
                # _, sequence_id, encoder = struct.unpack_from('=IQi', self.driver_data_raw)
                # ts, sequence_id, encoder, driver_data_body = b2t(self.driver_data_raw)
                (
                    sequence_id,
                    encoder,
                    host_ts,
                    device_ts_low,
                    device_ts_high,
                    driver_data_body,
                ) = b2t(self.driver_data_raw)
                data_msg = RfFrameMsg(
                    0,
                    RfFrameMemory(
                        RfFrameMeta(
                            encoder=encoder,
                            sequence_id=sequence_id,
                            robot_x=self.rtsi.pos[0],
                            robot_y=self.rtsi.pos[1],
                            robot_z=self.rtsi.pos[2],
                            robot_roll=self.rtsi.pos[3],
                            robot_pitch=self.rtsi.pos[4],
                            robot_yal=self.rtsi.pos[5],
                            robot_force_x=self.rtsi.force[0],
                            robot_force_y=self.rtsi.force[1],
                            robot_force_z=self.rtsi.force[2],
                            robot_force_roll=self.rtsi.force[3],
                            robot_force_pitch=self.rtsi.force[4],
                            robot_force_yal=self.rtsi.force[5],
                        ),
                        self.seq_meta_live,
                        driver_data_body,
                    ),
                )
            case _:
                raise NotImplementedError()
        # if (data_msg.data.__len__() // 2) != data_msg.rf_frame.prod():
        #     self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
        #     return
        self.rep_socket.send(
            BeamformerMsg(self.arg.encode_msg() + data_msg.encode_msg()).encode_msg()
        )

    def loop(self):
        t = Thread(target=self.p_thread).start()
        self.rtsi = RobotRtsiMsg(
            pos=(0, 0, 0, 0, 0, 0),
            force=(0, 0, 0, 0, 0, 0),
        )
        device_socket = self.context.socket(zmq.PULL)

        self.arg = ImageArgMsg("", t_start=0, t_end=1499)
        self.c.poller.register(device_socket, zmq.POLLIN)
        self.send(
            ImagingConfigNameListMsg(
                [path.stem for path in C.imaging_config_folder.glob("*.json")]
            )
        )
        while True:
            socks = dict(self.c.poller.poll())
            for k in socks:
                if k == device_socket:
                    pass
                if k == self.rep_socket:
                    self.handle_rep_socket()
                if k == self.c.sub:
                    msg = self.recv()
                    if isinstance(msg, KillMsg):
                        self.run_p_thread = False
                        if msg.name == "":
                            return
                    elif isinstance(msg, RfFrameMsg):
                        if msg.sender == 1:
                            self.playback_rf_msg = msg
                    elif isinstance(msg, ImageArgMsg):
                        self.arg = msg
                    elif isinstance(msg, SeqMetaMsg):
                        match msg.target:
                            case "live":
                                self.seq_meta_live = RfSequenceMeta.from_name(msg.name)
                    elif isinstance(msg, SetPlayMode):
                        logger.info(f"set playmode {msg}")
                        self.play_mode = msg.value
                    elif isinstance(msg, RobotRtsiMsg):
                        self.rtsi = msg
                    elif isinstance(msg, DeviceEnabledMsg):
                        self.device_enabled = msg.value