import logging
from queue import Queue
from typing import cast

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

from where_fi.collection import CSIMatrix

from ..application import CSIApplication
from ..config import config
from ..processing.aoa import AoA
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.csi_producer.is_live:
        raise NotImplementedError("This command only works with live data")

    antenna_order.main()


@cli.command()
def heatmap() -> None:
    app = CSIApplication(globals.csi_producer, visualise_raw=True)
    aoa = AoA()

    @app.on_sample
    def _(sample: npt.NDArray[np.complex64]) -> None:
        processed_tensor = torch.tensor(sample, device=device)
        aoa.update(processed_tensor)
        aoa.heatmap(visualiser=app.visualise_data)

    app.start()
    logger.info("Finished processing CSI data")


@cli.command()
def phase_analysis(
    subcarriers: list[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(globals.csi_producer)

    if isinstance(subcarriers, int):
        subcarriers = [subcarriers]
    print("Starting phase analysis on subcarriers: ", subcarriers)

    subcarrier_phase: dict[int, Queue[float]] = {
        x: Queue(config.collection_sample_rate) for x in subcarriers
    }

    @app.on_sample
    def _(sample: CSIMatrix) -> None:
        for subcarrier in subcarriers:
            phase = cast(float, np.angle(sample[subcarrier, rx_antenna, tx_antenna]))
            if subcarrier_phase[subcarrier].full():
                subcarrier_phase[subcarrier].get()
            subcarrier_phase[subcarrier].put(phase)

    @app.on_process
    def _(_: CSIMatrix) -> None:
        logger.info("Updating phase visualisation")
        app.visualise_data(
            np.array([x.queue for x in subcarrier_phase.values()]),
            visualise.figures.Figure.PHASE_ANALYSIS,
        )

    app.start()


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