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.aoa import AoA
from where_fi.visualise import server as visualise

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_before_resize = sample[:, :, 0].T
    sample_after = sample_before_resize.reshape(-1, 1)
    sample_tensor = torch.tensor(sample_after)
    historical[cnt] = sample_tensor @ torch.conj(sample_tensor).T
    print(
        f"{sample.shape} => {sample_before_resize.shape} => {sample_after.shape} "
        f"=> {historical[cnt].shape}"
    )
    cnt = (cnt + 1) % T


def get_steering(theta: float, tau: float) -> torch.Tensor:
    """
    Calculate the alpha value for the given angle and time delay.
    """
    sub, ant = np.indices((N_sub1, N_sub2))
    alpha = np.exp(
        -1j
        * (
            2 * np.pi * (sub * config.delta_f * tau)
            + 2
            * np.pi
            * (
                ant
                * config.antennas.spacing
                * np.sin(theta)
                * 299_792_458
                / (config.central_freq_hz + (sub - 28) * config.delta_f)
            )
        )
    )
    alpha = sub + 1j * ant
    alpha = alpha.reshape(-1, 1)
    return torch.tensor(alpha, dtype=torch.complex64)


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
    Rss = R

    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]
    # E_n_H = torch.conj(E_n.T)

    # heatmap = np.zeros(
    #     (config.music.heatmap.tof_resolution, config.music.heatmap.theta_resolution)
    # )
    # for i_theta, theta in enumerate(
    #     np.linspace(0, np.pi, config.music.heatmap.theta_resolution)
    # ):
    #     for i_tau, tau in enumerate(
    #         np.linspace(
    #             0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
    #         )
    #     ):
    #         steering = get_steering(theta, tau)
    #         steering_h = torch.conj(steering.T)
    #         c = 1 / (steering_h @ E_n @ E_n_H @ steering)
    #         heatmap[i_tau, i_theta] = torch.abs(c)
    # print(steering)
    # app.visualise_data(heatmap, visualise.figures.Figure.AOA_HEATMAP)
    # 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()