import logging
from datetime import datetime

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

from ..config import config

logger = logging.getLogger(__name__)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_device(device)


class AoA:
    def __init__(self) -> None:
        self.historical_autocorr = torch.tensor([], dtype=torch.complex64)
        self.N_subcarriers = -1
        self.N_rx = -1
        self.timestamp = datetime.now()
        pass

    def smooth(self, data: torch.Tensor) -> torch.Tensor:
        assert len(data.shape) == 3

        M = data.shape[0]  # Number of subcarriers
        N = data.shape[1]  # Number of RX antennas
        T = data.shape[2]  # Number of TX antennas

        self.N_subcarriers = M
        self.N_rx = N

        logger.debug(f"Smoothing: Subcarriers: {M}, RX antennas: {N}, TX antennas: {T}")

        # This only works with 1 TX antenna (i.e. no MIMO) - see #4 for more details
        assert T == 1, "The current implementation only supports 1 TX antenna"

        H_n = torch.zeros((N, M // 2, M // 2 + 1), dtype=torch.complex64)

        for i in range(N):
            for j in range(M // 2):
                H_n[i, j] = data[j : j + M // 2 + 1, i, 0]

        H_sm_rows = [torch.hstack(list(H_n[i : i + N // 2 + 1])) for i in range(N // 2)]
        H_sm = torch.vstack(H_sm_rows)

        logger.debug(f"Smoothed: {H_sm.shape}")

        return H_sm

    def update(self, data: torch.Tensor) -> None:
        self.timestamp = datetime.now()
        H_sm = self.smooth(data)

        auto_corr = H_sm @ torch.conj(H_sm).T
        # This matrix is by definition Hermitian.
        # Therefore, all of its eigenvectors are orthogonal.

        if len(self.historical_autocorr.shape) <= 1:
            self.historical_autocorr = torch.unsqueeze(auto_corr, 0)
        else:
            self.historical_autocorr = torch.cat(
                (self.historical_autocorr, torch.unsqueeze(auto_corr, 0))
            )

        WINDOW_SIZE = config.music.window_size
        if self.historical_autocorr.shape[0] > WINDOW_SIZE:
            self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]

        # Is the moving average also Hermitian?
        R = torch.mean(self.historical_autocorr, dim=0)

        # The smallest eigenvectors span the noise subspace,
        # and the largest span the signal subspace.
        eigvals, eigvecs = torch.linalg.eigh(R)
        self.E_n = (
            eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold].cpu().numpy()
        )

    def steering_vector(
        self, theta: float, tof: float
    ) -> npt.NDArray[np.complexfloating]:
        omega_t: npt.NDArray[np.complex64] = np.exp(-2j * np.pi * config.delta_f * tof)
        phi_theta: npt.NDArray[np.complex64] = np.exp(
            2j
            * np.pi
            * config.central_freq_hz
            * config.antennas.spacing
            * (1 - np.cos(theta))
            / 299_792_458
        )

        omega_t = np.expand_dims(omega_t, axis=-1)
        phi_theta = np.expand_dims(phi_theta, axis=-1)

        antenna_v = omega_t ** torch.arange(self.N_subcarriers // 2)
        phis = phi_theta ** torch.arange(self.N_rx // 2)
        antenna_v = np.expand_dims(antenna_v, axis=-1)
        print(antenna_v.shape, phis.shape)
        steering = antenna_v[0] * phis
        print(steering.shape)
        return steering.T.reshape(-1)

    def evaluate(self, theta: float, tof: float) -> float:
        try:
            steering = self.steering_vector(theta, tof)
            steering_h = np.conj(steering).T
        except Exception as e:
            logger.exception(e)
            return 0

        assert isinstance(self.E_n, torch.Tensor)
        E_n = self.E_n
        E_n_H = np.conj(E_n).T
        c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
        return np.abs(c.real).item()


def test_smoothing() -> None:
    row, col = np.indices((6, 4))
    data = row + 1j * col
    data = np.expand_dims(data, axis=2)
    np.set_printoptions(linewidth=200)
    print(data.shape)
    aoa = AoA()
    aoa.N_subcarriers = 6
    aoa.N_rx = 4
    smoothed = aoa.smooth(data)
    print(smoothed)


def test_steering_vector() -> None:
    aoa = AoA()
    aoa.N_subcarriers = 10
    aoa.N_rx = 2
    tau = 1
    theta = 0
    print(aoa.steering_vector(theta, tau))
    assert False