nbus-client-api/nbus_api/nbus_module_slave.py

import struct
import time
from threading import Thread, Lock
from typing import Optional
import pandas as pd

import time
import traceback
from threading import Thread, Lock
from typing import Optional, Tuple
import pandas as pd

from nbus_api.nbus_sensor import NBusSensor
from nbus_api.nbus_common_parser import NbusCommonParser
from nbus_hal.crc8 import crc8
from nbus_hal.nbus_generic_port import *
from nbus_types.nbus_address_type import NBusModuleAddress
from nbus_types.nbus_data_fomat import NBusDataValue, NBusDataFormat
from nbus_types.nbus_defines import *
from nbus_types.nbus_exceptions.nbus_api_exception import NBusErrorAPI, NBusErrorAPIType
from nbus_types.nbus_parameter_type import NBusParameterID, NBusParameterValue
from nbus_types.nbus_status_type import NBusStatusType
from nbus_types.nbus_sensor_count_type import NBusSensorCount
from nbus_types.nbus_info_type import NBusModuleInfo
from nbus_types.nbus_sensor_type import NBusSensorType


@beartype
class NBusSlaveModule:
    """
    Class representing nBus slave module.
    """

    def __init__(self, port: NBusPort, module_address: NBusModuleAddress, acquire_delay: float):
        """
        Constructor.

        :param port: serial port
        :param module_address: address of module
        :param device_cnt: number of devices
        """
        self.__port = port
        self.__module_addr = module_address
        self.__params = {}
        self.__devices = {}

        self._lock = Lock()  # thread lock
        self.__in_acquisition = False  # flag when in acquisition
        self.__acquire_delay = acquire_delay  # intermediate delay between data fetching when data not ready
        self.__acquire_thread = None  # thread for data acquisition
        self.__data_raw = bytearray()  # raw data buffer

        self.__df = pd.DataFrame()              # internal data frame
        self.__ts0 = None                       # 0-th timestamp
        self.__payload_size = 0
        self.__damaged_packets = []

    """
    ================================================================================================================
                                                Module General Methods
    ================================================================================================================
    """
    def init(self) -> None:
        """
        Initialize the module from hardware.

        :param load_format: flag to fetch data format
        """
        sensors = self.cmd_get_sensor_type()

        for sen_address, sen_type in sensors.items():
            self.__devices[sen_address] = NBusSensor(self.__port, self.__module_addr, sen_address)
            self.__devices[sen_address].type = sen_type

        self.cmd_get_format()
        self.__calculate_payload_size()

    def get_devices(self) -> dict[NBusSensorAddress, NBusSensor]:
        """
        Get module devices.

        :return: dictionary of connected devices
        """
        return self.__devices

    def get_damaged_packets(self) -> list[Tuple[int, bytearray]]:
        """
        Get damaged packets from module stream.
        :return: list of damaged packets (id, packet)
        """
        return self.__damaged_packets

    """
    ================================================================================================================
                                                Module Get Commands
    ================================================================================================================
    """
    def cmd_get_echo(self, message: bytearray) -> bool:
        """
        Get echo from module.

        :param message: message to send
        :return: status (True = echo, False = no echo)
        """
        _, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_ECHO, message)
        return response == list(message)

    def cmd_get_param(self, parameter: NBusParameterID) -> NBusParameterValue:
        """
        Get single module parameter.

        :param parameter: parameter id
        :return: parameter value
        """
        # get response
        resp_len, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_PARAM,
                                                         bytearray([parameter.value]))

        # parse parameter
        param_id, param_val = NbusCommonParser.parameters_from_response(resp_len, response)[0]

        # store parameter value
        self.__params[param_id] = param_val

        return param_val

    def cmd_get_all_params(self) -> dict[NBusParameterID, NBusParameterValue]:
        """
        Get all module parameters.

        :return: dict of parameter id and parameter value
        """

        # get response
        resp_len, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_PARAM, bytearray([]))

        # parse parameters
        params = NbusCommonParser.parameters_from_response(resp_len, response)

        for param_id, param_val in params:
            # store parameters
            self.__params[param_id] = param_val

        return self.__params.copy()

    def cmd_get_sensor_cnt(self) -> NBusSensorCount:
        """
        Get sensor count.

        :return: count of read-only and read-write sensors
        """
        _, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_SENSOR_CNT, bytearray([]))
        return NBusSensorCount(*response)

    def cmd_get_data(self) -> dict[NBusSensorAddress, list[NBusDataValue]]:
        """
        Get data from all module sensors.

        :return: dict of device addresses and data values
        """

        # get data
        resp_length, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_DATA, bytearray([]))

        # parse data
        begin_idx = 0
        data = {}

        while begin_idx < resp_length:
            device_id = response[begin_idx]

            # handle errors
            if self.__devices[device_id].data_format is None:  # check for format and params
                raise NBusErrorAPI(NBusErrorAPIType.FORMAT_NOT_LOADED)

            values, offset = NbusCommonParser.data_from_response(self.__devices[device_id].data_format,
                                                                 response[begin_idx:])
            data[device_id] = values
            begin_idx += offset + 1

        return data

    def cmd_get_info(self) -> NBusModuleInfo:
        """
        Get module info.

        :return: module info
        """

        response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_INFO, bytearray([]))

        name = str(response[1:9], "ascii")
        typ = str(response[9:12], "ascii")
        uuid = struct.unpack("<I", bytearray(response[12:16]))[0]
        hw = str(response[16:19], "ascii")
        fw = str(response[19:22], "ascii")
        mem_id = struct.unpack("<Q", bytearray(response[22:30]))[0]
        ro_count = int(response[30])
        rw_count = int(response[31])

        return NBusModuleInfo(module_name=name, module_type=typ, uuid=uuid, hw=hw, fw=fw, memory_id=mem_id,
                              read_only_sensors=ro_count, read_write_sensors=rw_count)

    def cmd_get_format(self) -> dict[NBusSensorAddress, NBusDataFormat]:
        """
        Get format of all on-board sensors.

        :return: dict of sensor addresses and sensor formats
        """

        # get response
        resp_length, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_FORMAT,
                                                            bytearray([]))
        begin_idx = 0
        formats = {}

        # parse format
        while begin_idx < resp_length:
            device_id = response[begin_idx]
            device_format = NbusCommonParser.format_from_response(response[begin_idx:begin_idx + 4])

            self.__devices[device_id].data_format = device_format
            formats[device_id] = device_format
            begin_idx += 4

        return formats

    def cmd_get_sensor_type(self) -> dict[NBusSensorAddress, NBusSensorType]:
        """
        Get type of all on-board sensors.

        :return: dict of sensor addresses and sensor types
        """
        # get response
        resp_length, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_GET_SENSOR_TYPE,
                                                            bytearray([]))
        # parse response
        types = {}
        i = 0

        while i < resp_length - 1:
            types[NBusSensorAddress(response[i])] = NBusSensorType(response[i + 1])
            i += 2

        return types

    """
    ================================================================================================================
                                                Module Set Commands
    ================================================================================================================
    """
    def cmd_set_module_stop(self) -> None:
        """
        Stop automatic measuring.

        :return: status
        """
        self.__port.send_module(self.__module_addr, NBusCommand.CMD_SET_STOP, bytearray([]))

    def cmd_set_module_start(self) -> None:
        """
        Start automatic measuring.

        :return: status
        """
        self.__port.send_module(self.__module_addr, NBusCommand.CMD_SET_START, bytearray([]))

    def cmd_set_param(self, param: NBusParameterID, value: NBusParameterValue) -> NBusStatusType:
        """
        Set module parameter.

        :param param: parameter ID
        :param value: parameter value
        :return: status
        """

        # create request packet
        param_id_raw = struct.pack("B", param.value)
        param_val_raw = struct.pack("<I",value)
        param_bytes = bytearray(param_id_raw) + bytearray(param_val_raw)

        # proceed request
        _, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_SET_PARAM, param_bytes)

        # if response is valid, store parameter
        if response[0] == param.value and response[1] == NBusStatusType.STATUS_SUCCESS:
            self.__params[param] = value

        return NBusStatusType(response[1])

    def cmd_set_multi_params(self, params: dict[NBusParameterID, NBusParameterValue]) \
            -> dict[NBusParameterID, NBusStatusType]:
        """
        Set multiple module parameters.

        :param params: parameters
        :return: dict od statuses
        """

        # create request packet
        param_bytes = NbusCommonParser.parameters_to_request(params)

        # send request
        resp_length, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_SET_PARAM,
                                                            bytearray(param_bytes), long_answer=1.0)
        # parse statuses
        statuses = {}

        for i in range(0, resp_length - 1, 2):
            p_id = NBusParameterID(response[i])
            if response[i+1] == NBusStatusType.STATUS_SUCCESS:
                self.__params[p_id] = params[p_id]  # if success, store param

            statuses[p_id] = NBusStatusType(response[i + 1])

        return statuses

    def cmd_set_calibrate(self) -> NBusStatusType:
        """
        Send calibration command.

        :return: calibration status
        """

        resp_length, *response = self.__port.request_module(self.__module_addr,
                                                            NBusCommand.CMD_SET_CALIBRATE, bytearray([]))
        print(response)
        return NBusStatusType(response[0])

    def cmd_set_data(self, data: dict[NBusSensorAddress, list[NBusDataValue]]) \
            -> dict[NBusSensorAddress, NBusStatusType]:
        """
        Set data to read-write sensors.

        :param data:
        :return: operation statuses
        """

        # create request packet
        request = []

        # transform data
        for addr in data.keys():
            # handle errors
            if self.__devices[addr].data_format is None:  # check for format and params
                raise NBusErrorAPI(NBusErrorAPIType.FORMAT_NOT_LOADED)

            raw_data = data[addr]
            request.append(addr)
            request.extend(NbusCommonParser.data_to_request(self.__devices[addr].data_format, raw_data))

        # send request
        resp_length, *response = self.__port.request_module(self.__module_addr, NBusCommand.CMD_SET_DATA,
                                                            bytearray(request))

        # return response statuses
        statuses = {}

        for i in range(0, resp_length - 1, 2):
            statuses[NBusSensorAddress(response[i])] = NBusStatusType(response[i + 1])

        return statuses

    """
    ================================================================================================================
                                            Automatic Data Stream Methods
    ================================================================================================================
    """

    def start_streaming(self) -> None:
        """
        Start data streaming (e.g. auto-cast).
        """
        self.cmd_set_module_start()
        self.__acquire_thread = Thread(target=self.__acquire_callback)
        self.__in_acquisition = True

        # end thread if running
        if self.__acquire_thread is not None and self.__acquire_thread.is_alive():
            self.__acquire_thread.join()

        self.__acquire_thread.start()

    def stop_streaming(self):
        """
        Stop data streaming (e.g. auto-cast).
        """
        self.cmd_set_module_stop()

        self.__in_acquisition = False

        if self.__acquire_thread is not None and self.__acquire_thread.is_alive():
            self.__acquire_thread.join()

        self.__port.flush()

    def fetch_stream_chunk(self) -> pd.DataFrame:
        """
        Fetch data from stream (e.g. auto-cast).
        Can be called anytime.
        It not erase internal dataframe.

        :return: stream data frame
        """
        with self._lock:
            packets = self.__data_raw.split(NBUS_BRIDGE_DATA_HDR)[1:]
            packet_cnt = len(packets) - self.__in_acquisition
            parsed_packets = []

            # parse packets
            for i in range(packet_cnt):
                data = self._parse_data_from_stream_packet(packets[i])
                if data is not None:
                    parsed_packets.append(data)
                else:
                    print("damaged: ", i, packets[i])

            # extend internal dataframe
            if parsed_packets:
                data_frame = pd.DataFrame(parsed_packets)
                self.__df = pd.concat([self.__df, data_frame], ignore_index=True).copy(deep=True)
            else:
                data_frame = pd.DataFrame()

            # erase raw data buffer
            if self.__in_acquisition:
                unparsed_bytes = len(packets[-1]) + NBUS_BRIDGE_DATA_HDR_SIZE
                self.__data_raw = self.__data_raw[-unparsed_bytes:]
            else:
                self.__data_raw = bytearray()

        self._transform_timestamp(data_frame)

        return data_frame

    def fetch_full_stream(self) -> pd.DataFrame:
        """
        Fetch all data from stream (e.g. auto-cast).
        Must be called after stop_streaming() method.
        It will erase internal dataframe.

        :return: stream dataframe
        """
        if self.__in_acquisition:
            return pd.DataFrame()

        self.fetch_stream_chunk()
        df = self.__df
        self._transform_timestamp(df)

        self.__df = pd.DataFrame()
        self.__ts0 = None

        return df

    """
    ================================================================================================================
                                            Internal Helper Methods
    ================================================================================================================
    """

    def _transform_timestamp(self, data_frame: pd.DataFrame) -> None:
        """
        Transform timestamp values in dataframe.
        :param data_frame: dataframe to transform
        """
        if not data_frame.empty and "TS" in self.__df.columns:
            if self.__ts0 is None:
                self.__ts0 = self.__df["TS"].iloc[0]

            data_frame["TS"] -= self.__ts0

    def _parse_data_from_stream_response(self, module_address: NBusModuleAddress, resp_length: int, response: bytearray) \
            -> dict[str, NBusDataValue]:
        """
        Parse data of slave from stream response.
        :param module_address: address of module
        :param resp_length: length of response
        :param response: raw data

        :return: dict of data values
        """
        data_offset = 0
        data = {}

        while data_offset < resp_length:
            device_id = response[data_offset]

            if self.__devices[device_id].data_format is None:
                raise NBusErrorAPI(NBusErrorAPIType.FORMAT_NOT_LOADED)

            values, offset = NbusCommonParser.data_from_response(self.__devices[device_id].data_format, response[data_offset:])
            data_tag = str(module_address) + "." + str(device_id)

            for i in range(len(values)):
                data[data_tag + "." + str(i + 1)] = values[i]

            data_offset += offset + NBUS_SA_SIZE

        return data


    def _parse_data_from_stream_packet(self, data_packet: bytearray) \
            -> Optional[dict[str, NBusDataValue]]:
        """
        Parse the stream-mode data packet.

        :param data_packet: packet to parse

        :return: dictionary of data values or None
        """
        packet_size = len(data_packet)
        packet_crc = crc8(data_packet[:NBUS_CRC_ADDR])
        # check validity

        if packet_size < NBUS_TS_SIZE + NBUS_CRC_SIZE or data_packet[NBUS_CRC_ADDR] != packet_crc:
            return None

        # parse data
        try:
            data_offset = 0
            ts = struct.unpack("<I", data_packet[data_offset: data_offset + NBUS_TS_SIZE])[0]
            data = {"TS": ts}
            data_offset += NBUS_TS_SIZE

            module_addr = data_packet[data_offset]
            data_offset += NBUS_MA_SIZE
            packet = data_packet[data_offset: data_offset + self.__payload_size]

            data |= self._parse_data_from_stream_response(module_addr, self.__payload_size, packet)
            data_offset += self.__payload_size

            return data

        except Exception:
            return None

    def __calculate_payload_size(self):
        """
        Calculate payload size of stream-mode packet.
        :return: payload size in bytes
        """
        for device in self.__devices.values():
            fmt = device.data_format
            self.__payload_size += fmt.byte_length * fmt.samples + NBUS_SA_SIZE

    def __acquire_callback(self) -> None:
        """
        Thread worker for data acquisition.
        """
        while self.__in_acquisition:
            data = self.__port.try_read()

            if not data:
                time.sleep(self.__acquire_delay)
                continue

            with self._lock:
                self.__data_raw.extend(data)