dissertation/where_fi/application.py

import logging
import multiprocessing as mp
import threading
import time
from typing import Any, Callable, NamedTuple

import numpy.typing as npt

from .collection import CSIMatrix, ingest
from .config import config
from .visualise import server as visualise


class Receiver(NamedTuple):
    ip: ingest.Host
    receiver: ingest.FeitReceiver
    thread: threading.Thread


class Scheduler:
    """
    A simple scheduler that runs a function at a given rate.
    """

    def __init__(self, rate: float, func: Callable[[], None]) -> None:
        self.rate = rate
        self.func = func
        self.active = True

    def run(self) -> None:
        """
        Run the scheduler in a loop, calling the function at the given rate.
        """
        while self.active:
            self.func()
            time.sleep(1 / self.rate)

    def start(self) -> None:
        """
        Start the scheduler in a separate thread.
        """
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

    def stop(self) -> None:
        """
        Stop the scheduler.
        """
        self.active = False


class CSIApplication:
    """
    A high-level application that manages CSI data ingestion and processing.

    This allows a simple interface to be used for the main CLI logic, without having to
    deal with the threads and queues directly.

    It can be used using decortors to register callbacks for different stages of the
    processing pipeline:

        - `on_pre_merge`: Called with the raw CSI data (including headers) from each
                          receiver, before it is merged into a single matrix.
        - `on_sample`: Called once per sample with the merged CSI data.
        - `on_process`: Called at the processing sample rate, without any data (this is
                        mostly used as a scheduler).
    """

    def create_receivers(self) -> list[Receiver]:
        """
        Create a thread for each receiver, so that data can be received in parallel.

        The receivers are created with a queue that is used to store the received data,
        to be processed later. The receivers are created based on the configuration
        file, which contains the IP addresses of the receivers.
        """
        feit_receivers = [
            ingest.FeitReceiver(ip, mp.Queue(config.collection_sample_rate))
            for ip in config.receive_hosts
        ]
        return [
            Receiver(
                ip,
                receiver,
                threading.Thread(target=receiver.listen),
            )
            for receiver, ip in zip(feit_receivers, config.receive_hosts, strict=True)
        ]

    def __init__(self) -> None:
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

        self.receivers = self.create_receivers()
        self.transmitter = ingest.FeitTransmitter()

        self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
            config.processing_sample_rate
        )
        self.webapp = visualise.Webapp()
        self.webapp_thread = threading.Thread(
            target=self.webapp.start, args=(self.webapp_queue,)
        )

        self.csi_callback = None
        self.pre_merge_callback = None
        self.processing_callback = None

        self.buffer = ingest.CSIAntennaJoin(
            {r.ip: r.receiver.queue for r in self.receivers}
        )

    def visualise_data(
        self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
    ) -> None:
        """
        Update a visualisation with the given data. The available visualisations are as
        per visualise.server.all_figures.
        """
        if not self.webapp_queue.full():
            self.webapp_queue.put(visualise.VisualiserData(data, dtype))

    def on_sample(
        self, func: Callable[[CSIMatrix], None]
    ) -> Callable[[CSIMatrix], None]:
        """
        Decorator to register a callback for the sample preprocessing.
        """

        def decorator(data: CSIMatrix) -> None:
            self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
            func(data)

        self.csi_callback = decorator
        return decorator

    def on_pre_merge(
        self,
    ) -> Callable[[Callable[[dict[ingest.Host, ingest.CSI]], None]], None]:
        def decorator(func: Callable[[dict[ingest.Host, ingest.CSI]], None]) -> None:
            self.pre_merge_callback = func

        return decorator

    def on_process(self, func: Callable[[], None]) -> Callable[[], None]:
        """
        Decorator to register a callback for the sample processing.
        """
        self.processing_callback = func
        return func

    def start(self) -> None:
        for receiver in self.receivers:
            receiver.thread.start()
        self.buffer_thread = threading.Thread(
            target=self.buffer.process_forever,
            args=(self.csi_callback, self.pre_merge_callback),
        )
        self.buffer_thread.start()
        self.webapp_thread.start()

        if self.processing_callback is not None:
            self.scheduler = Scheduler(
                config.processing_sample_rate, self.processing_callback
            )
            self.scheduler.start()
        try:
            self.webapp_thread.join()
        except KeyboardInterrupt:
            self.stop()

    def stop(self) -> None:
        self.logger.info("Stopping application")
        for receiver in self.receivers:
            receiver.receiver.active = False
            receiver.thread.join()

        self.transmitter.active = False

        if hasattr(self, "scheduler"):
            self.scheduler.stop()
            self.scheduler.thread.join()
        self.buffer.active = False
        self.buffer_thread.join()

        self.webapp.active = False
        self.webapp_thread.join()
        self.logger.info("Application stopped")