import logging
import multiprocessing as mp
from typing import Any

import numpy as np
import numpy.typing as npt
import torch
import typer

from ..config import config
from ..processing.aoa import AoA
from ..processing.preprocess import Preprocessor
from ..visualise import server as visualise
from . import file, globals

app = typer.Typer(callback=globals.main)
logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


@app.command()
def antennas() -> None:
    """Utility to help determine the order in which antennas are plugged in

    Once the script is running, unplug and replug antennas from left to right, to get
    the correct order. Every time an antenna is unplugged and replugged, the script will
    print the antenna identifier. When you are done, press Ctrl+C to stop the script and
    get the final order.
    """
    from ..utils import antenna_order

    if not globals.is_live:
        raise ValueError("This command only works with live data")

    antenna_order.main()


@app.command()
def heatmap() -> None:
    preprocessor = Preprocessor()
    aoa = AoA()

    # Start webapp in background process
    webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(config.sample_rate)
    webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
    webapp.start()

    def visualise_data(data: npt.NDArray[Any], dtype: visualise.DataType) -> None:
        if not webapp_queue.full():
            webapp_queue.put(visualise.VisualiserData(data, dtype))

    def callback(antenna_data: npt.NDArray[np.complex64]) -> None:
        logger.info(f"Got final CSI data with shape {antenna_data.shape}")
        visualise_data(antenna_data, visualise.DataType.RAW_CSI)
        processed = preprocessor.preprocess(antenna_data)
        visualise_data(processed, visualise.DataType.PROCESSED_CSI)
        logger.info(f"Processed CSI data with shape {processed.shape}")
        processed_tensor = torch.tensor(processed, device=device)
        aoa.update(processed_tensor)
        heatmap = aoa.heatmap()
        visualise_data(heatmap, visualise.DataType.HEATMAP)

    globals.csi_producer(csi_callback=callback)
    logger.info("Finished processing CSI data")


app.add_typer(file.app, name="file", help="Commands for working with CSI files")