import logging
import multiprocessing as mp
from queue import Queue
import threading
import time
from typing import Any, cast

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

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

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


@cli.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 NotImplementedError("This command only works with live data")

    antenna_order.main()


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

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

    def visualise_data(data: npt.NDArray[Any], dtype: visualise.figures.Figure) -> 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.figures.Figure.RAW_CSI)
        processed = preprocessor.preprocess(antenna_data, visualiser=visualise_data)
        visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
        logger.info(f"Processed CSI data with shape {processed.shape}")
        processed_tensor = torch.tensor(processed, device=device)
        aoa.update(processed_tensor)
        aoa.heatmap(visualiser=visualise_data)

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


@cli.command()
def phase_analysis(
    subcarrier: int = 0, rx_antenna: int = 0, tx_antenna: int = 0
) -> None:
    """
    Visualise the phase information in the CSI data received from the antennas.

    The data goes through the same preprocessing steps as the heatmap command, but
    instead of going through the AoA estimation, we simply analyse the phase of the
    selected subcarrier and antenna.
    """

    app = CSIApplication()
    preprocessor = Preprocessor()

    subcarrier_phase: Queue[float] = Queue(config.collection_sample_rate)

    @app.on_sample
    def _(antenna_data: npt.NDArray[np.complex64]) -> None:
        processed = preprocessor.preprocess(antenna_data, visualiser=app.visualise_data)
        app.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)

        phase = cast(float, np.angle(processed[subcarrier, rx_antenna, tx_antenna]))
        if subcarrier_phase.full():
            subcarrier_phase.get()
        subcarrier_phase.put(phase)

    @app.on_process
    def _() -> None:
        logger.info(f"Updating phase visualisation")
        app.visualise_data(
            np.array(subcarrier_phase.queue), visualise.figures.Figure.PHASE_ANALYSIS
        )

    # globals.csi_producer(csi_callback=callback)
    app.start()
    logger.info("Finished processing CSI data")
    while threading.active_count() > 1:
        names = [
            t.name for t in threading.enumerate() if t != threading.current_thread()
        ]
        logger.info("Waiting for threads to close: " + ",".join(names))
        time.sleep(2)


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