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

app = CSIApplication(visualise_raw=True)
visualise.figures.all_figures["median"] = visualise.figures.RandomVariable(
    "Median Phase"
)
visualise.figures.all_figures["median_magn"] = visualise.figures.RandomVariable(
    "Median Magnitude"
)
visualise.figures.all_figures["denoised"] = visualise.figures.PerAntennaFigure(
    [
        visualise.figures.SimpleLineChart("Denoised CSI Phase", "Subcarrier", "Phase"),
        visualise.figures.SimpleLineChart(
            "Denoised CSI Amplitude", "Subcarrier", "Amplitude"
        ),
    ],
    [np.angle, np.abs],
)
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)
    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()]),
        "median_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,
    )
    app.visualise_data(proc, "denoised")
    cnt += 1


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