dissertation/src/collection/ingest.py

import logging
import multiprocessing as mp
import socket
import struct
import subprocess
import threading
import time
from datetime import datetime
from typing import Callable, NamedTuple, NoReturn

import numpy as np
import numpy.typing as npt

from .. import config
from .csi_frame import CSI

Host = tuple[str, int]
CSICallback = Callable[[npt.NDArray[np.complex128]], None]


class FeitHost:
    def __init__(self, host: Host, command: str) -> None:
        self.command = command
        self.host = host
        self.logger = logging.getLogger(
            f"{__name__}.{self.__class__.__name__}-{self.host[0]}"
        )
        self.checker = threading.Thread(target=self.check_continuous)
        self.checker.start()

    def check_connection(self) -> bool:
        feitcsi_status = subprocess.run(
            f"ssh root@{self.host[0]} pgrep feitcsi",
            check=False,
            stdout=subprocess.DEVNULL,
            shell=True,
        )
        return feitcsi_status.returncode == 0

    def connect(self) -> None:
        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server.connect(self.host)
        self.server.send(b"stop\n")
        self.server.send(self.command.encode())
        self.logger.info(f"Connected to {self.host}")

    def check_continuous(self) -> NoReturn:
        """
        Repeatedly check if the FeitCSI service is running

        Because FeitCSI is using TCP, we will get no information if the service stops,
        or the computer is not reachable. In order to make debugging easier, this checks
        and logs continuously if the service is running.
        """
        last_status = False
        while True:
            if not self.check_connection():
                self.logger.error(f"FeitCSI is not running on {self.host[0]}")
                last_status = False
            else:
                if not last_status:
                    self.connect()
                last_status = True
            time.sleep(1)


class FeitTransmitter(FeitHost):
    def __init__(self) -> None:
        command = (
            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
            f"--channel-width {config.CHANNEL_WIDTH} "
            f"--format {config.FRAME_FORMAT} "
            f"--mode inject -s 1 --verbose "
            f"--inject-delay {1_000_000 // config.SAMPLE_RATE}"
        )
        super().__init__(config.INJECT_HOST, command)


class FeitReceiver(FeitHost):
    def __init__(self, host: Host) -> None:
        command = (
            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
            f"--channel-width {config.CHANNEL_WIDTH} "
            f"--format {config.FRAME_FORMAT} "
            f"--mode measure"
        )
        super().__init__(host, command)

    def listen(self, queue: "mp.Queue[CSI]") -> NoReturn:
        prev_time = datetime.now()
        self.logger.info("Listening for CSI data")
        while True:
            while not hasattr(self, "server"):
                time.sleep(0.1)
            # This is the max size of a UDP packet. The size of the actual CSI
            # packet will depend on the frame format and channel width, which
            # changes the number of subcarriers
            data = self.server.recv(65535)
            try:
                csidata = CSI(data)
                self.logger.debug(
                    f"Received CSI data after {datetime.now() - prev_time}"
                )
                prev_time = datetime.now()
                queue.put(csidata)
            except struct.error:
                self.logger.error("Failed to parse CSI data")


class CSIProcessor:
    def __init__(
        self,
        receiver_connections: dict[Host, "mp.Queue[CSI]"],
    ) -> None:
        self.pending_data: dict[Host, tuple[datetime, CSI]] = {}
        self.pending_data_lock = mp.Lock()
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
        self.last_processed = datetime.now()
        self.connections = receiver_connections

    def add_data(self, host: Host, data: CSI) -> None:
        if (
            host in self.pending_data
            and self.last_processed < self.pending_data[host][0]
        ):
            self.logger.warning(
                f"Skipping data from {host} at {self.pending_data[host][0]}"
            )

        with self.pending_data_lock:
            self.pending_data[host] = (datetime.now(), data)

    def is_ready(self) -> bool:
        for host in config.RECEIVE_HOSTS:
            if (
                host not in self.pending_data
                or self.pending_data[host][0] <= self.last_processed
            ):
                return False
        return True

    def process_data(
        self,
        callback: CSICallback | None = None,
        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
    ) -> None:
        self.last_processed = datetime.now()
        if pre_merge_callback is not None:
            pre_merge_callback(
                {host: data[1] for host, data in self.pending_data.items()}
            )
        antenna_data = [
            np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
            for ip, antenna in config.ANTENNA_ORDER
        ]

        # We have data from all servers
        all_data = np.concat(antenna_data, axis=1)

        if callback:
            callback(all_data)

    def process_forever(
        self,
        callback: CSICallback | None = None,
        pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
    ) -> NoReturn:
        while True:
            for ip, queue in self.connections.items():
                while not queue.empty():
                    self.add_data(ip, queue.get())
            if self.is_ready():
                self.process_data(callback, pre_merge_callback=pre_merge_callback)
            else:
                self.logger.debug("Not all data is ready")
            time.sleep(0.0005)


class Receiver(NamedTuple):
    ip: Host
    receiver: FeitReceiver
    queue: "mp.Queue[CSI]"


def start_processing(
    csi_callback: CSICallback | None = None,
    pre_merge_callback: Callable[[dict[Host, CSI]], None] | None = None,
) -> None:
    receivers = [
        Receiver(ip, FeitReceiver(ip), mp.Queue(config.SAMPLE_RATE))
        for ip in config.RECEIVE_HOSTS
    ]

    FeitTransmitter()

    # Start injecting CSI frames
    processor = CSIProcessor({r.ip: r.queue for r in receivers})

    receiver_processes = [
        threading.Thread(target=r.receiver.listen, args=(r.queue,)) for r in receivers
    ]

    for proc in receiver_processes:
        proc.start()

    processing_thread = mp.Process(
        target=processor.process_forever, args=(csi_callback, pre_merge_callback)
    )
    processing_thread.start()
    try:
        processing_thread.join()
    except KeyboardInterrupt:
        return