from itertools import product
from queue import Queue
from typing import cast

import numpy as np

from where_fi.application import CSIApplication
from where_fi.collection import CSIMatrix
from where_fi.collection.ingest import RealtimeCSIProducer
from where_fi.config import config
from where_fi.visualise import server as visualise

producer = RealtimeCSIProducer()
app = CSIApplication(producer, visualise_raw=True)
visualise.figures.all_figures["median"] = visualise.figures.RandomVariable(
    "Median Phase"
)
visualise.figures.all_figures["magn"] = visualise.figures.RandomVariable(
    "Median Magnitude"
)

measurements: list[np.complex64] = []


subcarriers = [0, 1]
rx_antenna = [0, 1]
tx_antenna = [0]

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

proc_phase: dict[tuple[int, int, int], Queue[float]] = {
    x: Queue(config.collection_sample_rate)
    for x in product(subcarriers, rx_antenna, tx_antenna)
}

proc_magn: dict[tuple[int, int, int], Queue[float]] = {
    x: Queue(config.collection_sample_rate)
    for x in product(subcarriers, rx_antenna, tx_antenna)
}


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


cnt = 0


@app.on_process
def _(proc: CSIMatrix) -> None:
    global cnt
    global proc_phase

    for i in product(subcarriers, rx_antenna, tx_antenna):
        phase = cast(float, np.angle(proc[i]))
        magn = cast(float, np.abs(proc[i]))
        if proc_phase[i].full():
            proc_phase[i].get()
        proc_phase[i].put(phase)
        if proc_magn[i].full():
            proc_magn[i].get()
        proc_magn[i].put(magn)
    if cnt % 100 == 0:
        print(sum(proc_phase[0, 0, 0].queue))
        app.visualise_data(
            np.array([x.queue for x in proc_phase.values()]),
            "median",
        )
        app.visualise_data(
            np.array([x.queue for x in proc_magn.values()]),
            "magn",
        )
        app.visualise_data(
            np.array([x.queue for x in subcarrier_phase.values()]),
            visualise.figures.Figure.PHASE_ANALYSIS,
        )
        app.visualise_data(
            np.array([x.queue for x in subcarrier_magn.values()]),
            visualise.figures.Figure.MAGN_ANALYSIS,
        )
    cnt += 1


if __name__ == "__main__":
    print("Starting app")
    app.start()