import logging

import numpy as np
import torch

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.processing import aoa

app = CSIApplication(visualise_raw=True)

logging.basicConfig(level=logging.INFO)

T = 500

N_sub1 = config.antennas.count // 2
N_sub2 = config.subcarriers // 2
L2 = config.subcarriers - N_sub2 + 1
L1 = config.antennas.count - N_sub1 + 1
N_sensors = config.subcarriers * config.antennas.count
historical = torch.zeros(T, N_sensors, N_sensors, dtype=torch.complex64)

cnt = 0


@app.on_sample
def _(sample: CSIMatrix) -> None:
    global cnt
    sample = sample.T.reshape(-1, 1)
    sample_tensor = torch.tensor(sample)
    historical[cnt] = sample_tensor @ torch.conj(sample_tensor).T
    cnt = (cnt + 1) % T


aoa = aoa.AoA()


@app.on_process
def _(_: CSIMatrix) -> None:
    R: torch.Tensor = torch.mean(historical, axis=0)
    Rss = torch.zeros(N_sub1 * N_sub2, N_sub1 * N_sub2, dtype=torch.complex64)
    for i in range(L1):
        for j in range(L2):
            Rss += R[i : i + N_sub1 * N_sub2, j : j + N_sub1 * N_sub2]
    Rss /= L1 * L2

    aoa.historical_autocorr = torch.unsqueeze(Rss, 0)
    aoa.heatmap(app.visualise_data)
    # eigvals, eigvecs = torch.linalg.eig(Rss)
    # app.visualise_data(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
    # E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]

    # print(E_n)
    # c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
    # return torch.abs(c)[:, 0, 0]


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