dissertation/where_fi/visualise/server/__init__.py

import logging
import multiprocessing as mp
import queue
import threading
import uuid
from concurrent import futures
from dataclasses import dataclass
from enum import Enum
from typing import Any, Generator

import grpc
import numpy as np
import numpy.typing as npt

from ...config import config
from .generated import figure_pb2, figure_pb2_grpc
from .generated.figure_type import heatmap_pb2, histogram_pb2, line_pb2

logger = logging.getLogger(__name__)

clients: dict[str, list[queue.Queue[figure_pb2.FigureData]]] = {}
clients_lock = threading.Lock()


class DataType(Enum):
    RAW_CSI = 1
    UNWRAPPED_PHASE = 2
    PROCESSED_CSI = 3
    HEATMAP = 4
    EIGENVALUES = 5


@dataclass
class VisualiserData:
    data: npt.NDArray[Any]
    dtype: DataType


figures = {
    "raw_phase": figure_pb2.Figure(
        uuid=str(uuid.uuid4()),
        title="Raw CSI Phase",
        x_label="Subcarrier",
        y_label="Phase",
    ),
    "unwrapped_phase": figure_pb2.Figure(
        uuid=str(uuid.uuid4()),
        title="Unwrapped CSI Phase",
        x_label="Subcarrier",
        y_label="Unwrapped phase",
    ),
    "aoa_eig": figure_pb2.Figure(
        uuid=str(uuid.uuid4()),
        title="AoA Eigenvalues",
        x_label="Eigenvalue",
        y_label="Frequency",
        logx=True,
    ),
    "processed_phase": figure_pb2.Figure(
        uuid=str(uuid.uuid4()),
        title="Preprocessed CSI Phase",
        x_label="Subcarrier",
        y_label="Phase",
    ),
    "aoa_heatmap": figure_pb2.Figure(
        uuid=str(uuid.uuid4()),
        title="AoA Heatmap",
        x_label="AoA",
        y_label="ToF",
    ),
}


class FigureServer(figure_pb2_grpc.FigureServiceServicer):
    def GetFigure(
        self, request: figure_pb2.FigureRequest, context: grpc.ServicerContext
    ) -> Generator[figure_pb2.Figure, None, None]:
        for figure in figures.values():
            yield figure

    def GetFigureUpdate(
        self, request: figure_pb2.FigureDataRequest, context: grpc.ServicerContext
    ) -> Generator[figure_pb2.FigureData, None, None]:
        logger.info(
            f"Received request for figure data stream for figure {request.uuid}"
        )

        with clients_lock:
            q: queue.Queue[figure_pb2.FigureData] = queue.Queue()
            if request.uuid not in clients:
                clients[request.uuid] = []
            clients[request.uuid].append(q)

        try:
            while context.is_active():
                try:
                    yield q.get(timeout=1)
                except queue.Empty:
                    pass
        finally:
            with clients_lock:
                clients[request.uuid].remove(q)


def add_data(dtype: DataType, new_data: npt.NDArray[np.complex128]) -> None:
    match dtype:
        case DataType.RAW_CSI | DataType.PROCESSED_CSI:
            uuid = (
                figures["raw_phase"].uuid
                if dtype == DataType.RAW_CSI
                else figures["processed_phase"].uuid
            )
            lines = [
                line_pb2.LineChartData.Line(
                    y=np.angle(new_data)[:, i, 0], label=f"Antenna {i}"
                )
                for i in range(new_data.shape[1])
            ]
            linechart = line_pb2.LineChartData(lines=lines)
            for q in clients.get(uuid, []):
                q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
        case DataType.HEATMAP:
            uuid = figures["aoa_heatmap"].uuid
            heatmap = heatmap_pb2.HeatmapData(
                uuid=uuid,
                data=new_data.flatten(),
                width=new_data.shape[1],
                height=new_data.shape[0],
                x_min=0,
                x_max=np.pi,
                y_min=0,
                y_max=config.music.heatmap.tof_max,
            )
            for q in clients.get(uuid, []):
                q.put(figure_pb2.FigureData(uuid=uuid, heatmap=heatmap))
            pass
        case DataType.UNWRAPPED_PHASE:
            uuid = figures["unwrapped_phase"].uuid
            lines = [
                line_pb2.LineChartData.Line(y=new_data[:, i, 0], label=f"Antenna {i}")
                for i in range(new_data.shape[1])
            ]
            linechart = line_pb2.LineChartData(lines=lines)
            for q in clients.get(uuid, []):
                q.put(figure_pb2.FigureData(uuid=uuid, line=linechart))
        case DataType.EIGENVALUES:
            uuid = figures["aoa_eig"].uuid
            magn = np.abs(new_data)
            bins = np.logspace(np.log10(magn.min()), np.log10(magn.max()), 10)
            hist, _ = np.histogram(magn, bins=bins)
            logger.info(f"bins: {bins}, hist: {hist}")
            logger.info(
                f"Updating histogram with {len(bins)} bins, and {len(hist)} bars"
            )
            logger.info(bins)
            series = [histogram_pb2.HistogramSeries(data=hist)]
            histogram = histogram_pb2.HistogramData(data=series, bins=bins)
            for q in clients.get(uuid, []):
                q.put(figure_pb2.FigureData(uuid=uuid, histogram=histogram))


def listen_for_data(data_queue: "mp.Queue[VisualiserData]") -> None:
    while True:
        data = data_queue.get()
        add_data(data.dtype, data.data)


def start(data_queue: "mp.Queue[VisualiserData]") -> None:
    logger.info("Visualisation server shut down")

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    figure_pb2_grpc.add_FigureServiceServicer_to_server(FigureServer(), server)
    server.add_insecure_port("[::]:50051")
    server.add_insecure_port("0.0.0.0:50051")
    logger.info("Starting server on port 50051")
    server.start()
    logger.info("Server started")
    threading.Thread(target=listen_for_data, args=(data_queue,), daemon=True).start()
    server.wait_for_termination()