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 where_fi.collection.protocols import CSIProducer

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 __init__(self, producer: CSIProducer) -> None:
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

        self.producer = producer

        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

    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 listen(self) -> None:
        """
        Listen for incoming data from the selected ingestor, and call the appropriate
        callback.
        """
        data_gen = self.producer()
        while True:
            try:
                sample = next(data_gen)
                if self.csi_callback is not None:
                    self.csi_callback(sample.matrix)
                if self.pre_merge_callback is not None:
                    self.pre_merge_callback(sample.frames)
            except StopIteration:
                break
            except Exception as e:
                self.logger.error(f"Error processing sample: {e}", exc_info=True)
                break

    def start(self) -> None:
        # for receiver in self.receivers:
        #     receiver.thread.start()
        self.buffer_thread = threading.Thread(
            target=self.listen,
        )
        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")

        self.producer.stop()

        if hasattr(self, "scheduler"):
            self.scheduler.stop()
            self.scheduler.thread.join()
        self.buffer_thread.join()

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